[checkpoint] fix: edge case in redundant checkpoint saving (#2622)

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

src/megatron/bridge/training/pretrain.py

@@ -17,7 +17,6 @@
 
 from megatron.bridge.data.utils import get_dataset_provider
 from megatron.bridge.training.callbacks import Callback, CallbackManager, normalize_callbacks
-from megatron.bridge.training.checkpointing import save_checkpoint
 from megatron.bridge.training.config import ConfigContainer, runtime_config_update
 from megatron.bridge.training.eval import evaluate_and_print_results
 from megatron.bridge.training.forward_step_func_types import ForwardStepCallable
@@ -154,17 +153,6 @@ def _pretrain(
             )
 
         barrier_and_log("after training is done")
-        ckpt_config = config.checkpoint
-        if ckpt_config.save and state.train_state.step != 0 and ckpt_config.save_interval != 0:
-            save_checkpoint(
-                state,
-                model,
-                optimizer,
-                scheduler,
-                state.train_state.floating_point_operations_so_far,
-                ckpt_context,
-                train_data_iterator=train_data_iterator,
-            )
 
     else:
         print_rank_0("skipping training ...")

src/megatron/bridge/training/train.py

@@ -588,6 +588,27 @@ def train(
         if should_exit:
             break
 
+    # Save final checkpoint when training completes normally and the last
+    # step wasn't already persisted by the interval-based save inside
+    # checkpoint_and_decide_exit.
+    if not should_exit:
+        ckpt_config = config.checkpoint
+        if (
+            ckpt_config.save
+            and global_state.train_state.step != 0
+            and ckpt_config.save_interval != 0
+            and (ckpt_config.save_interval is None or global_state.train_state.step % ckpt_config.save_interval != 0)
+        ):
+            save_checkpoint_and_time(
+                global_state,
+                model,
+                optimizer,
+                scheduler,
+                num_floating_point_operations_so_far,
+                checkpointing_context,
+                train_data_iterator=train_data_iterator,
+            )
+
     _delete_cuda_graphs(cuda_graph_helper)
 
     # Flush TensorBoard, WandB writers and one-logger.