create an explicit config-contract with optional list of moe experts per layer

Author: chochowski

src/megatron/bridge/models/conversion/model_bridge.py

@@ -259,18 +259,83 @@ def _install_pg_collection_on_mappings(
         mapping.set_process_groups_from_pg_collection(pg_collection)
 
 
+def _resolve_num_local_experts(config: TransformerConfig, param_name: str, ep_size: int) -> int:
+    """Resolve the per-EP-rank expert count for an expert parameter.
+
+    Expert numbering during conversion must match how the model was built, so it
+    is driven entirely by the model config — never by traversing an instantiated
+    module. This keeps name mapping independent of runtime layout (PP/VPP, MTP,
+    custom MoE wrappers) and of whether a live model is available.
+
+    The expert layout is an explicit config contract:
+
+    * Homogeneous models leave ``num_moe_experts_per_layer`` unset and every MoE
+      layer uses the scalar ``config.num_moe_experts``.
+    * Heterogeneous models set ``config.num_moe_experts_per_layer`` to a list of
+      global expert counts indexed by global decoder-layer number (``0`` for
+      dense layers).
+
+    Args:
+        config: The Megatron model config (authoritative expert layout).
+        param_name: Expert parameter name, already in *global* layer numbering
+            (the PP branch of :func:`_megatron_local_name_to_global` runs first).
+        ep_size: Expert-parallel group size.
+
+    Returns:
+        The number of experts owned by each expert-parallel rank for the layer
+        that ``param_name`` belongs to.
+
+    Raises:
+        ValueError: If the configured layout is missing or inconsistent with the
+            parameter being renumbered (unset scalar, per-layer list too short,
+            an expert parameter on a layer declared dense, or a global count not
+            divisible by ``ep_size``).
+    """
+    per_layer = getattr(config, "num_moe_experts_per_layer", None)
+
+    # Multi-token-prediction layers (``mtp.layers.N``) are not decoder layers and
+    # are not covered by the per-decoder-layer contract; match only decoder
+    # layers so MTP experts fall through to the scalar count.
+    layer_match = re.search(r"(?:^|\.)decoder\.layers\.(\d+)(?=\.)", param_name)
+
+    if per_layer is not None and layer_match is not None:
+        global_layer_idx = int(layer_match.group(1))
+        if global_layer_idx >= len(per_layer):
+            raise ValueError(
+                f"num_moe_experts_per_layer has length {len(per_layer)} but expert parameter "
+                f"{param_name!r} resolves to global layer {global_layer_idx}; the list must "
+                "cover every decoder layer (length == num_layers)."
+            )
+        num_global_experts = per_layer[global_layer_idx]
+        if not num_global_experts:
+            raise ValueError(
+                f"num_moe_experts_per_layer[{global_layer_idx}] == {num_global_experts}, but "
+                f"expert parameter {param_name!r} exists on that layer; layers declared dense "
+                "must not carry expert parameters."
+            )
+    else:
+        num_global_experts = config.num_moe_experts
+
+    if num_global_experts is None:
+        raise ValueError(
+            f"Cannot renumber expert parameter {param_name!r}: num_moe_experts is None and no "
+            "per-layer expert layout (num_moe_experts_per_layer) is configured."
+        )
+    if num_global_experts % ep_size != 0:
+        raise ValueError(
+            f"Global expert count {num_global_experts} for {param_name!r} is not divisible by "
+            f"expert-parallel size {ep_size}."
+        )
+    return num_global_experts // ep_size
+
+
 def _megatron_local_name_to_global(
     models: MegatronModule | List[MegatronModule],
     config: TransformerConfig,
     param_name: str,
     vp_stage: Optional[int] = None,
 ) -> str:
     """Adjust layer number and expert number from local to global numbering."""
-    # ``layer_module`` is resolved lazily and shared between the PP and EP
-    # branches. It stays ``None`` on the EP-only path (PP=1, EP>1) until the EP
-    # branch resolves it, so it must never be read before being assigned.
-    layer_module = None
-
     # PP
     pp_group = _get_pp_group(models)
     if "layers." in param_name and get_pg_size(pp_group) > 1:
@@ -295,25 +360,13 @@ def _megatron_local_name_to_global(
     is_expert_param = (is_grouped_expert_param or is_local_expert_param) and ".adapter." not in param_name
     ep_group = _get_ep_group(models) if is_expert_param else None
     if is_expert_param and ep_group is not None and get_pg_size(ep_group) > 1:
-        # Resolve the layer module independently of the PP branch above: on the
-        # EP-only path (PP=1, EP>1) that branch never runs, so ``layer_module``
-        # is still ``None`` here. Prefer the per-layer module count
-        # (heterogeneous-safe); fall back to the config-derived count when the
-        # model isn't available (e.g. name-only conversions in tests).
-        if layer_module is None:
-            layer_match = re.match(r"^(.+?\.layers\.\d+)", param_name)
-            if models is not None and layer_match is not None:
-                try:
-                    _, layer_module = get_module_and_param_from_name(
-                        models=models, param_name=layer_match.group(1), vp_stage=vp_stage
-                    )
-                except (ValueError, AttributeError):
-                    layer_module = None
-
-        if isinstance(layer_module, MegatronModule):
-            num_experts_per_rank = layer_module.mlp.num_local_experts  # per-layer, heterogeneous-safe
-        else:
-            num_experts_per_rank = config.num_moe_experts // get_pg_size(ep_group)
+        # Expert numbering follows the explicit config contract (scalar
+        # ``num_moe_experts`` for homogeneous models, ``num_moe_experts_per_layer``
+        # for heterogeneous ones). ``param_name`` is already in global layer
+        # numbering because the PP branch above runs first. This is intentionally
+        # independent of the instantiated module so conversion does not depend on
+        # runtime layout (PP/VPP, MTP, custom MoE wrappers) or on a live model.
+        num_experts_per_rank = _resolve_num_local_experts(config, param_name, get_pg_size(ep_group))
 
         def _update_grouped_expert_number(param_name: str, param_type: str) -> str:
             """Update expert number from local to global for weight or bias parameters."""

src/megatron/bridge/models/gpt_provider.py

@@ -181,6 +181,25 @@ class GPTModelProvider(TransformerConfig, ModelProviderMixin[MCoreGPTModel]):
 
     # MoE / FP8
     num_moe_experts: Optional[int] = None
+    num_moe_experts_per_layer: Optional[list[int]] = None
+    """Explicit per-layer expert layout for heterogeneous MoE models.
+
+    ``None`` (the default) means the model is homogeneous and every MoE layer
+    uses the scalar :attr:`num_moe_experts`. When set, it must be a list of
+    length ``num_layers`` giving the *global* number of routed experts for each
+    decoder layer, indexed by global layer number; use ``0`` for dense (non-MoE)
+    layers, and every non-zero entry must be divisible by the expert-parallel
+    size.
+
+    This is the authoritative, config-level contract for expert layout: it is
+    the layout the model is built from and the layout checkpoint conversion
+    (HF <-> Megatron) uses to renumber experts across expert-parallel ranks.
+    Conversion therefore never inspects an instantiated module to recover
+    per-layer expert counts, keeping name mapping decoupled from runtime layout
+    (PP/VPP, MTP, custom MoE wrappers) and from whether a live model exists.
+    Multi-token-prediction layers (``mtp.layers.*``) are not decoder layers and
+    are not covered by this list; they follow the scalar :attr:`num_moe_experts`.
+    """
     moe_grouped_gemm: bool = False
     qk_layernorm: bool = False
     fp8: Optional[str] = None

tests/unit_tests/models/test_quant_mapping.py

@@ -856,7 +856,9 @@ def test_regex_matches_expert_bias_indices(self, name):
         assert re.search(r"\.bias\d+$", name) is not None
         assert re.search(r"\.weight\d+$", name) is None
 
-    def _call_local_to_global(self, param_name, num_moe_experts=4, ep_size=2, ep_rank=0):
+    def _call_local_to_global(
+        self, param_name, num_moe_experts=4, ep_size=2, ep_rank=0, num_moe_experts_per_layer=None
+    ):
         """Invoke the real `_megatron_local_name_to_global` with parallel_state mocked.
 
         Mocks ``parallel_state.get_expert_model_parallel_group`` and
@@ -870,7 +872,7 @@ def _call_local_to_global(self, param_name, num_moe_experts=4, ep_size=2, ep_ran
         ep_group.rank.return_value = ep_rank
         pp_group = MagicMock()
 
-        config = SimpleNamespace(num_moe_experts=num_moe_experts)
+        config = SimpleNamespace(num_moe_experts=num_moe_experts, num_moe_experts_per_layer=num_moe_experts_per_layer)
 
         def _pg_size(group):
             if group is ep_group:
@@ -911,3 +913,85 @@ def test_expert_bias_renumbered_under_ep(self):
             ep_rank=1,
         )
         assert out == "decoder.layers.0.mlp.experts.linear_fc1.bias3"
+
+
+class TestHeterogeneousExpertLayout(TestEPRenumberRegexSkipsAmax):
+    """The per-layer expert-count contract (`num_moe_experts_per_layer`).
+
+    Conversion must read the expert layout from the model config, not from an
+    instantiated module. Homogeneous models keep using the scalar
+    ``num_moe_experts``; heterogeneous models declare a per-decoder-layer list of
+    global expert counts and conversion renumbers each layer with its own count.
+    """
+
+    def test_homogeneous_falls_back_to_scalar(self):
+        """With the per-layer list unset, every layer uses the scalar count."""
+        out = self._call_local_to_global(
+            "decoder.layers.3.mlp.experts.linear_fc1.weight0",
+            num_moe_experts=4,
+            ep_size=2,
+            ep_rank=1,
+            num_moe_experts_per_layer=None,
+        )
+        assert out == "decoder.layers.3.mlp.experts.linear_fc1.weight2"
+
+    def test_heterogeneous_uses_per_layer_count(self):
+        """Different layers renumber with their own global expert count.
+
+        Layer 1 has 8 experts (4 per rank on EP=2); layer 2 has 4 experts (2 per
+        rank). The same local ``weight0`` therefore maps to different globals.
+        """
+        per_layer = [0, 8, 4, 8]  # layer 0 dense, others MoE with differing counts
+        out_layer1 = self._call_local_to_global(
+            "decoder.layers.1.mlp.experts.linear_fc1.weight0",
+            ep_size=2,
+            ep_rank=1,
+            num_moe_experts_per_layer=per_layer,
+        )
+        assert out_layer1 == "decoder.layers.1.mlp.experts.linear_fc1.weight4"
+
+        out_layer2 = self._call_local_to_global(
+            "decoder.layers.2.mlp.experts.linear_fc1.weight0",
+            ep_size=2,
+            ep_rank=1,
+            num_moe_experts_per_layer=per_layer,
+        )
+        assert out_layer2 == "decoder.layers.2.mlp.experts.linear_fc1.weight2"
+
+    def test_mtp_layer_uses_scalar_not_per_layer_list(self):
+        """MTP layers are outside the per-decoder-layer contract; they use the scalar."""
+        out = self._call_local_to_global(
+            "mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight0",
+            num_moe_experts=4,
+            ep_size=2,
+            ep_rank=1,
+            num_moe_experts_per_layer=[8, 8],
+        )
+        assert out == "mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight2"
+
+    def test_per_layer_list_too_short_raises(self):
+        """A list that does not cover the resolved layer is a hard error."""
+        with pytest.raises(ValueError, match="must cover every decoder layer"):
+            self._call_local_to_global(
+                "decoder.layers.5.mlp.experts.linear_fc1.weight0",
+                ep_size=2,
+                num_moe_experts_per_layer=[8, 8],
+            )
+
+    def test_expert_param_on_dense_layer_raises(self):
+        """An expert parameter on a layer declared dense (count 0) is a hard error."""
+        with pytest.raises(ValueError, match="declared dense"):
+            self._call_local_to_global(
+                "decoder.layers.0.mlp.experts.linear_fc1.weight0",
+                ep_size=2,
+                num_moe_experts_per_layer=[0, 8],
+            )
+
+    def test_count_not_divisible_by_ep_raises(self):
+        """A global expert count not divisible by EP size is a hard error."""
+        with pytest.raises(ValueError, match="not divisible by expert-parallel size"):
+            self._call_local_to_global(
+                "decoder.layers.1.mlp.experts.linear_fc1.weight0",
+                ep_size=4,
+                num_moe_experts_per_layer=[0, 6],
+            )