From 1634ac82e0607589aaf4256375428c820647e845 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Tue, 14 May 2024 08:48:39 -0400
Subject: [PATCH] make sure to save on the last step (#1615)

---
 src/axolotl/utils/callbacks/__init__.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/axolotl/utils/callbacks/__init__.py b/src/axolotl/utils/callbacks/__init__.py
index e66f165a53..2965ac1e29 100644
--- 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
     ):