Hello everyone, Was wondering if there was any documentation on best practices for doing model checkpointing, particularly in the context of multi-node distributed training? I came across this here - https://github.com/zillow/metaflow/blob/25cf84301fed372d80d3f9cbd788e86b3ef32132/metaflow/plugins/frameworks/checkpoints.py

we should document all this properly - it has been at the experimental status this far but it is maturing quickly. in the absence of better examples/docs, you can maybe get inspired by this piece of code

Thanks Ville - I'll check it out, thanks.

Another approach that may help is pushing to HuggingFace Hub, which this piece of code shows. It assumes the model state is on a single machine, so some extra care would need to be taken to sync the gradients if you are doing model parallel strategy.

Wow thanks Eddie! And cool flow using PEFT. Didn't know you could do that on HuggingFace. Probably not feasible for proprietary models though. Great for personal projects though!

ahh yes, proprietary case this won't work. where are you trying to dump the checkpoint? S3?

Yup S3 🙂 Just looking for a good example to follow.

Would love to know if that tutorial covers your needs, but suspect you may need a bit more heavy duty approach if the checkpointing is happening a lot. I have some time I can allocate tomorrow to try and debug a bit. Is it this same workflow from your thread earlier this week? Would like to bring myself up to speed on both issues and come up with more solid suggestions for you & others.

Yes that's correct. Same workflow as before, yup. We are gearing up to support training of large foundational models now, and both checkpointing and the progress bar not being displayed are the most critical pieces right now. Customers haven't asked for these things yet, but I suspect they will when they start onboarding to the training platform, shortly. Trying to stay ahead of the curve so they don't become blockers for them.

Ok thanks for the nice problem statements and your patience with these features. Will dive in tomorrow and writeup issues/workarounds I run into.

No problem. Appreciate your help. No apologies necessary. Everyone at Outerbounds has been so supportive and gotten us this far! We should be thanking YOU.

Was wondering how the checkpointing logic worked. Do you create checkpoints every

N

epochs for the Flow based on

FlowName/RunID/epoch_number

and if the job fails in the middle of training retry automatically using

@retry

based on the latest added model from

FlowName/RunID/epoch_number

?

yeah, that's what we have been prototyping. Even better, when you

resume

it resumes from the latest checkpoint. Basically in your code you just check if a checkpoint is available. If it is, use it. If it isn't, start from scratch