Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions src/megatron/bridge/recipes/moonlight/h100/moonlight_16b.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
32 changes: 32 additions & 0 deletions tests/unit_tests/recipes/test_moonlight_recipes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading