[OMNIML-6076] [13/21] Describe expert parallelism in the model specs

modelopt/torch/models/__init__.py

@@ -33,6 +33,7 @@
 
 from .moe import is_moe
 from .specs import (
+    EPSpec,
     ExportSpec,
     ModelSpec,
     MoESpec,
@@ -52,6 +53,7 @@
 # cannot drift silently: an entry that no longer exists fails at import, and one that is
 # no longer used fails lint.
 __all__ = [
+    "EPSpec",
     "ExportSpec",
     "MoESpec",
     "ModelSpec",
@@ -87,6 +89,9 @@
     phimoe,
     nemotron,
     nemotron_h,
+    minimax_m3_vl,
+    llama4_text,
+    nemotron_h_omni,
     qwen2_moe,
     qwen3,
     qwen3_5_moe,

modelopt/torch/models/deepseek_v3/specs.py

@@ -19,7 +19,7 @@
 ``DeepseekMoE`` block of DeepSeek-MoE is a separate spec under ``deepseek``.
 """
 
-from ..specs import ExportSpec, ModelSpec, MoESpec, register
+from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register
 
 __all__: list[str] = []
 
@@ -37,6 +37,16 @@
         model_type="deepseek_v3",
         min_transformers_version="5.5",
         export_spec=ExportSpec(grouped_expert_export=True),
+        ep_spec=EPSpec(
+            # Stock transformers ships no plan for deepseek_v3. No gate entry on purpose: the
+            # router is logits-only and the MoE block computes GLOBAL top-k ids, so the experts
+            # already receive unmasked global ids -- what the all-to-all dispatch wants.
+            base_model_ep_plan={
+                "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+                "layers.*.mlp.experts.down_proj": "grouped_gemm",
+                "layers.*.mlp.experts": "moe_tp_experts",
+            },
+        ),
         moe_spec=MoESpec(
             block_names=("DeepseekV3MoE",),
             expert_linear_names=("gate_proj", "down_proj", "up_proj"),

modelopt/torch/models/gpt_oss/specs.py

@@ -15,14 +15,21 @@
 
 """GPT-OSS specs (HF model type ``gpt_oss``)."""
 
-from ..specs import ModelSpec, MoESpec, register
+from ..specs import EPSpec, ModelSpec, MoESpec, register
 
 __all__: list[str] = []
 
 register(
     ModelSpec(
         model_type="gpt_oss",
         min_transformers_version="5.5",
+        ep_spec=EPSpec(
+            unsupported=(
+                "GPT-OSS stores its experts transposed, (num_experts, in, out), with gate and up "
+                "interleaved in gate_up_proj, and the expert-parallel export cannot write that "
+                "layout."
+            ),
+        ),
         moe_spec=MoESpec(
             # GPT-OSS fuses gate and up into a single gate_up_proj.
             # transformers names the block GptOssMLP; GptOssMoE is kept for the

modelopt/torch/models/llama4/specs.py

@@ -15,10 +15,18 @@
 
 """Llama4 specs (HF model type ``llama4``)."""
 
-from ..specs import ModelSpec, MoESpec, register
+from ..specs import EPSpec, ModelSpec, MoESpec, register
 
 __all__: list[str] = []
 
+# transformers ships a base_model_ep_plan for Llama-4, which alone would route it into the
+# expert-parallel path.
+EP_UNSUPPORTED = (
+    "Llama-4's MoE block routes tokens itself and runs every expert densely, rather than "
+    "dispatching its top-k through the experts interface, and its experts are stored transposed, "
+    "(num_experts, in, out), which the expert-parallel export cannot write."
+)
+
 # Llama4TextExperts is fused: one module holding 3-D gate_up_proj and down_proj
 # parameters, run through torch.bmm. There is no (gate, up) pair left to fuse, and the
 # grouped-export path (get_experts_list) does not apply. The block also carries a
@@ -27,6 +35,7 @@
     ModelSpec(
         model_type="llama4",
         min_transformers_version="5.5",
+        ep_spec=EPSpec(unsupported=EP_UNSUPPORTED),
         moe_spec=MoESpec(
             block_names=("Llama4TextMoe",),
             expert_linear_names=("gate_up_proj", "down_proj"),

modelopt/torch/models/llama4_text/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""llama4_text specs (HF model type ``llama4_text``)."""
+
+from . import specs

modelopt/torch/models/llama4_text/specs.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Llama-4 text-only specs (HF model type ``llama4_text``).
+
+The same architecture as ``llama4`` reached through ``Llama4ForCausalLM``, with its own
+``base_model_ep_plan`` on the root config, so it is refused expert parallelism the same way.
+"""
+
+from ..llama4.specs import EP_UNSUPPORTED
+from ..specs import EPSpec, ModelSpec, register
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type="llama4_text",
+        min_transformers_version="5.5",
+        ep_spec=EPSpec(unsupported=EP_UNSUPPORTED),
+    )
+)

modelopt/torch/models/minimax_m3_vl/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""minimax_m3_vl specs (HF model type ``minimax_m3_vl``)."""
+
+from . import specs

modelopt/torch/models/minimax_m3_vl/specs.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MiniMax-M3 specs (HF model type ``minimax_m3_vl``).
+
+transformers ships everything M3 needs to run expert-parallel: its MoE block calls the experts with
+the router's top-k, its experts dispatch through the experts interface, and its text config carries a
+``base_model_ep_plan``. The generic all-to-all transport takes over unchanged.
+"""
+
+from ..specs import EPSpec, ModelSpec, register
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type="minimax_m3_vl",
+        min_transformers_version="5.13",
+        ep_spec=EPSpec(requires_transformers="5.13"),
+    )
+)

modelopt/torch/models/nemotron_h/specs.py

@@ -15,15 +15,26 @@
 
 """Nemotron-H specs (HF model type ``nemotron_h``)."""
 
-from ..specs import ExportSpec, ModelSpec, MoESpec, register
+from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register
 
 __all__: list[str] = []
 
+# Ungated experts, so ``up_proj``/``down_proj`` rather than a gate/up pair. The gate emits logits and
+# the MoE block runs its own group-limited top-k, so the plan has no router entry and the experts
+# receive global expert ids -- which the expert-parallel all-to-all experts forward routes on. The
+# loader maps ``nemotron_experts_ep``, a name no transformers style carries, to that forward.
+BASE_MODEL_EP_PLAN = {
+    "layers.*.mixer.experts.up_proj": "grouped_gemm",
+    "layers.*.mixer.experts.down_proj": "grouped_gemm",
+    "layers.*.mixer.experts": "nemotron_experts_ep",
+}
+
 register(
     ModelSpec(
         model_type="nemotron_h",
         min_transformers_version="5.5",
         export_spec=ExportSpec(grouped_expert_export=True),
+        ep_spec=EPSpec(base_model_ep_plan=BASE_MODEL_EP_PLAN),
         moe_spec=MoESpec(
             # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate).
             block_names=("NemotronHMOE",),

modelopt/torch/models/nemotron_h_omni/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""NemotronH-Omni specs (HF model type ``nemotron_h_omni``)."""
+
+from . import specs

modelopt/torch/models/nemotron_h_omni/specs.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""NemotronH-Omni specs (HF model type ``nemotron_h_omni``).
+
+The Omni VLM's language sub-model is a native NemotronH, so its MoE facts are NemotronH's.
+What differs is where the expert-parallel plan has to be attached: on the nested ``llm_config``,
+not the VLM root. transformers aggregates each sub-model's plan up to the root WITH its module
+prefix, so a plan on ``llm_config`` surfaces as ``language_model.<base>.layers.*`` and matches the
+real parameter paths -- while a plan on the top config reaches no expert at all.
+"""
+
+from ..nemotron_h.specs import BASE_MODEL_EP_PLAN
+from ..specs import EPSpec, ModelSpec, register
+
+__all__: list[str] = []
+
+register(
+    ModelSpec(
+        model_type="nemotron_h_omni",
+        # transformers ships the Omni VLM itself from 5.18. Expert parallelism still runs on the
+        # checkpoint's own code (``trust_remote_code=True``), whose sub-model config is the
+        # ``llm_config`` the plan below targets; transformers' port names it ``text_config``.
+        min_transformers_version="5.18",
+        ep_spec=EPSpec(
+            base_model_ep_plan=BASE_MODEL_EP_PLAN,
+            plan_config_attr="llm_config",
+        ),
+    )
+)

modelopt/torch/models/qwen3_moe/specs.py

@@ -15,14 +15,20 @@
 
 """Qwen3-MoE specs (HF model type ``qwen3_moe``)."""
 
-from ..specs import ExportSpec, ModelSpec, MoESpec, register
+from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register
 
 __all__: list[str] = []
 
 register(
     ModelSpec(
         model_type="qwen3_moe",
         min_transformers_version="5.5",
+        ep_spec=EPSpec(
+            # No plan here: qwen3_moe's own base_model_ep_plan ships in transformers, and it is
+            # the architecture the all-to-all path was built around. Its plan lands in 5.7, later
+            # than the 5.6 EP framework floor, so the model carries the higher requirement.
+            requires_transformers="5.7",
+        ),
         moe_spec=MoESpec(
             block_names=("Qwen3MoeSparseMoeBlock",),
             expert_linear_names=("gate_proj", "down_proj", "up_proj"),