[recipe] fix: Resync Kimi pipeline layout after overrides (#5008)

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

Author: yaoyu-33

scripts/training/recipe_runner.py

@@ -368,6 +368,34 @@ def sync_finetuning_cp_invariants(config: ConfigContainer, *, mode: str) -> Conf
     return config
 
 
+def sync_model_pipeline_layout(
+    config: ConfigContainer,
+    *,
+    cli_overrides: list[str],
+) -> ConfigContainer:
+    """Rebuild a recipe-owned pipeline layout after PP or VP overrides."""
+    override_fields = {override.lstrip("+~").split("=", 1)[0] for override in cli_overrides}
+    topology_fields = {
+        "model.pipeline_model_parallel_size",
+        "model.virtual_pipeline_model_parallel_size",
+    }
+    if not override_fields.intersection(topology_fields):
+        return config
+    if "model.pipeline_model_parallel_layout" in override_fields:
+        return config
+
+    model = getattr(config, "model", None)
+    layout_builder = getattr(model, "_pipeline_model_parallel_layout_builder", None)
+    if layout_builder is None:
+        return config
+
+    model.pipeline_model_parallel_layout = layout_builder(
+        model.pipeline_model_parallel_size,
+        model.virtual_pipeline_model_parallel_size,
+    )
+    return config
+
+
 def sync_offline_packing_alignment(config: ConfigContainer) -> ConfigContainer:
     """Align offline-packed samples to the resolved length and parallel topology."""
     dataset = getattr(config, "dataset", None)

scripts/training/run_recipe.py

@@ -100,6 +100,7 @@
     run_config,
     sync_finetuning_cp_invariants,
     sync_model_dataset_sequence_length,
+    sync_model_pipeline_layout,
     sync_offline_packing_alignment,
 )
 
@@ -418,6 +419,7 @@ def main(argv: list[str] | None = None) -> None:
     recipe = _apply_dataset(recipe, args)
     recipe = apply_determinism(recipe, deterministic=args.deterministic)
     recipe = apply_cli_overrides(recipe, cli_overrides)
+    recipe = sync_model_pipeline_layout(recipe, cli_overrides=cli_overrides)
     if benchmark_metadata is not None:
         recipe = _apply_benchmark_runtime_defaults(recipe, benchmark_metadata, cli_overrides)
     configuration_mode = _train_mode(args.mode)

src/megatron/bridge/recipes/kimi/h100/kimi_k2.py

@@ -26,7 +26,7 @@
 from megatron.bridge.training.mixed_precision import MixedPrecisionConfig
 
 
