fix: Epoch checkpointing index off by 1 (#2911)

Author: bogdansalyp

recipes/full_finetune_distributed.py

@@ -915,8 +915,9 @@ def validate(self) -> dict[str, float]:
         return log_dict
 
     def save_checkpoint(self, *, epoch: int, full_tensors: bool):
+        training_progress_epoch = epoch
         if self.global_step % self._steps_per_epoch == 0:
-            epoch += 1
+            training_progress_epoch += 1
 
         self._checkpoint_client.save_checkpoint(
             model=self._model,
@@ -927,7 +928,7 @@ def save_checkpoint(self, *, epoch: int, full_tensors: bool):
             ),
             training_progress=TrainingProgress(
                 seed=self.seed,
-                epochs_run=epoch,
+                epochs_run=training_progress_epoch,
                 total_epochs=self.total_epochs,
                 max_steps_per_epoch=self.max_steps_per_epoch,
                 steps_run=self.global_step,

recipes/full_finetune_single_device.py

@@ -526,14 +526,16 @@ def save_checkpoint(self, *, epoch: int, step: int, full_tensors: bool) -> None:
         correctly creating the checkpoint dict and passing to the checkpointer.
         """
         # Since we might save at an epoch boundary, we need to increment the epoch counter
+        training_progress_epoch = epoch
         if step % self._steps_per_epoch == 0:
-            epoch += 1
+            training_progress_epoch += 1
+
         self._checkpoint_client.save_checkpoint(
             model=self._model,
             optimizer=self.optimizer,
             training_progress=TrainingProgress(
                 seed=self.seed,
-                epochs_run=epoch,
+                epochs_run=training_progress_epoch,
                 total_epochs=self.total_epochs,
                 max_steps_per_epoch=self.max_steps_per_epoch,
                 dataloader_state_dict=self._dataloader.state_dict(),

tests/recipes/test_full_finetune_distributed.py

@@ -397,7 +397,7 @@ def test_training_state_on_resume_from_distributed_checkpoint_single_rank(
 
         resumed_log_dir = (tmpdir / "resumed/").mkdir()
         resumed_log_file = gen_log_file_name(resumed_log_dir)
-        shutil.rmtree((tmpdir / "epoch_2"))
+        shutil.rmtree((tmpdir / "epoch_1"))
 
         # Resume training
         cmd_2 = f"""
@@ -504,7 +504,7 @@ def test_training_state_on_resume_from_distributed_checkpoint_multi_rank(
 
         resumed_log_dir = (tmpdir / "resumed/").mkdir()
         resumed_log_file = gen_log_file_name(resumed_log_dir)
-        shutil.rmtree((tmpdir / "epoch_2"))
+        shutil.rmtree((tmpdir / "epoch_1"))
 
         # Resume training
         cmd_2 = f"""

tests/recipes/test_full_finetune_single_device.py

@@ -204,8 +204,8 @@ def test_training_state_on_resume(self, tmpdir, use_steps, monkeypatch):
             final_ckpt_dir = "step_4"
             prev_ckpt_dir = "step_2"
         else:
-            final_ckpt_dir = "epoch_2"
-            prev_ckpt_dir = "epoch_1"
+            final_ckpt_dir = "epoch_1"
+            prev_ckpt_dir = "epoch_0"
         cmd_1 = cmd_1 + self._get_test_config_overrides() + model_config
 
         monkeypatch.setattr(sys, "argv", cmd_1)