fix(checkpoint): use global FSDP pipeline layer keys

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/training/checkpointing.py

@@ -139,6 +139,12 @@
 except ImportError:
     HAVE_MEGATRON_FSDP = False
 
+try:
+    from megatron.core.transformer.fsdp_dtensor_checkpoint import handle_pipeline_parallel_in_state_dict
+except ImportError:
+    handle_pipeline_parallel_in_state_dict = None
+
+
 try:
     from megatron.core.transformer.fsdp_dtensor_checkpoint import handle_gdn_in_state_dict
 except ImportError:
@@ -2291,7 +2297,7 @@ def preprocess_fsdp_dtensor_state_dict(cfg, raw_state_dict: dict[str, Any], mode
         getattr(model_config, "gated_linear_unit", False) and getattr(model_config, "activation_func", None) is F.silu
     )
 
-    def apply(handler):
+    def apply(handler: Callable[..., tuple[dict[str, Any], dict[str, Any] | None]]) -> None:
         """Run a state dict handler over the model and, when present, the optimizer state."""
         model_state_dict, optimizer_state_dict = handler(model, state_dict["model"], state_dict.get("optimizer"))
         state_dict["model"] = model_state_dict
@@ -2324,6 +2330,15 @@ def apply(handler):
     if handle_mtp_in_state_dict is not None:
         apply(handle_mtp_in_state_dict)
 
+    # Globalize layer names last: earlier handlers resolve stage-local module paths.
+    if handle_pipeline_parallel_in_state_dict is not None:
+        apply(handle_pipeline_parallel_in_state_dict)
+    elif model_config.pipeline_model_parallel_size > 1:
+        raise RuntimeError(
+            "FSDP pipeline checkpointing requires a Megatron-Core version with "
+            "handle_pipeline_parallel_in_state_dict; older versions collide stage-local layer names."
+        )
+
     preprocess_state_dict_for_uneven_dtensor(state_dict)
 
     return state_dict

tests/unit_tests/training/test_checkpointing.py

@@ -4946,6 +4946,7 @@ def test_preprocess_fsdp_dtensor_state_dict(self):
             mock_model_config = Mock()
             mock_model_config.gated_linear_unit = False
             mock_model_config.num_moe_experts = None
+            mock_model_config.pipeline_model_parallel_size = 1
             mock_get_config.return_value = mock_model_config
 
             with (
@@ -4956,15 +4957,47 @@ def test_preprocess_fsdp_dtensor_state_dict(self):
                 patch("megatron.bridge.training.checkpointing.handle_gdn_in_state_dict", None),
                 patch("megatron.bridge.training.checkpointing.handle_mla_down_proj_in_state_dict", None),
                 patch("megatron.bridge.training.checkpointing.handle_mtp_in_state_dict", None),
+                patch(
+                    "megatron.bridge.training.checkpointing.handle_pipeline_parallel_in_state_dict",
+                    side_effect=lambda model, state, optim: (state, optim),
+                ) as mock_pipeline,
             ):
                 raw_state_dict = {"model": {"test_param": torch.tensor([1.0])}}
                 result = preprocess_fsdp_dtensor_state_dict(mock_cfg, raw_state_dict, mock_model)
 
                 # Should call FP8 and uneven dtensor preprocessing
                 mock_fp8.assert_called_once()
                 mock_uneven.assert_called_once()
+                mock_pipeline.assert_called_once_with(mock_model, raw_state_dict["model"], None)
                 assert "model" in result
 
+    @pytest.mark.parametrize("pp", [1, 2])
+    def test_fsdp_pipeline_checkpoint_requires_core_handler(self, pp: int) -> None:
+        """Old Core remains usable for PP=1 but must not silently corrupt PP>1."""
+        from types import SimpleNamespace
+
+        from megatron.bridge.training import checkpointing
+
+        model_config = SimpleNamespace(gated_linear_unit=False, num_moe_experts=None, pipeline_model_parallel_size=pp)
+        with (
+            patch("megatron.core.utils.get_model_config", return_value=model_config),
+            patch.object(checkpointing, "handle_fp8_extra_state_case"),
+            patch.object(checkpointing, "preprocess_state_dict_for_uneven_dtensor"),
+            patch.multiple(
+                checkpointing,
+                handle_gdn_in_state_dict=None,
+                handle_mla_down_proj_in_state_dict=None,
+                handle_mtp_in_state_dict=None,
+                handle_pipeline_parallel_in_state_dict=None,
+            ),
+        ):
+            state = {"model": {"weight": torch.ones(1)}}
+            if pp > 1:
+                with pytest.raises(RuntimeError, match="collide stage-local layer names"):
+                    checkpointing.preprocess_fsdp_dtensor_state_dict(None, state, None)
+            else:
+                assert checkpointing.preprocess_fsdp_dtensor_state_dict(None, state, None) == state
+
     @patch("megatron.bridge.training.checkpointing.HAVE_MEGATRON_FSDP", True)
     def test_save_fsdp_dtensor_checkpoint_preprocesses_and_saves(self, tmp_path):
         """Test public FSDP DTensor save helper preprocesses and calls PyTorch DCP."""