-def _get_kimi_k2_pipeline_layout(pp_size: int, vp_size: int):
+def _get_kimi_k2_pipeline_layout(pp_size: int, vp_size: int | None) -> list[list[str]] | None:
     """Get pipeline layout for Kimi-K2 based on PP and VP size."""
     map_pp_vp_to_layout = {
         (1, 1): None,
@@ -121,6 +121,7 @@ def kimi_k2_pretrain_512gpu_h100_bf16_config() -> ConfigContainer:
     cfg.model.num_layers_in_last_pipeline_stage = None
 
     # Set pipeline layout
+    cfg.model._pipeline_model_parallel_layout_builder = _get_kimi_k2_pipeline_layout
     cfg.model.pipeline_model_parallel_layout = _get_kimi_k2_pipeline_layout(16, 1)
 
     # Tokenizer - uses NullTokenizer with model vocab_size

tests/unit_tests/recipes/kimi/test_kimi_k2.py

@@ -271,3 +271,4 @@ def test_pretrain_config_pipeline_layout(self):
         # Default PP=16, VP=None (1), should have a layout
         expected_layout = _get_kimi_k2_pipeline_layout(16, 1)
         assert cfg.model.pipeline_model_parallel_layout == expected_layout
+        assert cfg.model._pipeline_model_parallel_layout_builder is _get_kimi_k2_pipeline_layout

tests/unit_tests/scripts/training/test_recipe_runner.py

@@ -244,6 +244,56 @@ def test_sync_model_dataset_sequence_length_uses_canonical_dataset_field(recipe_
     assert config.model.seq_length == 256
 
 
+def test_sync_model_pipeline_layout_uses_overridden_topology(recipe_runner: ModuleType) -> None:
+    layout = [["first"], ["middle"], ["middle"], ["last"]]
+    layout_builder = Mock(return_value=layout)
+    config = SimpleNamespace(
+        model=SimpleNamespace(
+            pipeline_model_parallel_size=4,
+            virtual_pipeline_model_parallel_size=1,
+            pipeline_model_parallel_layout=[["stale"]] * 16,
+            _pipeline_model_parallel_layout_builder=layout_builder,
+        )
+    )
+
+    assert (
+        recipe_runner.sync_model_pipeline_layout(
+            config,
+            cli_overrides=[
+                "model.pipeline_model_parallel_size=4",
+                "model.virtual_pipeline_model_parallel_size=1",
+            ],
+        )
+        is config
+    )
+    assert config.model.pipeline_model_parallel_layout == layout
+    layout_builder.assert_called_once_with(4, 1)
+
+
+def test_sync_model_pipeline_layout_preserves_explicit_layout_override(recipe_runner: ModuleType) -> None:
+    layout = [["custom"]]
+    layout_builder = Mock()
+    config = SimpleNamespace(
+        model=SimpleNamespace(
+            pipeline_model_parallel_size=1,
+            virtual_pipeline_model_parallel_size=1,
+            pipeline_model_parallel_layout=layout,
+            _pipeline_model_parallel_layout_builder=layout_builder,
+        )
+    )
+
+    recipe_runner.sync_model_pipeline_layout(
+        config,
+        cli_overrides=[
+            "model.pipeline_model_parallel_size=1",
+            "model.pipeline_model_parallel_layout=[[custom]]",
+        ],
+    )
+
+    assert config.model.pipeline_model_parallel_layout == layout
+    layout_builder.assert_not_called()
+
+
 @pytest.mark.parametrize(
     "dataset",
     [

tests/unit_tests/scripts/training/test_run_recipe.py

@@ -47,6 +47,7 @@ def _load_module():
         "load_recipe",
         "run_config",
         "sync_finetuning_cp_invariants",
+        "sync_model_pipeline_layout",
         "sync_offline_packing_alignment",
         "sync_model_dataset_sequence_length",
     ):
@@ -56,6 +57,7 @@ def _load_module():
     recipe_runner.apply_runtime_environment.side_effect = lambda config: config
     recipe_runner.bootstrap_recipe_environment.side_effect = lambda config, **_: config
     recipe_runner.sync_finetuning_cp_invariants.side_effect = lambda config, **_: config
+    recipe_runner.sync_model_pipeline_layout.side_effect = lambda config, **_: config
     recipe_runner.sync_offline_packing_alignment.side_effect = lambda config: config
     recipe_runner.sync_model_dataset_sequence_length.side_effect = lambda config: config
     recipe_runner.load_forward_step.return_value = object()
@@ -224,6 +226,59 @@ def test_full_recipe_uses_library_recipe_and_default_llm_step():
     handles.recipe_runner.load_forward_step.assert_called_once_with("llm_step", mode="pretrain")
 
 
+def test_kimi_supported_pp_vp_override_refreshes_pipeline_layout():
+    module, handles = _load_module()
+    default_layout = [[f"stage-{index}"] for index in range(16)]
+    overridden_layout = [[f"stage-{index}"] for index in range(4)]
+    config = SimpleNamespace(
+        model=SimpleNamespace(
+            pipeline_model_parallel_size=16,
+            virtual_pipeline_model_parallel_size=None,
+            pipeline_model_parallel_layout=default_layout,
+            _pipeline_model_parallel_layout_builder=lambda pp, vp: overridden_layout
+            if (pp, vp) == (4, 1)
+            else default_layout,
+        )
+    )
+    handles.recipe_runner.load_recipe.return_value = config
+
+    def apply_override(current_config, overrides):
+        assert overrides == [
+            "model.pipeline_model_parallel_size=4",
+            "model.virtual_pipeline_model_parallel_size=1",
+        ]
+        current_config.model.pipeline_model_parallel_size = 4
+        current_config.model.virtual_pipeline_model_parallel_size = 1
+        return current_config
+
+    def refresh_layout(current_config, *, cli_overrides):
+        assert cli_overrides == [
+            "model.pipeline_model_parallel_size=4",
+            "model.virtual_pipeline_model_parallel_size=1",
+        ]
+        model = current_config.model
+        model.pipeline_model_parallel_layout = model._pipeline_model_parallel_layout_builder(
+            model.pipeline_model_parallel_size,
+            model.virtual_pipeline_model_parallel_size,
+        )
+        return current_config
+
+    def validate_layout(*, config, **_):
+        model = config.model
+        detected_vp = len(model.pipeline_model_parallel_layout) // model.pipeline_model_parallel_size
+        assert detected_vp == model.virtual_pipeline_model_parallel_size, (
+            "virtual_pipeline_model_parallel_size conflicts with the pipeline layout"
+        )
+
+    handles.recipe_runner.apply_cli_overrides.side_effect = apply_override
+    handles.recipe_runner.sync_model_pipeline_layout.side_effect = refresh_layout
+    handles.recipe_runner.run_config.side_effect = validate_layout
+
+    module.main(["--recipe", "kimi_k2_pretrain_config", "-pp", "4", "-vp", "1"])
+
+    assert config.model.pipeline_model_parallel_layout == overridden_layout
+
+
 @pytest.mark.parametrize(
     ("recipe_name", "mode", "step_name"),
     [