[ckpt] fix: Skip unavailable RNG shards after DP resize (#5988)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/checkpointing.py

@@ -610,13 +610,14 @@ def get_rng_state(
     return rng_state_list
 
 
-def _align_rng_state_sharded_metadata(rng_state: ShardedObject, checkpoint_name: str) -> ShardedObject:
+def _align_rng_state_sharded_metadata(rng_state: ShardedObject, checkpoint_name: str) -> ShardedObject | None:
     """Align RNG load metadata with the layout stored in a torch-dist checkpoint.
 
     Newer MCore checkpoints shard RNG state across PP, TP, and DP/CP, while
     older checkpoints encode DP as a replica ID. Keep the generated metadata
     when its exact key exists; otherwise adopt a unique stored layout whose
-    PP/TP prefix matches this rank.
+    PP/TP prefix matches this rank. Return ``None`` when this rank is outside
+    the stored DP/CP extent so the caller can keep its freshly initialized RNG.
     """
     checkpoint_path = Path(checkpoint_name)
     if not (checkpoint_path / ".metadata").is_file():
@@ -627,17 +628,21 @@ def _align_rng_state_sharded_metadata(rng_state: ShardedObject, checkpoint_name:
         return rng_state
 
     prefix = rng_state.global_offset[:2]
-    matches = [
+    compatible_layouts = [
         metadata
         for metadata in sharded_metadata.values()
         if isinstance(metadata, ShardedObject)
         and metadata.key == rng_state.key
         and metadata.global_offset[:2] == prefix
         and len(metadata.global_offset) == len(rng_state.global_offset) + 1
-        and metadata.global_offset[-1] == rng_state.replica_id
         and metadata.replica_id == 0
     ]
+    matches = [metadata for metadata in compatible_layouts if metadata.global_offset[-1] == rng_state.replica_id]
     if len(matches) != 1:
+        if compatible_layouts and all(
+            rng_state.replica_id >= metadata.global_shape[-1] for metadata in compatible_layouts
+        ):
+            return None
         return rng_state
 
     stored = matches[0]
@@ -3015,6 +3020,9 @@ def _load_checkpoint_from_path(
             )
             if ckpt_type != CheckpointType.LOCAL:
                 gen_sd_rng_state = _align_rng_state_sharded_metadata(gen_sd_rng_state, checkpoint_name)
+                if gen_sd_rng_state is None:
+                    ignore_rng_state = True
+                    print_rank_0("RNG state has no shard for this DP/CP rank; using freshly initialized RNG state")
         else:
             ignore_rng_state = True
             gen_sd_rng_state = None

tests/unit_tests/training/test_checkpointing.py

@@ -6051,6 +6051,17 @@ def test_adopts_unique_dp_cp_sharded_layout(self, mock_strategy, _mock_is_file):
         assert result.global_offset == (1, 3, 5)
         assert result.replica_id == 0
 
+    @patch("megatron.bridge.training.checkpointing.Path.is_file", return_value=True)
+    @patch("megatron.bridge.training.checkpointing.TorchDistLoadShardedStrategy")
+    def test_returns_none_without_a_dp_cp_shard_for_this_rank(self, mock_strategy, _mock_is_file):
+        rng_state = self._rng(replica_id=5)
+        stored = [self._rng(global_offset=(1, 3, rank), global_shape=(2, 4, 4), replica_id=0) for rank in range(4)]
+        mock_strategy.return_value.load_sharded_metadata.return_value = {item.unique_key: item for item in stored}
+
+        result = _align_rng_state_sharded_metadata(rng_state, "/checkpoint")
+
+        assert result is None
+
     @pytest.mark.parametrize("stored", [{}, None])
     @patch("megatron.bridge.training.checkpointing.Path.is_file", return_value=True)
     @patch("megatron.bridge.training.checkpointing.TorchDistLoadShardedStrategy")