fix(checkpoint): use global FSDP v1 pipeline layer keys

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,16 @@ def apply(handler):
     if handle_mtp_in_state_dict is not None:
         apply(handle_mtp_in_state_dict)
 
+    # V2 optimizer keys are already global. Only v1 needs this final renaming.
+    if getattr(getattr(cfg, "ddp", None), "megatron_fsdp_version", 1) == 1:
+        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

@@ -4930,8 +4930,9 @@ def test_generate_state_dict_fsdp_dtensor_no_preprocessing(self):
                 3.1 if hasattr(distrib_optimizer, "get_legacy_grad_dtypes") else 3.0
             )
 
+    @pytest.mark.parametrize("version", [1, 2])
     @patch("megatron.bridge.training.checkpointing.HAVE_MEGATRON_FSDP", True)
-    def test_preprocess_fsdp_dtensor_state_dict(self):
+    def test_preprocess_fsdp_dtensor_state_dict(self, version: int) -> None:
         """Test preprocess_fsdp_dtensor_state_dict applies all preprocessing steps."""
         from unittest.mock import Mock
 
@@ -4940,12 +4941,14 @@ def test_preprocess_fsdp_dtensor_state_dict(self):
         # Create mock model and config
         mock_model = Mock()
         mock_cfg = Mock()
+        mock_cfg.ddp.megatron_fsdp_version = version
 
         # Mock model config without swiglu or experts
         with patch("megatron.core.utils.get_model_config") as mock_get_config:
             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 +4959,52 @@ 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()
+                if version == 1:
+                    mock_pipeline.assert_called_once_with(mock_model, raw_state_dict["model"], None)
+                else:
+                    mock_pipeline.assert_not_called()
                 assert "model" in result
 
+    @pytest.mark.parametrize("pp", [1, 2])
+    @pytest.mark.parametrize("version", [1, 2])
+    def test_fsdp_pipeline_checkpoint_requires_core_handler(self, pp: int, version: 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)}}
+            cfg = SimpleNamespace(ddp=SimpleNamespace(megatron_fsdp_version=version))
+            if pp > 1 and version == 1:
+                with pytest.raises(RuntimeError, match="collide stage-local layer names"):
+                    checkpointing.preprocess_fsdp_dtensor_state_dict(cfg, state, None)
+            else:
+                assert checkpointing.preprocess_fsdp_dtensor_state_dict(cfg, 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."""