make sure to save on the last step (#1615)

axolotl-ai-cloud · May 14, 2024 · 1634ac8 · 1634ac8
1 parent 0298273
commit 1634ac8
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py
@@ -778,6 +778,17 @@ def on_train_begin(
 class SaveModelOnTrainEndCallback(TrainerCallback):
     """Callback to save model on train end"""
 
+    def on_step_end(  # pylint: disable=unused-argument
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        # Save
+        if state.global_step >= state.max_steps:
+            control.should_save = True
+
     def on_train_end(  # pylint: disable=unused-argument
         self, args, state, control, **kwargs
     ):