Repository: huggingface/transformers
Pull request: Wait for main process in _save_checkpoint to ensure best checkpoint exists

# What does this PR do?

In 4.50.0 a bug was introduced with a refactor of the best checkpoint process when we try to load the best model at end and  a non-main process has a null `self.state.best_model_checkpoint`. 

Before we run `_load_best_model()` there is a barrier to make sure the main process saves the model, however inside `_save_checkpoint`, this is missing. The model is saved [here](https://github.com/huggingface/transformers/blob/0b057e66b52556da3a1cbc29e2a98c0784ea9c33/src/transformers/trainer.py#L3197) and then a couple lines later the `self.state.best_model_checkpoint` is set only if the `best_checkpoint_dir` exists. When another process gets to this line first, `self.state.best_model_checkpoint` will be null. 

Finally, when training reaches the end and it is time to load the best model [here](https://github.com/huggingface/transformers/blob/0b057e66b52556da3a1cbc29e2a98c0784ea9c33/src/transformers/trainer.py#L2666) only the main process enters the if statement and the dist.barrier causes it to wait forever for process 1 to arrive which doesn't have a best_model_checkpoint.

In 4.48.3 it didn't check for existence before assigning the best model checkpoint. They refactored how this works, I can't get it to fail on a non-cube example yet. https://github.com/huggingface/transformers/blob/298b3f19303294293f7af075609481d64cb13de3/src/transformers/trainer.py#L3190 (edited)

Note: I didn't use `self.accelerator.wait_for_everyone()` since it doesn't seem to cover the sagemaker case, but I could make a wait function in the trainer since it is used twice now. 

I believe this fixes #38008


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?
@zach-huggingface, @SunMarc and @qgallouedec


Changed files:
filename: src/transformers/trainer.py
patch: @@ -2765,14 +2765,6 @@ def _inner_training_loop(
 
         logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
         if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
-            # Wait for everyone to get here so we are sure the model has been saved by process 0.
-            if is_torch_xla_available():
-                xm.rendezvous("load_best_model_at_end")
-            elif args.parallel_mode == ParallelMode.DISTRIBUTED:
-                dist.barrier()
-            elif is_sagemaker_mp_enabled():
-                smp.barrier()
-
             self._load_best_model()
 
         # add remaining tr_loss
@@ -3289,6 +3281,15 @@ def _save_checkpoint(self, model, trial):
         self.save_model(output_dir, _internal_call=True)
 
         if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH] and self.state.best_global_step:
+            # Wait for everyone to get here so we are sure the model has been saved by process 0
+            # before we check if the best_checkpoint_dir exists
+            if is_torch_xla_available():
+                xm.rendezvous("load_best_model_at_end")
+            elif self.args.parallel_mode == ParallelMode.DISTRIBUTED:
+                dist.barrier()
+            elif is_sagemaker_mp_enabled():
+                smp.barrier()
+
             best_checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.best_global_step}"
             best_checkpoint_dir = os.path.join(run_dir, best_checkpoint_folder)