only main process should call _save on deepspeed zero3

[zjjMaiMai]

Background

trainer._save call on all process after #25817. will raise FileExistsError when model save.

What does this PR do?

this pr fix it, trainer._save will call on main process only.

[amyeroberts]

cc @muellerzr @pacman100

[muellerzr]

This does not feel like the right solution, as should_save will check if we are the main process or not, and we should only ever be saving once. You can have self.model_wrapped.save_checkpoint under that self.args.should_save I believe, but I'm not sure this is the right "fix". Definitely cc @pacman100 here

[zjjMaiMai]

self.model_wrapped.save_checkpoint needs to be called on each process, because each process only have part of model weight under deepspeed zero3.

[muellerzr]

Got it, thanks for the clarification!

[pacman100]

Hello, Thank you for the fix! However, it needs correction as explained in the comment,

[pacman100]

This would remove the legitimate model checkpoint when stage3_gather_16bit_weights_on_model_save=True. remove_dummy_checkpoint should only be called in the exception block.

[zjjMaiMai]

thanks for your explanation, it has been fixed.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[zjjMaiMai]

@pacman100 any things i need to do?

[pacman100]

Thank you for iterating, LGTM!

[amyeroberts]

@zjjMaiMai One of the hub tests are failing, complaining that the base_model is empty when pushing to the hub. Could you try running this test locally to see whether it's a result of the changes in this PR?

[zjjMaiMai]

@zjjMaiMai One of the hub tests are failing, complaining that the base_model is empty when pushing to the hub. Could you try running this test locally to see whether it's a result of the changes in this PR?

$ pytest tests/trainer/test_trainer.py -k 'test_push_to_hub'
================================================================================================================================ test session starts ================================================================================================================================
platform linux -- Python 3.9.2, pytest-7.4.1, pluggy-1.3.0
configfile: setup.cfg
plugins: timeout-2.1.0, hypothesis-6.84.2, dash-2.13.0, xdist-3.3.1, anyio-3.7.1
collected 95 items / 91 deselected / 4 selected                                                                                                                                                                                                                                     

tests/trainer/test_trainer.py ssss                                                                                                                                                                                                                                            [100%]

================================================================================================================================= warnings summary ==================================================================================================================================
../../../../../../home/.local/lib/python3.9/site-packages/_pytest/config/__init__.py:1376
  /home/.local/lib/python3.9/site-packages/_pytest/config/__init__.py:1376: PytestConfigWarning: Unknown config option: doctest_glob
  
    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=================================================================================================================== 4 skipped, 91 deselected, 1 warning in 1.69s ====================================================================================================================
$ git branch 
* fix_save_deepspeed_3
  main

[amyeroberts]

@zjjMaiMai Could you try and rebase on main? This should resolve the failing tests.

[zjjMaiMai]

All green! @amyeroberts