diff --git a/docs/training/checkpointing.md b/docs/training/checkpointing.md index 9334ff8893..455c6ace16 100644 --- a/docs/training/checkpointing.md +++ b/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 diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 1fc2ff158a..85db05553d 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/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( diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 49c94625c2..4dfec5c299 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/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())