[ckpt] revert: clear incompatible saved pipeline layouts (#5838)

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

Author: yaoyu-33

src/megatron/bridge/training/model_load_save.py

@@ -458,7 +458,6 @@ def load_megatron_model(
         otherwise returns a dictionary containing the full, unsharded model state_dict.
     """
     model_cfg, mlm_args = load_model_config(checkpoint_path)
-    saved_pipeline_model_parallel_size = model_cfg.pipeline_model_parallel_size
     # If in single GPU environment, reset additional parallel settings
     model_cfg.tensor_model_parallel_size = 1
     model_cfg.pipeline_model_parallel_size = 1
@@ -483,16 +482,9 @@ def load_megatron_model(
     # Apply model-parallel overrides if provided
     if mp_overrides:
         for key, value in mp_overrides.items():
-            if hasattr(model_cfg, key) and (value is not None or key == "pipeline_model_parallel_layout"):
+            if hasattr(model_cfg, key) and value is not None:
                 setattr(model_cfg, key, value)
 
-        if (
-            "pipeline_model_parallel_size" in mp_overrides
-            and model_cfg.pipeline_model_parallel_size != saved_pipeline_model_parallel_size
-            and "pipeline_model_parallel_layout" not in mp_overrides
-        ):
-            model_cfg.pipeline_model_parallel_layout = None
-
     # A saved flexible layout describes PP/VPP stage ownership. It must not be
     # reinterpreted as virtual pipeline chunks after collapsing to one rank.
     if model_cfg.pipeline_model_parallel_size == 1 and model_cfg.virtual_pipeline_model_parallel_size is None:

tests/unit_tests/training/test_model_load_save.py

@@ -310,56 +310,6 @@ def _built_layer_count(checkpoint_path, model_cfg, *args):
 
         assert built_layer_count == provider.num_layers
 
-    @pytest.mark.parametrize(
-        ("mp_overrides", "expect_saved_layout"),
-        [
-            ({"pipeline_model_parallel_size": 2}, True),
-            ({"pipeline_model_parallel_size": 4}, False),
-            (
-                {
-                    "pipeline_model_parallel_size": 2,
-                    "pipeline_model_parallel_layout": None,
-                },
-                False,
-            ),
-        ],
-        ids=["same-pp", "changed-pp", "explicit-clear"],
-    )
-    @patch("megatron.bridge.training.model_load_save.build_and_load_model")
-    @patch("megatron.bridge.training.model_load_save.load_model_config")
-    def test_pipeline_override_does_not_reuse_incompatible_saved_layout(
-        self,
-        mock_load_model_config,
-        mock_build_and_load_model,
-        mp_overrides,
-        expect_saved_layout,
-    ):
-        """A saved PP layout is retained only for a compatible requested topology."""
-        provider = GPTModelProvider(
-            num_layers=4,
-            hidden_size=16,
-            num_attention_heads=2,
-            pipeline_model_parallel_size=2,
-            pipeline_model_parallel_layout=[
-                ["embedding", "decoder", "decoder"],
-                ["decoder", "decoder", "loss"],
-            ],
-        )
-        mock_load_model_config.return_value = (provider, None)
-
-        def _finalized_layout(checkpoint_path, model_cfg, *args):
-            model_cfg.finalize()
-            return model_cfg.pipeline_model_parallel_layout
-
-        mock_build_and_load_model.side_effect = _finalized_layout
-
-        result = load_megatron_model("/ckpt", mp_overrides=mp_overrides)
-
-        if expect_saved_layout:
-            assert isinstance(result, PipelineParallelLayerLayout)
-        else:
-            assert result is None
-
     @patch("megatron.bridge.training.model_load_save.temporary_distributed_context")
     @patch("megatron.bridge.training.checkpointing._load_model_weights_from_checkpoint")
     @patch("megatron.bridge.utils.instantiate_utils.instantiate")