Repository: huggingface/transformers
Pull request: fix load_best_model_checkpoint_at_end do not load the best model chec…

…kpoint when `save_strategy` is `best`

# What does this PR do?
fix load_best_model_checkpoint_at_end do not load the best model checkpoint at the end when `save_strategy` is `"best"`


Fixes # (issue)
fix load_best_model_checkpoint_at_end do not load the best model checkpoint at the end when `save_strategy` is `"best"`.
This bug is due to the `SaveStrategy.BEST` is not added to some if statements resulting in the 1self.state.best_model_checkpoint1 is always None, which leads to the code not loading the best model checkpoint once done training. 

## Before submitting
- [x] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR.

@SunMarc
@stevhliu




Changed files:
filename: src/transformers/trainer.py
patch: @@ -3045,7 +3045,10 @@ def _save_checkpoint(self, model: nn.Module, trial: "optuna.Trial | dict[str, An
         output_dir = os.path.join(run_dir, checkpoint_folder)
         self.save_model(output_dir, _internal_call=True)
 
-        if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH] and self.state.best_global_step:
+        if (
+            self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH, SaveStrategy.BEST]
+            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():
@@ -3126,7 +3129,7 @@ def _determine_best_metric(self, metrics: dict[str, float], trial: "optuna.Trial
             if operator(metric_value, self.state.best_metric):
                 self.state.best_metric = metric_value
 
-                if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH]:
+                if self.args.save_strategy in [SaveStrategy.STEPS, SaveStrategy.EPOCH, SaveStrategy.BEST]:
                     self.state.best_global_step = self.state.global_step
 
                 is_new_best_metric = True

filename: src/transformers/training_args.py
patch: @@ -511,7 +511,7 @@ class TrainingArguments:
             Load the best checkpoint at the end of training. Requires `eval_strategy` to be set.
             When enabled, the best checkpoint is always saved (see `save_total_limit`).
             <Tip>
-            When `True`, `save_strategy` must match `eval_strategy`, and if using `"steps"`,
+            When `True`, `save_strategy` must match `eval_strategy` (unless `save_strategy` is `"best"`), and if using `"steps"`,
             `save_steps` must be a multiple of `eval_steps`.
             </Tip>
         metric_for_best_model (`str`, *optional*):
@@ -1666,7 +1666,7 @@ def _validate_args(self):
         if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST:
             if self.eval_strategy != self.save_strategy:
                 raise ValueError(
-                    "--load_best_model_at_end requires the save and eval strategy to match, but found\n- Evaluation "
+                    '--load_best_model_at_end requires the save and eval strategy to match, except when --save_strategy="best", but found\n- Evaluation '
                     f"strategy: {self.eval_strategy}\n- Save strategy: {self.save_strategy}"
                 )
             if self.eval_strategy == IntervalStrategy.STEPS and self.save_steps % self.eval_steps != 0:

filename: tests/trainer/test_trainer_checkpointing.py
patch: @@ -1594,6 +1594,35 @@ def setUp(self):
         self.n_epochs = args.num_train_epochs
         self.batch_size = args.train_batch_size
 
+    def test_load_best_model_with_save_best(self):
+        # Regression test: when save_strategy="best", the best model checkpoint should
+        # be loaded at the end of training, not the last one.
+        tmp_dir = self.get_auto_remove_tmp_dir()
+        trainer = get_regression_trainer(
+            output_dir=tmp_dir,
+            save_strategy="best",
+            eval_strategy="steps",
+            eval_steps=5,
+            load_best_model_at_end=True,
+            save_total_limit=2,
+            max_steps=11,
+        )
+        trainer.train()
+
+        # Check that best_model_checkpoint was set
+        assert trainer.state.best_model_checkpoint is not None, (
+            "trainer.state.best_model_checkpoint is None. Cannot load the best model checkpoint."
+        )
+
+        # Check that the right model was loaded in at the end of training —
+        # trainer.model weights should match the best checkpoint, not the last one saved
+        model_state = trainer.model.state_dict()
+        final_model_weights = safetensors.torch.load_file(
+            os.path.join(trainer.state.best_model_checkpoint, "model.safetensors")
+        )
+        for k, v in model_state.items():
+            assert torch.allclose(v, final_model_weights[k]), f"{k} is not the same"
+
     def test_load_best_model_with_save(self):
         tmp_dir = self.get_auto_remove_tmp_dir()
         trainer = get_regression_trainer(