test(checkpointing): Added num_seen_steps_previous_run to keep every k tests.

Co-authored-by: Copilot <copilot@github.com>

Author: BlueCrescent

tests/checkpointing/test_checkpoint_strategies.py

@@ -52,15 +52,19 @@ def test_checkpoint_strategy_k(
 
 
 @pytest.mark.parametrize(
-    "k, num_recent_checkpoints_to_keep, num_steps",
+    "k, num_recent_checkpoints_to_keep, num_steps, num_seen_steps_previous_run, num_seen_tokens_previous_run",
     [
-        (3, 2, 11),
-        (2, 1, 10),
-        (4, 3, 15),
+        (3, 2, 11, 0, 0),
+        (2, 1, 10, 2, 4),
+        (4, 3, 15, 3, 6),
     ],
 )
 def test_keep_every_k_strategy_has_no_unexpected_checkpoints(
-    k: int, num_recent_checkpoints_to_keep: int, num_steps: int
+    k: int,
+    num_recent_checkpoints_to_keep: int,
+    num_steps: int,
+    num_seen_steps_previous_run: int,
+    num_seen_tokens_previous_run: int,
 ) -> None:
     checkpoint_strategy = KeepEveryKStepsAndMMostRecentCheckpointingStrategy(
         k=k, num_recent_checkpoints_to_keep=num_recent_checkpoints_to_keep
@@ -70,31 +74,42 @@ def test_keep_every_k_strategy_has_no_unexpected_checkpoints(
         num_seen_tokens_current_run=0,
         num_target_steps=20,
         num_target_tokens=40,
+        num_seen_steps_previous_run=num_seen_steps_previous_run,
+        num_seen_tokens_previous_run=num_seen_tokens_previous_run,
     )
 
     # Simulate training progress and checkpointing
     simulator = _CheckpointSavingSimulator()
-    for step in range(1, num_steps + 1):
+    for step in range(num_steps + 1):
         training_progress.num_seen_steps_current_run = step
         checkpoint_instruction = checkpoint_strategy.get_checkpoint_instruction(training_progress=training_progress)
         simulator.simulate_training_step(training_progress, checkpoint_instruction)
 
     for ckpt in simulator.saved_checkpoints:
         # Check that only checkpoints that are divisible by k or the most recent ones are kept.
-        last_checkpoints = set(range(num_steps - num_recent_checkpoints_to_keep + 1, num_steps + 1))
+        last_checkpoints = set(
+            range(
+                num_seen_steps_previous_run + num_steps - num_recent_checkpoints_to_keep + 1,
+                num_seen_steps_previous_run + num_steps + 1,
+            )
+        )
         assert ckpt.num_seen_steps_total % k == 0 or ckpt.num_seen_steps_total in last_checkpoints
 
 
 @pytest.mark.parametrize(
-    "k, num_recent_checkpoints_to_keep, num_steps",
+    "k, num_recent_checkpoints_to_keep, num_steps, num_seen_steps_previous_run, num_seen_tokens_previous_run",
     [
-        (3, 2, 11),
-        (2, 1, 10),
-        (4, 3, 15),
+        (3, 2, 11, 0, 0),
+        (2, 1, 10, 2, 4),
+        (4, 3, 15, 3, 6),
     ],
 )
 def test_keep_every_k_strategy_has_no_unexpected_deletions(
-    k: int, num_recent_checkpoints_to_keep: int, num_steps: int
+    k: int,
+    num_recent_checkpoints_to_keep: int,
+    num_steps: int,
+    num_seen_steps_previous_run: int,
+    num_seen_tokens_previous_run: int,
 ) -> None:
     checkpoint_strategy = KeepEveryKStepsAndMMostRecentCheckpointingStrategy(
         k=k, num_recent_checkpoints_to_keep=num_recent_checkpoints_to_keep
@@ -104,6 +119,8 @@ def test_keep_every_k_strategy_has_no_unexpected_deletions(
         num_seen_tokens_current_run=0,
         num_target_steps=20,
         num_target_tokens=40,
+        num_seen_steps_previous_run=num_seen_steps_previous_run,
+        num_seen_tokens_previous_run=num_seen_tokens_previous_run,
     )
 
     # Simulate training progress and checkpointing
@@ -113,9 +130,9 @@ def test_keep_every_k_strategy_has_no_unexpected_deletions(
         checkpoint_instruction = checkpoint_strategy.get_checkpoint_instruction(training_progress=training_progress)
         simulator.simulate_training_step(training_progress, checkpoint_instruction)
 
-    for i in range(1, num_steps + 1):
+    for i in range(num_seen_steps_previous_run + 1, num_seen_steps_previous_run + num_steps + 1):
         # Check that checkpoints that are divisible by k or the most recent ones are not deleted.
-        if i % k == 0 or i > num_steps - num_recent_checkpoints_to_keep:
+        if i % k == 0 or i > num_seen_steps_previous_run + num_steps - num_recent_checkpoints_to_keep:
             assert any(ckpt.num_seen_steps_total == i for ckpt in simulator.saved_checkpoints)