From 862874f72a2d3acd027a4a9c7657a111f95e5b6e Mon Sep 17 00:00:00 2001 From: Pruthviraj Prakash Date: Thu, 8 Oct 2026 11:38:29 -0700 Subject: [PATCH] fix(recipe): give Moonlight-16B 8-chunk pipeline layouts 27 decoder layers The PP=8/VP=1 and PP=4/VP=2 layouts in _get_moonlight_pipeline_layout described 4 + 4*6 + 3 = 31 decoder layers for a 27-layer model, so Megatron-Core's validate_layer_layout rejects them at config finalization. Re-split both as a front-loaded [4, 4, 4, 3, 3, 3, 3, 3]. Adds a unit test that runs every supported Moonlight layout through Megatron-Core layout validation at the model depth (red-green verified). Signed-off-by: Pruthviraj Prakash --- .../recipes/moonlight/h100/moonlight_16b.py | 10 ++++-- .../recipes/test_moonlight_recipes.py | 32 +++++++++++++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py b/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py index 7b20560f9c..ec0279aef4 100644 --- a/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py +++ b/src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py @@ -138,9 +138,15 @@ def _get_moonlight_pipeline_layout(pp_size: int, vp_size: int | None): (1, 1): None, (2, 1): [["embedding"] + ["decoder"] * 14, ["decoder"] * 13 + ["loss"]], (4, 1): [["embedding"] + ["decoder"] * 7] + [["decoder"] * 7] * 2 + [["decoder"] * 6 + ["loss"]], - (8, 1): [["embedding"] + ["decoder"] * 4] + [["decoder"] * 4] * 6 + [["decoder"] * 3 + ["loss"]], + (8, 1): [["embedding"] + ["decoder"] * 4] + + [["decoder"] * 4] * 2 + + [["decoder"] * 3] * 4 + + [["decoder"] * 3 + ["loss"]], (2, 2): [["embedding"] + ["decoder"] * 7] + [["decoder"] * 7] * 2 + [["decoder"] * 6 + ["loss"]], - (4, 2): [["embedding"] + ["decoder"] * 4] + [["decoder"] * 4] * 6 + [["decoder"] * 3 + ["loss"]], + (4, 2): [["embedding"] + ["decoder"] * 4] + + [["decoder"] * 4] * 2 + + [["decoder"] * 3] * 4 + + [["decoder"] * 3 + ["loss"]], } vp_size = 1 if vp_size is None else vp_size if (pp_size, vp_size) not in map_pp_vp_to_layout: diff --git a/tests/unit_tests/recipes/test_moonlight_recipes.py b/tests/unit_tests/recipes/test_moonlight_recipes.py index 8df3971a01..2161918c53 100644 --- a/tests/unit_tests/recipes/test_moonlight_recipes.py +++ b/tests/unit_tests/recipes/test_moonlight_recipes.py @@ -531,6 +531,38 @@ def test_moonlight_pipeline_layout_tracks_supported_runner_override(monkeypatch: assert cfg.model.moe_hybridep_pad_uneven_dispatch_inputs is True +_MOONLIGHT_16B_NUM_LAYERS = 27 + + +def _supported_moonlight_pp_vp_layouts(): + """Return every (pp, vp, layout) the Moonlight-16B layout helper supports.""" + from megatron.bridge.recipes.moonlight.h100.moonlight_16b import _get_moonlight_pipeline_layout + + supported = [] + for pp in range(1, 17): + for vp in range(1, 9): + try: + layout = _get_moonlight_pipeline_layout(pp, vp) + except ValueError: + continue + if layout is not None: + supported.append((pp, vp, layout)) + return supported + + +def test_moonlight_pipeline_layouts_match_model_depth(): + """Every supported Moonlight-16B pipeline layout must pass MCore layout validation.""" + from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout + + supported = _supported_moonlight_pp_vp_layouts() + assert supported + + for pp, vp, layout in supported: + parsed = PipelineParallelLayerLayout(layout, pipeline_model_parallel_size=pp) + assert parsed.virtual_pipeline_model_parallel_size == vp, (pp, vp) + parsed.validate_layer_layout(num_layers=_MOONLIGHT_16B_NUM_LAYERS, mtp_num_layers=None) + + def test_moonlight_16b_peft_convergence_contract(monkeypatch: pytest.MonkeyPatch): """Test the exact 4-GPU LoRA convergence contract and frozen base.""" from megatron.bridge.peft.lora import LoRA