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

src/axolotl/utils/callbacks/__init__.py

@@ -778,6 +778,17 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
 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
     ):