[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.

[SilverSoldier]

@SunMarc do we need approval from anyone else? Thanks!

[SunMarc]

I think @muellerzr wanted to check if there were no issues in multi-node scenario also.

[muellerzr]

Thanks! Small scale test seems fine, but let's keep an eye out on issues to see if we need to redo this approach again 🤗

[SilverSoldier]

Thanks @muellerzr @SunMarc!

[wizyoung]

I call for reverting this pull request, not everyone has enough space in the tmp path to store model weights, this approach is too inappropriate.

[muellerzr]

@wizyoung the tmp path is in the same directory location to where the final model weights will be stored. We're looking into other issues related to it, however what you are describing is not how this PR works at all

[wizyoung]

@muellerzr You are right. This PR has broken my model saving, causing me some frustration, and I haven't reviewed the code carefully, sorry about that.

[muellerzr]

@wizyoung no worries, we'll have a solution by EOD :) (which may include reverting, we'll see)

[muellerzr]

@wizyoung can you try with pip install git+https://github.com/huggingface/transformers@muellerzr-os-err?

[wizyoung]

@muellerzr Using your branch, the model save failure is gone!

[wizyoung]

@muellerzr New issue arrised. In this patch, under deepspeed zero2 8xH100 training, the model saves the model incompletely. Using the original save mode:

With you patch:

[acdart]

call for reverting this pull request

this pr do not consider any multi-process save problem, like rng states

@muellerzr @SunMarc

[rwightman]

This approach will not work if it's being used where multiple processes write a checkpoint. To work in that scenario it'd neeed synchronization across the processes and that will increase the testing requirements significantly.

Even the proposal to to check the existence of files becomes much more complicated. Since there wasn
t an awareness of it being used in this context there likely isn't exhaustive understanding of all the files you'd have to look for in all of the various scenarios.

That probably leaves the catch / detect error on load and try the the previous checkpoint as the best option, though that could also fail to work for distributed cases. Some processes could have succeeded and some fail but they all need to moec to the same other checkpoint tro try, again synchronization and lots of testing needed

[myron]

New issue arrised. In this patch, under deepspeed zero2 8xH100 training, the model saves the model incompletely. Using the original save mode:
@wizyoung @muellerzr

same issue, with deepspeed zero2, it saved incomplete checkpoint (with a single file bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt).
Can we please open a bug about it, or revert this PR. thanks

[Elenore1997]

@wizyoung can you try with pip install git+https://github.com/huggingface/transformers@muellerzr-os-err?

my checkpoint savings are incomplete when training qwen2_5_vl model in zero2 with 8*a800

[SunMarc]

Thanks for everyone feedbacks ! We will most probably revert the PR today until @muellerzr manages to find a fix before