diff --git a/src/megatron/bridge/training/checkpointing.py b/src/megatron/bridge/training/checkpointing.py index 8936e64ed9..ed710d0878 100644 --- a/src/megatron/bridge/training/checkpointing.py +++ b/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 diff --git a/tests/unit_tests/training/test_checkpointing.py b/tests/unit_tests/training/test_checkpointing.py index dc07cc689b..6a2eb98d59 100644 --- a/tests/unit_tests/training/test_checkpointing.py +++ b/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,6 +4959,10 @@ 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) @@ -4963,8 +4970,41 @@ def test_preprocess_fsdp_dtensor_state_dict(self): # 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."""