[ckpt] fix: Restore checkpoints without Energon state

docs/training/checkpointing.md

@@ -312,8 +312,8 @@ The load default follows the actual source because the checkpoint restored on re
 
 ### Restore Behavior and Failure Modes
 
-- If the dataloader state directory is **absent** (e.g. a checkpoint saved before this feature existed), the dataloader starts from the beginning and a message is logged.
-- If the directory **exists** but the current rank's per-DP-rank file is **missing**, restore **raises**. This almost always means the data-parallel size changed since the checkpoint was saved; resuming would silently change the data order, so it fails loudly instead.
+- If the dataloader state root or the selected `iter_N` directory is **absent** (e.g. the selected checkpoint was saved before this feature existed), the dataloader starts from the beginning and a message is logged. Other state generations under the same root do not change this behavior.
+- If the selected `iter_N` directory **exists** but the current rank's per-DP-rank file is **missing**, restore **raises**. This almost always means the data-parallel size changed since the checkpoint was saved; resuming would silently change the data order, so it fails loudly instead.
 
 ### Determinism Requirement
 

src/megatron/bridge/training/checkpointing.py

@@ -1803,11 +1803,11 @@ def maybe_load_dataloader_state(
     on *every* rank: each tensor/pipeline/context rank pulls from its own data iterator (e.g.
     ``qwen3_vl`` ``get_batch``), so all of them must be rewound to the saved position.
 
-    Restore failure modes are deliberately loud. If the dataloader state directory is absent
-    entirely, the checkpoint predates dataloader-state saving and the dataloader starts fresh. But
-    if the directory exists while the current rank's state file does not, the data-parallel size
-    almost certainly changed since the checkpoint was saved; rather than silently resume with a
-    different data order, this raises.
+    Restore failure modes distinguish checkpoint generations. If the dataloader state root or the
+    selected iteration is absent, that checkpoint predates dataloader-state saving and the
+    dataloader starts fresh. If the selected iteration exists while the current rank's state file
+    does not, the data-parallel size almost certainly changed since the checkpoint was saved;
+    rather than silently resume with a different data order, this raises.
 
     Restoring is only correct when the task encoder is deterministic per sample (Energon replays the
     samples since the last checkpoint by re-running the pipeline) — see
@@ -1843,8 +1843,13 @@ def maybe_load_dataloader_state(
         print_rank_0(f"no dataloader state under {dataloader_load_path}; dataloader starts from the beginning")
         return
 
-    dp_rank = get_pg_rank(pg_collection.dp)
     iter_dir = get_checkpoint_name(dataloader_load_path, iteration)
+    if not is_dir(iter_dir):
+        # This checkpoint generation predates dataloader-state saving. Start from scratch.
+        print_rank_0(f"no dataloader state for iteration {iteration}; dataloader starts from the beginning")
+        return
+
+    dp_rank = get_pg_rank(pg_collection.dp)
     data_state_load_path = join_paths(iter_dir, f"train_dataloader_dprank{dp_rank:03d}.pt")
     if not is_file(data_state_load_path):
         raise RuntimeError(

tests/unit_tests/training/test_checkpointing.py

@@ -5052,10 +5052,20 @@ def test_missing_dir_warns_and_skips(self, tmp_path):
         maybe_load_dataloader_state(train_iterator, 10, str(missing), pg_collection=self._pg())
         train_iterator.iterable.restore_state.assert_not_called()
 
+    def test_missing_selected_iteration_warns_and_skips(self, tmp_path):
+        """An unrelated state generation does not prevent an older checkpoint from resuming."""
+        train_iterator = Mock()
+        os.makedirs(get_checkpoint_name(str(tmp_path), 20))
+
+        maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg())
+
+        train_iterator.iterable.restore_state.assert_not_called()
+
     def test_existing_dir_missing_file_raises(self, tmp_path):
         """If the state dir exists but this rank's file does not, fail loudly (likely a DP-size change)."""
         train_iterator = Mock()
-        # tmp_path exists (the energon root) but contains no per-rank file for this iteration.
+        os.makedirs(get_checkpoint_name(str(tmp_path), 10))
+        # The selected iteration exists but contains no state file for this data-parallel rank.
         with pytest.raises(RuntimeError, match="data-parallel size"):
             maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg())