From cc7016affc7618d049afe0f5f744942f58b6086d Mon Sep 17 00:00:00 2001 From: Yu Yao Date: Tue, 18 Aug 2026 02:28:52 -0700 Subject: [PATCH] [ckpt] fix: Reject Energon resume after DP changes Signed-off-by: Yu Yao --- src/megatron/bridge/training/checkpointing.py | 14 +++++++++++ .../unit_tests/training/test_checkpointing.py | 24 ++++++++++++++++++- 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 7c81c0286a..2ff584b5a4 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/src/megatron/bridge/training/checkpointing.py @@ -1857,6 +1857,20 @@ def maybe_load_dataloader_state( print_rank_0(f"no dataloader state for iteration {iteration}; dataloader starts from the beginning") return + state_file_pattern = join_paths(iter_dir, "train_dataloader_dprank*.pt") + saved_dp_size = ( + len(msc.glob(state_file_pattern)) + if msc is not None + else len(list(Path(iter_dir).glob("train_dataloader_dprank*.pt"))) + ) + current_dp_size = get_pg_size(pg_collection.dp) + if saved_dp_size != current_dp_size: + raise RuntimeError( + f"Dataloader state at {iter_dir} was saved for data-parallel size {saved_dp_size}, but the current " + f"data-parallel size is {current_dp_size}. Resuming would silently change the training data order; " + "refusing to continue." + ) + 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): diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index 70130116bd..0352baae9b 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/tests/unit_tests/training/test_checkpointing.py @@ -5011,7 +5011,7 @@ def _dist_initialized(self): yield @staticmethod - def _pg(cp: int | None = 0, pp: int = 0, tp: int = 0, dp: int = 0): + def _pg(cp: int | None = 0, pp: int = 0, tp: int = 0, dp: int = 0, dp_size: int = 1): """Mock a ProcessGroupCollection with the given per-dimension ranks. Only cp may be None — modeling a collection configured without context parallelism (tp/pp/dp are always populated); the real get_pg_rank then reads the absent cp group as rank 0.""" @@ -5023,6 +5023,7 @@ def _pg(cp: int | None = 0, pp: int = 0, tp: int = 0, dp: int = 0): pg.pp.rank.return_value = pp pg.tp.rank.return_value = tp pg.dp.rank.return_value = dp + pg.dp.size.return_value = dp_size return pg def test_noop_when_no_path(self): @@ -5069,6 +5070,27 @@ def test_existing_dir_missing_file_raises(self, tmp_path): with pytest.raises(RuntimeError, match="data-parallel size"): maybe_load_dataloader_state(train_iterator, 10, str(tmp_path), pg_collection=self._pg()) + @patch("megatron.bridge.training.checkpointing.energon_torch_load") + def test_extra_saved_dp_rank_file_raises(self, mock_load, tmp_path): + """A smaller current DP size must not silently discard saved rank state.""" + train_iterator = Mock() + iter_dir = Path(get_checkpoint_name(str(tmp_path), 10)) + iter_dir.mkdir(parents=True) + (iter_dir / "train_dataloader_dprank000.pt").touch() + (iter_dir / "train_dataloader_dprank001.pt").touch() + pg_collection = self._pg(dp=0, dp_size=1) + + with pytest.raises(RuntimeError, match="data-parallel size"): + maybe_load_dataloader_state( + train_iterator, + 10, + str(tmp_path), + pg_collection=pg_collection, + ) + + mock_load.assert_not_called() + train_iterator.iterable.restore_state.assert_not_called() + @patch("megatron.bridge.training.checkpointing.energon_torch_load") def test_restores_on_every_cp_tp_pp_rank(self, mock_load, tmp_path): """Unlike save, restore must run on every cp/tp/pp because