Save checkpoint to temporary directory to handle partial saves during failures

[SilverSoldier]

What does this PR do?

When auto-resuming from checkpoint, currently the "max" folder is picked as the checkpoint folder to resume from. If there are some missing files (like model, config or trainer_state), this results in FileNotFoundError. This PR picks the latest checkpoint folder (should ideally be last or second last) with all required files allowing training to resume instead of just throwing an error requiring manual removing of files.

Fixes #35782

Currently, I test for atleast one of the model weights along with definitely requiring config and trainer_state. From what I understood the others (like optimizer) were optional, but please let me know if I missed something.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Trainer: @muellerzr and @SunMarc

[SunMarc]

Thanks for the PR ! Left a comment. Can you double check @muellerzr ?

[SunMarc]

I think it can be good to display a warning to the user, specifying that we decided to skip a particular checkpoint because it is missing files.

[rwightman]

Checking for existence of files is only a partial solution to this problem, you can have partially written files in a crash during save.

The needed fixes include

• error handling on load: detect a failure during any part of the load (crash on loading a partial/corrupted checkpoint) and execute logic to find and try the preceeding one
• improve checkpoint writing: write to temporary files (or folders in this case) and swap to correct, final name(s) once all writes have succesfully finished. This is a common practice when writing checkpoints / restore states

[SilverSoldier]

@rwightman thanks for your comments.

Checking for existence of files is only a partial solution to this problem, you can have partially written files in a crash during save.

You're right, but, I believe the underlying write to file calls usually handle this issue of partial file when crashing. For instance, safetensors.save uses numpy.save which this SO issue mentions doesn't create file until success. Even in my own runs, I hadn't faced that issue. It could happen due to edge cases, but I believe much less likely than this case.

error handling on load: detect a failure during any part of the load (crash on loading a partial/corrupted checkpoint) and execute logic to find and try the preceeding one

I initially planned to fix it that way (try catch around the load and change dir), but unfortunately, the resume_from_checkpoint dir, once selected, is used in too many places to easily catch and then replace everything prior as well. Another problem is that the user could pass that variable as well (get_last_checkpoint is only called in case of auto-resume, in other cases the variable is set to what the user passes), in which case we don't want to change the directory and should probably raise an error (maybe a better error message than FileNotFound though).

improve checkpoint writing: write to temporary files (or folders in this case) and swap to correct, final name(s) once all writes have succesfully finished. This is a common practice when writing checkpoints / restore states

This is a good suggestion and could also solve this problem. At a quick glance, all the saving happens in _save_checkpoint, so writing everything to a tmp dir and finally moving everything should also work (with the small possibility of crash in the middle of the move).

[SunMarc]

improve checkpoint writing: write to temporary files (or folders in this case) and swap to correct, final name(s) once all writes have succesfully finished. This is a common practice when writing checkpoints / restore states

This solution from @rwightman seems to be the most robust one. Would you like to try to implement this @SilverSoldier ?

[SilverSoldier]

Sure, let me change the PR to implement the save to tmp dir approach

[rwightman]

@SilverSoldier in these cases, the usual situation is to do the writes to a distinct temp name in the same location as the final destination. So if you normally save checkpoints to /mycheckpoints/checkpoint_0 etc... then while writing the temp you'd use /mycheckpoints/_temp_checkpoint ... when that is successful you can then do a os.rename _temp_checkpoint -> checkpoint_0, etc.

For most filesystem renaming a file/folder on the same drive is an atomic operation and very low risk of failure.

I looked at the numpy save and it's a chunked write, pretty sure you can cause partial write failures in many situations.

[SilverSoldier]

Have changed to the save to tmp approach using @rwightman's idea (thanks!).

1. Using TemporaryDirectory in the same checkpoint directory with name tmp-checkpoint-. Because of the execution context, if this fails somewhere in the middle, the tmp directory is cleaned up. Finally, renaming to the correct name.
2. There is one corner case, if the last epoch is already saved and save_checkpoint is called again at the end of training. In that case, we get the File or Directory already exists error for the checkpoint-50 or whatever dir (I got diff errors when I tried with different setups). In that case, we need to remove the old one and rename the new one. I think both should be exactly the same (checked for one run). If they are 100% the same, we can catch it before even checkpointing (which should save some time as well).

@SunMarc could you please review the PR now?

[SunMarc]

Thanks for iterating ! Left a few comments. Can you double check @muellerzr ?

[SunMarc]

Make sure to fix the CI with make style

[SilverSoldier]

Fixed the code style issues.
I believe the failing tests after merging main (qwen2_5_vl.test_modeling_qwen2_5_vl.Qwen2_5_VLModelTest) are unrelated.

[SunMarc]

LGTM !

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.