fix load_model_end = true work when save_steps < eval_steps

src/transformers/trainer.py

@@ -3098,6 +3098,8 @@ def _maybe_log_save_evaluate(
         if self.control.should_evaluate:
             metrics = self._evaluate(trial, ignore_keys_for_eval)
             is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial)
+            if is_new_best_metric:
+                self._save_checkpoint(self._get_output_dir(trial), trial)
 
             if self.args.save_strategy == SaveStrategy.BEST:
                 self.control.should_save = is_new_best_metric
@@ -4115,18 +4117,15 @@ def _rotate_checkpoints(self, use_mtime=False, output_dir=None) -> None:
         if len(checkpoints_sorted) <= self.args.save_total_limit:
             return
 
-        # If save_total_limit=1 with load_best_model_at_end=True, we could end up deleting the last checkpoint, which
-        # we don't do to allow resuming.
         save_total_limit = self.args.save_total_limit
-        if (
-            self.state.best_model_checkpoint is not None
-            and self.args.save_total_limit == 1
-            and checkpoints_sorted[-1] != self.state.best_model_checkpoint
-        ):
-            save_total_limit = 2
+        best_model_checkpoint = self.state.best_model_checkpoint
+        preserve_best = getattr(self.args, "preserve_best_model", True)
 
         number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - save_total_limit)
         checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
+        if preserve_best and best_model_checkpoint is not None:
+            checkpoints_to_be_deleted = [ckpt for ckpt in checkpoints_to_be_deleted if ckpt != best_model_checkpoint]
+
         for checkpoint in checkpoints_to_be_deleted:
             logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
             shutil.rmtree(checkpoint, ignore_errors=True)

src/transformers/training_args.py

@@ -1024,6 +1024,12 @@ class TrainingArguments:
             )
         },
     )
+    preserve_best_model: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to preserve the best model checkpoint when rotating checkpoints due to save_total_limit."
+        },
+    )
     restore_callback_states_from_checkpoint: bool = field(
         default=False,
         metadata={
@@ -1664,26 +1670,15 @@ def __post_init__(self):
                     "--load_best_model_at_end requires the save and eval strategy to match, 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:
-                if self.eval_steps < 1 or self.save_steps < 1:
-                    if not (self.eval_steps < 1 and self.save_steps < 1):
-                        raise ValueError(
-                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
-                            "steps, which cannot get guaranteed when mixing ratio and absolute steps for save_steps "
-                            f"{self.save_steps} and eval_steps {self.eval_steps}."
-                        )
-                    # Work around floating point precision issues
-                    LARGE_MULTIPLIER = 1_000_000
-                    if (self.save_steps * LARGE_MULTIPLIER) % (self.eval_steps * LARGE_MULTIPLIER) != 0:
-                        raise ValueError(
-                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
-                            f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}."
-                        )
-                else:
-                    raise ValueError(
-                        "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation "
-                        f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}."
+            if self.eval_strategy == IntervalStrategy.STEPS:
+                if self.save_steps > 0 and self.eval_steps > 0:
+                    steps_aligned = (self.save_steps % self.eval_steps == 0) or (
+                        self.eval_steps % self.save_steps == 0
                     )
+                    if not steps_aligned:
+                        warnings.warn(
+                            "Warning: save_steps and eval_steps are misaligned. Best model will be saved at evaluation, not at a scheduled save step."
+                        )
 
         safetensors_available = is_safetensors_available()
         if self.save_safetensors and not safetensors_available: