[OMNIML-6076] [1/3] Describe expert parallelism in the model specs

modelopt/torch/models/README.md

@@ -75,6 +75,7 @@ A `ModelSpec` holds one attribute per section, each `None` unless the model fill
 |---|---|---|
 | `moe_spec` (`MoESpec`) | MoE architecture facts only — block classes, expert projection naming, which layout that naming describes | the model is dense |
 | `export_spec` (`ExportSpec`) | Per-model data *and policy* of the unified HF export path, e.g. whether grouped expert export is validated for this model | export needs nothing model-specific |
+| `ep_spec` (`EPSpec`) | Per-model data of the expert-parallel PTQ path: the `base_model_ep_plan` to attach and the config it goes on, the transformers release it needs, or why the model cannot be expert-parallel yet | the model is dense, or stock transformers already ships all of it |
 
 `None` means "the model has nothing to say about this", which stays distinct from a
 filled-in-but-empty section.

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,17 @@
         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 before 5.13; from 5.13 the config's
+            # own plan wins (see apply_distributed_support) and this one goes unused. 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,38 @@
+# 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.
+
+Only that refusal is registered under this type, deliberately: no MoE or export section, so its
+MoE and export lookups find none, as they did before this type had a spec. Giving it ``llama4``'s
+MoE section would change how text-only Llama-4 checkpoints match MoE blocks in quantization and
+export, which belongs with tests of those paths, not with its expert-parallel refusal.
+"""
+
+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,51 @@
+# 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 language-model
+config, not the VLM root. transformers aggregates each sub-model's plan up to the root WITH its
+module prefix, so a plan there surfaces as ``language_model.<base>.layers.*`` and matches the real
+parameter paths -- while a plan on the top config reaches no expert at all.
+
+Only that plan is registered under the root type, so lookups of the MoE and export sections by
+``nemotron_h_omni`` find neither. No export depends on them: on transformers 5 the language
+model's experts are one fused container, which the export writes through its own handler, and
+every step those lookups gate leaves a fused container as it is.
+"""
+
+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. The checkpoint's own code
+        # (``trust_remote_code=True``) names the language-model config ``llm_config``, and
+        # transformers' port ``text_config``: the plan goes on whichever the config has, tried in
+        # the order the load reads them.
+        min_transformers_version="5.18",
+        ep_spec=EPSpec(
+            # Its own copy: the registry is process-wide, and a spec holding NemotronH's dict would
+            # carry any change to one plan into the other.
+            base_model_ep_plan=dict(BASE_MODEL_EP_PLAN),
+            plan_config_attr=("text_config", "llm_config"),
+        ),
+    )
+)

modelopt/torch/models/qwen3_moe/specs.py

@@ -15,14 +15,19 @@
 
 """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.
+            requires_transformers="5.7",
+        ),
         moe_spec=MoESpec(
             block_names=("Qwen3MoeSparseMoeBlock",),
             expert_linear_names=("gate_proj", "down_proj", "up_proj"),