From 0e2561b166d710263a7f00e60a46c6b586db644a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:53 +0000 Subject: [PATCH 01/23] Describe expert parallelism in the model specs The DP=EP path needs three things from a model before it can shard its experts: which modules the routed experts are and how they shard, which transformers release first ships that information, and -- for a VLM -- which nested config the plan has to be attached to. All three are facts about the architecture, so they belong with the model's other facts. So this adds an EPSpec section to ModelSpec, and apply_distributed_support, the one applier that reads it. Most models become data in the spec package they already had: deepseek_v3 carries the plan stock transformers does not ship, while qwen3_moe and gpt_oss carry only the release their own plan needs. An existing plan on the config always wins, since overwriting it would substitute our idea of the layout for the model's. Two things resisted being data and kept code, deliberately. NemotronH needs a custom parallel style, because its gate emits logits only and the group-limited top-k happens inside the MoE block -- so the usual ep_router has nothing to remap and the remap has to move to the experts' input. transformers resolves styles through a global registry rather than from config, so that registration is code, in nemotron_h/ep.py. And EPSpec.prepare is the escape hatch for checkpoints carrying their own modeling code. NemotronH-Omni is registered separately because the difference is real: its experts live in the language sub-model, and transformers aggregates a sub-model's plan up to the root WITH its module prefix -- so a plan on llm_config matches the actual parameter paths while a plan on the VLM root matches nothing at all, which fails as silence rather than as an error. A test pins both the placement and that the two spellings share one plan. Signed-off-by: Shengliang Xu --- modelopt/torch/models/__init__.py | 3 + modelopt/torch/models/deepseek_v3/specs.py | 12 +- modelopt/torch/models/gpt_oss/specs.py | 8 +- modelopt/torch/models/nemotron_h/ep.py | 84 +++++++++ modelopt/torch/models/nemotron_h/specs.py | 4 +- .../torch/models/nemotron_h_omni/__init__.py | 19 +++ .../torch/models/nemotron_h_omni/specs.py | 43 +++++ modelopt/torch/models/qwen3_moe/specs.py | 8 +- modelopt/torch/models/specs.py | 59 +++++++ modelopt/torch/utils/moe.py | 56 ++++++ tests/unit/torch/models/test_ep_specs.py | 161 ++++++++++++++++++ 11 files changed, 453 insertions(+), 4 deletions(-) create mode 100644 modelopt/torch/models/nemotron_h/ep.py create mode 100644 modelopt/torch/models/nemotron_h_omni/__init__.py create mode 100644 modelopt/torch/models/nemotron_h_omni/specs.py create mode 100644 tests/unit/torch/models/test_ep_specs.py diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index 98aaf9e57da..a9d106d36c2 100644 --- a/modelopt/torch/models/__init__.py +++ b/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,7 @@ phimoe, nemotron, nemotron_h, + nemotron_h_omni, qwen2_moe, qwen3, qwen3_5_moe, diff --git a/modelopt/torch/models/deepseek_v3/specs.py b/modelopt/torch/models/deepseek_v3/specs.py index 3a774ecb931..dab92dab2eb 100644 --- a/modelopt/torch/models/deepseek_v3/specs.py +++ b/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"), diff --git a/modelopt/torch/models/gpt_oss/specs.py b/modelopt/torch/models/gpt_oss/specs.py index be8419781c7..ae9c605994e 100644 --- a/modelopt/torch/models/gpt_oss/specs.py +++ b/modelopt/torch/models/gpt_oss/specs.py @@ -15,7 +15,7 @@ """GPT-OSS specs (HF model type ``gpt_oss``).""" -from ..specs import ModelSpec, MoESpec, register +from ..specs import EPSpec, ModelSpec, MoESpec, register __all__: list[str] = [] @@ -23,6 +23,12 @@ ModelSpec( model_type="gpt_oss", min_transformers_version="5.5", + ep_spec=EPSpec( + # Rides the generic transport unmodified: transformers ships the plan, and GptOssExperts + # already dispatches through ExpertsInterface. Its two specifics -- per-expert biases and + # native mxfp4 weights -- are handled generically in the transport and the loader. + requires_transformers="5.6", + ), 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 diff --git a/modelopt/torch/models/nemotron_h/ep.py b/modelopt/torch/models/nemotron_h/ep.py new file mode 100644 index 00000000000..b10d58cfee9 --- /dev/null +++ b/modelopt/torch/models/nemotron_h/ep.py @@ -0,0 +1,84 @@ +# 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. + +"""The expert-parallel style NemotronH's block-routed MoE needs. + +Everything else in this package is data; this is the one piece that has to be code, because +transformers resolves parallel styles through a global registry rather than from a config. +""" + +_REGISTERED = False + + +def _register_style() -> None: + """Register ``NemotronHExpertsEP`` under the name the spec's plan refers to. + + Imported lazily so this package stays importable without torch or transformers, and + idempotent because a PTQ run may prepare several models in one process. + """ + global _REGISTERED + if _REGISTERED: + return + + import torch + from transformers.integrations.tensor_parallel import ( + ALL_PARALLEL_STYLES, + MoeTensorParalellExperts, + all_reduce_forward, + ) + + class NemotronHExpertsEP(MoeTensorParalellExperts): + """Expert parallelism for a gate that emits logits only. + + The usual ``ep_router`` style remaps expert indices on the ROUTER's output, which + requires the router to return ``(scores, indices)``. NemotronH's gate returns logits + and the group-limited top-k happens inside the MoE block, so there is nothing for + ``ep_router`` to remap and the experts would receive global indices they cannot use. + + So the remap happens at the EXPERTS' input instead: a token routed to one of this + rank's local experts gets a local index, and a token belonging elsewhere is pointed at + local index 0 with weight zero, which makes the experts' ``index_add`` a no-op for it. + ``grouped_gemm`` has already set ``num_experts`` to the local count, and the inherited + ``all_reduce_forward`` sums each rank's partial expert output. + """ + + def _prepare_input_fn(self, mod, inputs, device_mesh): + hidden, top_k_index, top_k_weights = inputs[0], inputs[1], inputs[2] + local = mod.num_experts # local count, after grouped_gemm shards dim 0 + start = device_mesh.get_local_rank() * local + is_local = (top_k_index >= start) & (top_k_index < start + local) + local_index = torch.where(is_local, top_k_index - start, torch.zeros_like(top_k_index)) + local_weights = torch.where(is_local, top_k_weights, torch.zeros_like(top_k_weights)) + return (hidden, local_index, local_weights) + + def _prepare_output_fn(self, mod, outputs, device_mesh): + return all_reduce_forward(outputs, device_mesh) + + ALL_PARALLEL_STYLES.register("nemotron_experts_ep", NemotronHExpertsEP()) + _REGISTERED = True + + +def prepare(hf_config, **_): + """Register the style the plan names, then hand the config back unchanged.""" + _register_style() + return hf_config + + +BASE_MODEL_EP_PLAN = { + "layers.*.mixer.experts.up_proj": "grouped_gemm", + "layers.*.mixer.experts.down_proj": "grouped_gemm", + "layers.*.mixer.experts": "nemotron_experts_ep", +} +"""Ungated experts, so ``up_proj``/``down_proj`` rather than a gate/up pair.""" diff --git a/modelopt/torch/models/nemotron_h/specs.py b/modelopt/torch/models/nemotron_h/specs.py index 7198c2b31b3..dd7db03a5b6 100644 --- a/modelopt/torch/models/nemotron_h/specs.py +++ b/modelopt/torch/models/nemotron_h/specs.py @@ -15,7 +15,8 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" -from ..specs import ExportSpec, ModelSpec, MoESpec, register +from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register +from .ep import BASE_MODEL_EP_PLAN, prepare __all__: list[str] = [] @@ -24,6 +25,7 @@ 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, prepare=prepare), moe_spec=MoESpec( # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). block_names=("NemotronHMOE",), diff --git a/modelopt/torch/models/nemotron_h_omni/__init__.py b/modelopt/torch/models/nemotron_h_omni/__init__.py new file mode 100644 index 00000000000..5b9b820672c --- /dev/null +++ b/modelopt/torch/models/nemotron_h_omni/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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 diff --git a/modelopt/torch/models/nemotron_h_omni/specs.py b/modelopt/torch/models/nemotron_h_omni/specs.py new file mode 100644 index 00000000000..e022f5b8430 --- /dev/null +++ b/modelopt/torch/models/nemotron_h_omni/specs.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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..layers.*`` and matches the +real parameter paths -- while a plan on the top config reaches no expert at all. +""" + +from ..nemotron_h.ep import BASE_MODEL_EP_PLAN, prepare +from ..specs import EPSpec, ModelSpec, register + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="nemotron_h_omni", + # The Omni VLM ships its own modeling code with the checkpoint; only its language + # sub-model is the transformers-native NemotronH. + modeling_source="remote_code", + ep_spec=EPSpec( + base_model_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="llm_config", + prepare=prepare, + ), + ) +) diff --git a/modelopt/torch/models/qwen3_moe/specs.py b/modelopt/torch/models/qwen3_moe/specs.py index c8b9bcd3e5b..fadee69a880 100644 --- a/modelopt/torch/models/qwen3_moe/specs.py +++ b/modelopt/torch/models/qwen3_moe/specs.py @@ -15,7 +15,7 @@ """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] = [] @@ -23,6 +23,12 @@ 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"), diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 0a203780b7f..0d881644958 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -44,6 +44,9 @@ from typing import TYPE_CHECKING, Literal, get_args, get_type_hints if TYPE_CHECKING: + from collections.abc import Callable + from typing import Any + import torch.nn as nn __all__ = [ @@ -155,6 +158,58 @@ def expert_linear_names_for(self, module, fused: bool | None = None) -> tuple[st return self.expert_linear_names +@dataclass(kw_only=True) +class EPSpec(SpecSection): + """Subsystem section: per-model data of the expert-parallel (DP=EP) PTQ path. + + Calibrating a large MoE data-parallel with its experts sharded needs three things from + the model, and all three are facts about the architecture rather than about any run: + which modules are the routed experts and how they shard, which transformers release + first ships that information, and -- for a VLM -- which nested config carries it. + + Most models need only the data below. A checkpoint whose modeling code it ships itself + may need actual code run against it before it can be loaded expert-parallel (patching a + remote-code class, decompressing packed experts); that is what ``prepare`` is for, and it + is deliberately the exception rather than the interface. + """ + + base_model_ep_plan: dict[str, str] | None = None + """The ``base_model_ep_plan`` transformers needs to shard this model's experts, applied + to the config when it does not already carry one. + + Stock transformers ships a plan for some MoE models and not others; where it does, leave + this ``None`` and let the model's own plan win. Entries map a parameter glob to a + parallel style, e.g. ``{"layers.*.mlp.experts.gate_up_proj": "grouped_gemm"}``. + + A router entry is not always wanted. DeepSeek-V3's router is logits-only and its MoE + block computes global top-k ids, so the experts already receive unmasked global ids -- + which is exactly what the all-to-all dispatch needs, and adding a gate entry would + mask them.""" + + requires_transformers: str | None = None + """Earliest transformers release that supports THIS model expert-parallel, when it is + later than the framework floor. + + Distinct from ``ModelSpec.min_transformers_version``, which is about the modeling code + matching the spec. A model can be fully modelled and still not be shardable: qwen3_moe's + own ``base_model_ep_plan`` lands in 5.7, while the shared EP framework is 5.6.""" + + plan_config_attr: str | None = None + """The nested config attribute carrying the expert-parallel plan, for a VLM whose MoE is + its text sub-model (``"text_config"``, ``"llm_config"``). ``None`` means the root config. + + The plan has to sit on whichever config describes the module tree the globs are written + against, so a plan on the wrong config silently matches nothing.""" + + prepare: "Callable[..., Any] | None" = None + """Escape hatch for a model that cannot be made expert-parallel by data alone. + + Called as ``prepare(hf_config, ckpt_path=..., trust_remote_code=..., keep_compressed=...)`` + and returns the config to load. Reserved for checkpoints carrying their own modeling code, + where the work is patching classes or installing on-the-fly weight decompression. Runs + after the declarative fields above have been applied.""" + + @dataclass(kw_only=True) class ExportSpec(SpecSection): """Subsystem section: per-model data of the unified HF export path. @@ -233,6 +288,10 @@ class ModelSpec: """Per-model data of the unified HF export path, or ``None`` when export needs nothing model-specific.""" + ep_spec: EPSpec | None = None + """Per-model data of the expert-parallel PTQ path, or ``None`` when the model is dense + or needs nothing beyond what stock transformers already ships.""" + _SPECS: dict[str, ModelSpec] = {} diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 7ded9dd8eb9..2a10951d746 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -141,6 +141,62 @@ def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict } +def apply_distributed_support( + model_type: str, + hf_config, + *, + ckpt_path: "str | None" = None, + trust_remote_code: bool = False, + keep_compressed: bool = False, +): + """Prepare ``hf_config`` so ``model_type`` can be loaded expert-parallel, from its model spec. + + Reads the model's :class:`~modelopt.torch.models.EPSpec` and applies it: enforce the + transformers floor this model's expert parallelism needs, then attach the + ``base_model_ep_plan`` it requires to whichever config carries its MoE -- the root for a plain + MoE, the nested text/LLM config for a VLM whose experts live in its language sub-model, since + the plan's globs are written against that module tree and a plan on the wrong config matches + nothing at all rather than failing. + + An existing plan always wins: stock transformers ships one for some models, and overwriting it + would substitute our idea of the layout for the model's own. + + Returns the config to load, unchanged for a model type with no spec or no EP section. The + generic all-to-all transport is not wired here -- it self-registers in + :func:`load_and_prepare_ep`, so nothing about it is per-model. + """ + from modelopt.torch.models import get_spec + + spec = get_spec(model_type) + ep = getattr(spec, "ep_spec", None) if spec is not None else None + if ep is None: + return hf_config + + if ep.requires_transformers: + require_transformers(ep.requires_transformers, f"{model_type} expert-parallel PTQ") + + if ep.base_model_ep_plan: + target = hf_config + if ep.plan_config_attr: + target = getattr(hf_config, ep.plan_config_attr, None) + if target is None: + raise ValueError( + f"{model_type} expects its expert-parallel plan on " + f"config.{ep.plan_config_attr}, which this checkpoint's config does not have." + ) + if not getattr(target, "base_model_ep_plan", None): + target.base_model_ep_plan = dict(ep.base_model_ep_plan) + + if ep.prepare is not None: + hf_config = ep.prepare( + hf_config, + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + keep_compressed=keep_compressed, + ) + return hf_config + + # --------------------------------------------------------------------------------------------------- # DP=EP all-to-all expert transport (`torch_a2a`) # diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py new file mode 100644 index 00000000000..6044432814a --- /dev/null +++ b/tests/unit/torch/models/test_ep_specs.py @@ -0,0 +1,161 @@ +# 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. + +"""The expert-parallel section of the model specs, and the applier that reads it.""" + +from types import SimpleNamespace + +import pytest + +from modelopt.torch.models import EPSpec, ModelSpec, get_spec, get_specs +from modelopt.torch.utils.moe import apply_distributed_support + +EP_MODEL_TYPES = [s.model_type for s in get_specs() if s.ep_spec is not None] + + +def cfg(**kw): + return SimpleNamespace(**kw) + + +# --------------------------------------------------------------------------- the registered specs + + +def test_some_model_declares_expert_parallel_support(): + """A guard on the wiring itself: an empty list here means no spec reached the registry.""" + assert EP_MODEL_TYPES + + +@pytest.mark.parametrize("model_type", EP_MODEL_TYPES) +def test_every_ep_plan_is_well_formed(model_type): + """Globs map to style names; a malformed entry matches nothing and fails silently at load.""" + plan = get_spec(model_type).ep_spec.base_model_ep_plan + if plan is None: + return + assert plan, f"{model_type} declares an empty plan, which is not the same as declaring none" + for glob, style in plan.items(): + assert isinstance(glob, str) and glob, f"{model_type}: bad glob {glob!r}" + assert isinstance(style, str) and style, f"{model_type}: bad style for {glob!r}" + + +@pytest.mark.parametrize("model_type", EP_MODEL_TYPES) +def test_a_spec_declares_either_a_plan_a_floor_or_a_prepare(model_type): + """An EP section that says nothing would claim support the model does not have.""" + ep = get_spec(model_type).ep_spec + assert ep.base_model_ep_plan or ep.requires_transformers or ep.prepare, ( + f"{model_type} has an EP section that declares nothing" + ) + + +def test_the_omni_vlm_reuses_the_text_models_plan(): + """Two spellings of one architecture must not drift into two different plans.""" + assert ( + get_spec("nemotron_h_omni").ep_spec.base_model_ep_plan + == get_spec("nemotron_h").ep_spec.base_model_ep_plan + ) + + +def test_the_omni_vlm_puts_its_plan_on_the_nested_config(): + """Its experts live in the language sub-model, and the globs are written against that tree.""" + assert get_spec("nemotron_h_omni").ep_spec.plan_config_attr == "llm_config" + assert get_spec("nemotron_h").ep_spec.plan_config_attr is None + + +# ----------------------------------------------------------------------------------- the applier + + +def test_a_model_type_with_no_spec_is_left_alone(): + config = cfg() + assert apply_distributed_support("not_a_real_model", config) is config + assert not hasattr(config, "base_model_ep_plan") + + +def test_the_plan_is_attached_when_the_config_has_none(): + config = cfg() + apply_distributed_support("deepseek_v3", config) + assert config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +def test_the_models_own_plan_always_wins(): + """Stock transformers ships a plan for some models; overwriting it would substitute our idea + of the layout for the model's own.""" + config = cfg(base_model_ep_plan={"layers.*.mlp.experts": "stock_style"}) + apply_distributed_support("deepseek_v3", config) + assert config.base_model_ep_plan == {"layers.*.mlp.experts": "stock_style"} + + +def test_the_attached_plan_is_a_copy(): + """Two models prepared in one process must not share one mutable dict.""" + a, b = cfg(), cfg() + apply_distributed_support("deepseek_v3", a) + apply_distributed_support("deepseek_v3", b) + a.base_model_ep_plan["injected"] = "oops" + assert "injected" not in b.base_model_ep_plan + assert "injected" not in get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +def test_a_vlm_plan_lands_on_the_nested_config_only(): + """On the root config the globs match no parameter, and the failure is silence, not an error.""" + config = cfg(llm_config=cfg()) + apply_distributed_support("nemotron_h_omni", config) + assert config.llm_config.base_model_ep_plan + assert not hasattr(config, "base_model_ep_plan") + + +def test_a_missing_nested_config_is_an_error_not_a_silent_skip(): + with pytest.raises(ValueError, match="llm_config"): + apply_distributed_support("nemotron_h_omni", cfg()) + + +def test_prepare_runs_and_its_config_is_returned(monkeypatch): + seen = {} + + def prepare(hf_config, **kwargs): + seen.update(kwargs) + seen["config"] = hf_config + return "replaced" + + spec = ModelSpec(model_type="_probe", ep_spec=EPSpec(prepare=prepare)) + monkeypatch.setattr( + "modelopt.torch.models.get_spec", lambda mt: spec if mt == "_probe" else None + ) + + config = cfg() + assert ( + apply_distributed_support("_probe", config, ckpt_path="/c", keep_compressed=True) + == "replaced" + ) + assert seen["config"] is config + assert seen["ckpt_path"] == "/c" + assert seen["keep_compressed"] is True + + +def test_prepare_runs_after_the_plan_is_attached(monkeypatch): + """A prepare hook may inspect or adjust the plan, so it must not run before there is one.""" + observed = {} + + def prepare(hf_config, **_): + observed["plan_present"] = hasattr(hf_config, "base_model_ep_plan") + return hf_config + + spec = ModelSpec( + model_type="_probe", + ep_spec=EPSpec(base_model_ep_plan={"layers.*.experts": "grouped_gemm"}, prepare=prepare), + ) + monkeypatch.setattr( + "modelopt.torch.models.get_spec", lambda mt: spec if mt == "_probe" else None + ) + + apply_distributed_support("_probe", cfg()) + assert observed["plan_present"] From d9732e7f4d6adb7beeb0fbce905adfd6e5a02ca2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 03:17:07 +0000 Subject: [PATCH 02/23] Stop threading keep_compressed through apply_distributed_support keep_compressed asked a model's prepare hook to keep its experts packed, for a frozen speculative-decoding target. PTQ always needs the decompressed weights and hf_ptq never set it, so the knob and its pass-through go. prepare now receives ckpt_path and trust_remote_code only, which the test pins. Signed-off-by: Shengliang Xu --- modelopt/torch/models/specs.py | 8 ++++---- modelopt/torch/utils/moe.py | 8 +------- tests/unit/torch/models/test_ep_specs.py | 6 ++---- 3 files changed, 7 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 0d881644958..7aba26e7ce8 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -204,10 +204,10 @@ class EPSpec(SpecSection): prepare: "Callable[..., Any] | None" = None """Escape hatch for a model that cannot be made expert-parallel by data alone. - Called as ``prepare(hf_config, ckpt_path=..., trust_remote_code=..., keep_compressed=...)`` - and returns the config to load. Reserved for checkpoints carrying their own modeling code, - where the work is patching classes or installing on-the-fly weight decompression. Runs - after the declarative fields above have been applied.""" + Called as ``prepare(hf_config, ckpt_path=..., trust_remote_code=...)`` and returns the + config to load. Reserved for checkpoints carrying their own modeling code, where the work is + patching classes or installing on-the-fly weight decompression. Runs after the declarative + fields above have been applied.""" @dataclass(kw_only=True) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 2a10951d746..469283d869e 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -147,7 +147,6 @@ def apply_distributed_support( *, ckpt_path: "str | None" = None, trust_remote_code: bool = False, - keep_compressed: bool = False, ): """Prepare ``hf_config`` so ``model_type`` can be loaded expert-parallel, from its model spec. @@ -188,12 +187,7 @@ def apply_distributed_support( target.base_model_ep_plan = dict(ep.base_model_ep_plan) if ep.prepare is not None: - hf_config = ep.prepare( - hf_config, - ckpt_path=ckpt_path, - trust_remote_code=trust_remote_code, - keep_compressed=keep_compressed, - ) + hf_config = ep.prepare(hf_config, ckpt_path=ckpt_path, trust_remote_code=trust_remote_code) return hf_config diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 6044432814a..539ad56f229 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -133,12 +133,10 @@ def prepare(hf_config, **kwargs): config = cfg() assert ( - apply_distributed_support("_probe", config, ckpt_path="/c", keep_compressed=True) + apply_distributed_support("_probe", config, ckpt_path="/c", trust_remote_code=True) == "replaced" ) - assert seen["config"] is config - assert seen["ckpt_path"] == "/c" - assert seen["keep_compressed"] is True + assert seen == {"config": config, "ckpt_path": "/c", "trust_remote_code": True} def test_prepare_runs_after_the_plan_is_attached(monkeypatch): From fbebd4a967bca502b1bb289f6ba4d991df1b1276 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:12:47 +0000 Subject: [PATCH 03/23] Skip the EP spec tests without transformers The applier they exercise enforces each model's transformers floor, which imports transformers. The sibling model tests skip without it the same way. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_ep_specs.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 539ad56f229..5cf61d57cc4 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -19,6 +19,8 @@ import pytest +pytest.importorskip("transformers") + from modelopt.torch.models import EPSpec, ModelSpec, get_spec, get_specs from modelopt.torch.utils.moe import apply_distributed_support From 9709b952d6b7721073e40117c58c4efeb5c17b81 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:23:21 +0000 Subject: [PATCH 04/23] Refuse expert-parallel PTQ for GPT-OSS until its experts can be exported transformers ships a base_model_ep_plan for gpt_oss, so a torchrun run took the expert-parallel path and failed only at quantization, minutes into the load. The loader drops GptOssExperts' own quantized wrapper so that the generic grouped-GEMM one takes over, and that one refuses experts stored transposed: the fused-experts export cannot write them. It would split gate and up on the wrong axis, and GPT-OSS also interleaves them. Add EPSpec.unsupported. apply_distributed_support refuses a model that names a reason, before anything is applied or loaded, and points at FSDP2. Mark gpt_oss with it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/gpt_oss/specs.py | 9 +++--- modelopt/torch/models/specs.py | 7 ++++ modelopt/torch/utils/moe.py | 5 +++ tests/unit/torch/models/test_ep_specs.py | 41 ++++++++++++++++++++++-- 4 files changed, 56 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/models/gpt_oss/specs.py b/modelopt/torch/models/gpt_oss/specs.py index ae9c605994e..a4b95c790e9 100644 --- a/modelopt/torch/models/gpt_oss/specs.py +++ b/modelopt/torch/models/gpt_oss/specs.py @@ -24,10 +24,11 @@ model_type="gpt_oss", min_transformers_version="5.5", ep_spec=EPSpec( - # Rides the generic transport unmodified: transformers ships the plan, and GptOssExperts - # already dispatches through ExpertsInterface. Its two specifics -- per-expert biases and - # native mxfp4 weights -- are handled generically in the transport and the loader. - requires_transformers="5.6", + 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. diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 7aba26e7ce8..2e1804095da 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -209,6 +209,13 @@ class EPSpec(SpecSection): patching classes or installing on-the-fly weight decompression. Runs after the declarative fields above have been applied.""" + unsupported: str | None = None + """Why this model cannot be quantized expert-parallel yet, when its config would route it there. + + Stock transformers ships a ``base_model_ep_plan`` for some MoEs the expert-parallel path + cannot handle, and the plan alone makes a model expert-parallel capable. Naming the reason + here refuses the run before any weight is loaded instead of failing deep in quantization.""" + @dataclass(kw_only=True) class ExportSpec(SpecSection): diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 469283d869e..c8b3e155852 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -171,6 +171,11 @@ def apply_distributed_support( if ep is None: return hf_config + if ep.unsupported: + raise NotImplementedError( + f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " + "without expert parallelism, e.g. with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py)." + ) if ep.requires_transformers: require_transformers(ep.requires_transformers, f"{model_type} expert-parallel PTQ") diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 5cf61d57cc4..c52553097c9 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -25,6 +25,7 @@ from modelopt.torch.utils.moe import apply_distributed_support EP_MODEL_TYPES = [s.model_type for s in get_specs() if s.ep_spec is not None] +EP_REFUSED_MODEL_TYPES = [mt for mt in EP_MODEL_TYPES if get_spec(mt).ep_spec.unsupported] def cfg(**kw): @@ -52,14 +53,19 @@ def test_every_ep_plan_is_well_formed(model_type): @pytest.mark.parametrize("model_type", EP_MODEL_TYPES) -def test_a_spec_declares_either_a_plan_a_floor_or_a_prepare(model_type): +def test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal(model_type): """An EP section that says nothing would claim support the model does not have.""" ep = get_spec(model_type).ep_spec - assert ep.base_model_ep_plan or ep.requires_transformers or ep.prepare, ( + assert ep.base_model_ep_plan or ep.requires_transformers or ep.prepare or ep.unsupported, ( f"{model_type} has an EP section that declares nothing" ) +def test_models_whose_own_plan_would_misroute_them_are_refused(): + """GPT-OSS ships a plan of its own, so only its refusal keeps it off the expert-parallel path.""" + assert "gpt_oss" in EP_REFUSED_MODEL_TYPES + + def test_the_omni_vlm_reuses_the_text_models_plan(): """Two spellings of one architecture must not drift into two different plans.""" assert ( @@ -120,6 +126,37 @@ def test_a_missing_nested_config_is_an_error_not_a_silent_skip(): apply_distributed_support("nemotron_h_omni", cfg()) +@pytest.mark.parametrize("model_type", EP_REFUSED_MODEL_TYPES) +def test_an_unsupported_model_is_refused_with_its_config_untouched(model_type): + config = cfg() + + with pytest.raises(NotImplementedError, match="cannot be quantized expert-parallel yet"): + apply_distributed_support(model_type, config) + assert vars(config) == {} + + +def test_a_refusal_comes_before_the_rest_of_the_section(monkeypatch): + """Nothing is attached or patched for a run that is about to be refused.""" + prepared = [] + spec = ModelSpec( + model_type="_probe", + ep_spec=EPSpec( + base_model_ep_plan={"layers.*.experts": "grouped_gemm"}, + prepare=lambda hf_config, **_: prepared.append(hf_config), + unsupported="It is not ready", + ), + ) + monkeypatch.setattr( + "modelopt.torch.models.get_spec", lambda mt: spec if mt == "_probe" else None + ) + config = cfg() + + with pytest.raises(NotImplementedError, match="It is not ready"): + apply_distributed_support("_probe", config) + assert not prepared + assert vars(config) == {} + + def test_prepare_runs_and_its_config_is_returned(monkeypatch): seen = {} From c471bd884aebede66f77f8dec99b0f84961e8517 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 10 Oct 2026 21:22:59 +0000 Subject: [PATCH 05/23] Add the Llama-4 and MiniMax-M3 expert-parallel adapters Per-model expert-parallel support in the same structure as the model specs: one package per HF model type under modelopt/torch/models, data in specs.py, and code in ep.py only where the model genuinely needs it. These two need code. Llama-4 is the one MoE that does not follow the transformers sparse-dispatch convention -- its block routes inline and computes every expert densely, so there is no router-to-experts seam for the transport to wrap; ep.py adapts it to the ExpertsInterface contract instead. llama4_text reuses llama4's adapter and differs only in where the plan lives -- the config itself rather than a nested text config, since there is no VLM wrapper. MiniMax-M3 already follows the convention, with fused [E, out, in] expert params, but hardcodes a per-expert loop in its experts forward, so it only needs its experts class routed through the ExpertsInterface. The declarative half is covered by the spec tests. The prepare hooks patch real transformers classes, so they are exercised by the PTQ runs rather than by unit tests. Signed-off-by: Shengliang Xu --- modelopt/torch/models/__init__.py | 2 + modelopt/torch/models/llama4/ep.py | 128 ++++++++++++++++++ modelopt/torch/models/llama4/specs.py | 13 +- modelopt/torch/models/llama4_text/__init__.py | 18 +++ modelopt/torch/models/llama4_text/specs.py | 33 +++++ .../torch/models/minimax_m3_vl/__init__.py | 18 +++ modelopt/torch/models/minimax_m3_vl/ep.py | 90 ++++++++++++ modelopt/torch/models/minimax_m3_vl/specs.py | 39 ++++++ 8 files changed, 339 insertions(+), 2 deletions(-) create mode 100644 modelopt/torch/models/llama4/ep.py create mode 100644 modelopt/torch/models/llama4_text/__init__.py create mode 100644 modelopt/torch/models/llama4_text/specs.py create mode 100644 modelopt/torch/models/minimax_m3_vl/__init__.py create mode 100644 modelopt/torch/models/minimax_m3_vl/ep.py create mode 100644 modelopt/torch/models/minimax_m3_vl/specs.py diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index a9d106d36c2..7d71f9e106f 100644 --- a/modelopt/torch/models/__init__.py +++ b/modelopt/torch/models/__init__.py @@ -89,6 +89,8 @@ phimoe, nemotron, nemotron_h, + minimax_m3_vl, + llama4_text, nemotron_h_omni, qwen2_moe, qwen3, diff --git a/modelopt/torch/models/llama4/ep.py b/modelopt/torch/models/llama4/ep.py new file mode 100644 index 00000000000..b3264d1d98e --- /dev/null +++ b/modelopt/torch/models/llama4/ep.py @@ -0,0 +1,128 @@ +# 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. + +"""DP=EP a2a support for transformers-native ``llama4`` (Llama-4 Maverick/Scout). + +Unlike every other MoE we run on a2a, Llama-4's ``Llama4TextMoe`` does NOT use the transformers +``ExpertsInterface`` sparse-dispatch convention. Its reference forward: + + router_scores, router_logits = self.router(hidden) # router returns TWO values + routed_in = hidden.repeat(num_experts, 1) * router_scores.T # compute ALL experts, DENSELY + routed_out = self.experts(routed_in) # experts() takes ONE pre-routed tensor + out = shared_expert(hidden) + routed_out.reshape(E, -1, H).sum(0) + +i.e. routing is done inline (not via a wrappable ``RouterParallel``), the experts are computed densely +(every token through every expert, masked by the sigmoid router scores), and the experts module is +called with a single already-routed tensor -- so a2a has no ``router -> experts(hidden, top_k_index, +top_k_weights)`` seam to wrap. (This is why Llama-4 was the one MoE that didn't fit the generic path.) + +This plugin monkey-patches transformers' llama4 modeling to fit the ExpertsInterface contract so the +generic ``torch_a2a`` transport takes over unchanged: + +1. **Decorate ``Llama4TextExperts`` with ``use_experts_implementation``.** Llama-4 stores the expert + weights as ``gate_up_proj`` ``[E, hidden, 2*inter]`` and ``down_proj`` ``[E, inter, hidden]`` -- the + ``(num_experts, input_dim, output_dim)`` layout, i.e. ``is_transposed=True`` for ``_grouped_linear`` + (the opposite of qwen3/deepseek's ``[E, out, in]``). The decorator makes ``forward`` dispatch through + ``ExpertsInterface`` and exposes ``has_gate``/``has_bias``/``is_transposed``/``_apply_gate`` -- exactly + what ``moe.torch_a2a_experts_forward`` reads. The ``grouped_gemm`` shard updates ``num_experts`` to the + per-rank local count, so ``num_global = num_local * ep`` is correct. + +2. **Replace ``Llama4TextMoe.forward`` with the sparse convention:** extract the real top-k from the + router (``topk(router_logits)`` -> global expert ids; ``sigmoid`` of the top values as weights, matching + Llama-4's gating), call ``self.experts(hidden, top_k_index, top_k_weights)`` (the a2a dispatch/combine, + which applies the weights and sums over experts -> ``[T, H]``), then add the replicated shared expert. + Doing the routing here also bypasses ``RouterParallel`` entirely -- Llama4's router returning 2 values + (not the 3 ``RouterParallel._prepare_output_fn`` expects) is what broke the stock-EP attempt. + +3. **Drop the ``ep_router`` entry from ``base_model_ep_plan``** so the router stays a plain replicated + ``Linear`` (we route inline). The grouped_gemm expert-param entries stay; the a2a swap derives the + ``...experts`` container as ``torch_a2a_experts`` and drops the attention / shared-FFN plain-TP styles + (replicated, no TP dim in the DP=EP mesh). + +Because llama4 stores its experts transposed, this also needs ModelOpt's ``_QuantFusedExperts`` grouped +quantization to accept the ``is_transposed`` layout -- supported for per-tensor weight quant (FP8); NVFP4 +(per-block) transposed experts are not yet supported (the block axis would be wrong) and will raise. + +Version floor: transformers >= 5.6 (the ExpertsInterface / EP framework; enforced also by +``moe.load_and_prepare_ep``). Validated FP8 end-to-end on a tiny llama4 (Llama4ForCausalLM, EP4/1 node, +``validation/make_tiny_llama4.py``). Not yet run on real Llama-4-Maverick (4 nodes) or for NVFP4. +""" + +from modelopt.torch.utils.moe import require_transformers + +_PATCHED = False + + +def _patch_llama4_modeling() -> None: + """Idempotently adapt transformers' llama4 MoE to the ExpertsInterface sparse-dispatch convention.""" + global _PATCHED + if _PATCHED: + return + import torch + from transformers.integrations.moe import use_experts_implementation + from transformers.models.llama4 import modeling_llama4 + + # (1) Make Llama4TextExperts an ExpertsInterface experts class: dispatch forward through + # ALL_EXPERTS_FUNCTIONS (so config._experts_implementation="torch_a2a" swaps in the a2a forward) and + # expose has_gate/has_bias/is_transposed/_apply_gate. is_transposed=True: gate_up_proj/down_proj are + # stored [E, in, out] (Llama-4 does bmm(hidden, gate_up_proj)), unlike qwen3/deepseek's [E, out, in]. + use_experts_implementation( + modeling_llama4.Llama4TextExperts, + is_transposed=True, + has_bias=False, + has_gate=True, + ) + + # (2) Sparse-convention MoE forward: route inline, dispatch the real top-k into the experts (a2a + # applies the weights + sums), add the replicated shared expert. Bypasses RouterParallel. + def _moe_forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_out = self.router(hidden_states) + # Llama4Router returns (router_scores, router_logits); use the raw logits for the top-k. + router_logits = router_out[-1] if isinstance(router_out, tuple) else router_out + top_values, top_index = torch.topk(router_logits, self.top_k, dim=-1) + # Llama-4 gates with sigmoid of the selected logits (== the non-zero entries of router_scores). + top_weights = torch.sigmoid(top_values.float()).to(hidden_states.dtype) + routed_out = self.experts( + hidden_states, top_index, top_weights + ) # a2a dispatch/combine -> [T, H] + out = self.shared_expert(hidden_states) + out = out + routed_out.view(out.shape) + return out, router_logits + + modeling_llama4.Llama4TextMoe.forward = _moe_forward + _PATCHED = True + + +def prepare(hf_config, **_): + """Patch llama4 modeling for a2a and drop the router ep_plan entry. Returns the config.""" + require_transformers("5.6", "llama4 expert-parallel PTQ (ExpertsInterface a2a adapter)") + _patch_llama4_modeling() + # base_model_ep_plan lives on the text config for the Llama-4 VLM (Llama4TextConfig). + text_cfg = getattr(hf_config, "text_config", None) or hf_config + plan = dict(getattr(text_cfg, "base_model_ep_plan", None) or {}) + # We route inline in the patched Llama4TextMoe.forward, so keep the router a plain replicated Linear + # (no torch_a2a_router / RouterParallel wrapping). The a2a swap keeps the grouped_gemm expert params + # and derives the experts container; the attention / shared-FFN plain-TP styles are dropped (replicated). + plan = {k: v for k, v in plan.items() if v != "ep_router"} + text_cfg.base_model_ep_plan = plan + # The decorated Llama4TextExperts.forward dispatches on ``config._experts_implementation``. Because our + # decoration is applied at runtime (not in the source file), transformers' ``_can_set_experts_ + # implementation`` source heuristic may not see it and so from_pretrained's ``experts_implementation`` + # kwarg wouldn't stick; set it on the config now so it's present pre-load (load_and_prepare_ep also + # force-sets it post-load, but only where the attr already exists). ``torch_a2a`` is registered into + # ExpertsInterface by ``register_torch_a2a()`` (called from load_and_prepare_ep) before any forward runs. + text_cfg._experts_implementation = "torch_a2a" + return hf_config diff --git a/modelopt/torch/models/llama4/specs.py b/modelopt/torch/models/llama4/specs.py index dc733398d87..8657a0fdade 100644 --- a/modelopt/torch/models/llama4/specs.py +++ b/modelopt/torch/models/llama4/specs.py @@ -13,9 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Llama4 specs (HF model type ``llama4``).""" +"""Llama4 specs (HF model type ``llama4``). -from ..specs import ModelSpec, MoESpec, register +Llama-4 is the one MoE that does not follow the transformers sparse-dispatch convention -- its MoE +block routes inline and computes every expert densely -- so making it expert-parallel is an adapter +rather than a plan, and lives in ``ep.py``. +""" + +from ..specs import EPSpec, ModelSpec, MoESpec, register +from .ep import prepare __all__: list[str] = [] @@ -27,6 +33,9 @@ ModelSpec( model_type="llama4", min_transformers_version="5.5", + ep_spec=EPSpec( + plan_config_attr="text_config", requires_transformers="5.6", prepare=prepare + ), moe_spec=MoESpec( block_names=("Llama4TextMoe",), expert_linear_names=("gate_up_proj", "down_proj"), diff --git a/modelopt/torch/models/llama4_text/__init__.py b/modelopt/torch/models/llama4_text/__init__.py new file mode 100644 index 00000000000..89c3f2a5103 --- /dev/null +++ b/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 diff --git a/modelopt/torch/models/llama4_text/specs.py b/modelopt/torch/models/llama4_text/specs.py new file mode 100644 index 00000000000..63299712f86 --- /dev/null +++ b/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``, so it shares the adapter. +The plan sits on the config itself rather than a nested text config -- there is no VLM wrapper here. +""" + +from ..llama4.ep import prepare +from ..specs import EPSpec, ModelSpec, register + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="llama4_text", + min_transformers_version="5.6", + ep_spec=EPSpec(requires_transformers="5.6", prepare=prepare), + ) +) diff --git a/modelopt/torch/models/minimax_m3_vl/__init__.py b/modelopt/torch/models/minimax_m3_vl/__init__.py new file mode 100644 index 00000000000..d8916c327a4 --- /dev/null +++ b/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 diff --git a/modelopt/torch/models/minimax_m3_vl/ep.py b/modelopt/torch/models/minimax_m3_vl/ep.py new file mode 100644 index 00000000000..7716ddbfcca --- /dev/null +++ b/modelopt/torch/models/minimax_m3_vl/ep.py @@ -0,0 +1,90 @@ +# 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. + +"""DP=EP all-to-all support for transformers-native ``minimax_m3_vl`` (MiniMax-M3, a VLM MoE). + +M3's text MoE (``MiniMaxM3VLSparseMoeBlock``) already follows the ExpertsInterface **sparse-dispatch** +convention -- its forward runs the router inline and calls ``self.experts(hidden, top_k_index, +top_k_weights)`` -- and its experts store the FUSED grouped params ``gate_up_proj [E, 2*inter, hidden]`` +/ ``down_proj [E, hidden, inter]`` (the standard ``[E, out, in]`` layout, i.e. ``is_transposed=False``, +so NVFP4 per-block quant works). The only gap is that ``MiniMaxM3VLExperts.forward`` is HARDCODED (a +per-expert Python loop) instead of dispatching through the transformers ``ExpertsInterface``, so the +generic ``torch_a2a`` transport (``modelopt.torch.utils.moe``) has no seam to take over. This plugin: + +1. Decorates ``MiniMaxM3VLExperts`` with ``use_experts_implementation`` (``is_transposed=False``, + ``has_gate=True``, ``has_bias=False``) so its forward dispatches through ExpertsInterface and honors + ``config._experts_implementation`` at forward time. The generic ``torch_a2a`` experts forward then + takes over -- NCCL all-to-all dispatch -> grouped-GEMM ``_grouped_linear`` -> combine -- preserving + M3's custom SwiGLU via the module's own ``_apply_gate``. Unlike llama4, M3's MoE block already calls + ``experts(hidden, index, weights)``, so NO block-forward rewrite is needed. + +2. Sets ``base_model_ep_plan`` on ``text_config`` (M3 is a VLM; the plan bubbles up under + ``language_model.*``): the fused expert params -> ``grouped_gemm``; the experts container -> + ``moe_tp_experts`` (the a2a swap in ``load_and_prepare_ep`` derives ``torch_a2a_experts``). NO + ``ep_router`` entry -- M3's ``TopKRouter`` runs inline over the REPLICATED full router weight and + emits GLOBAL expert ids, so it stays a plain replicated module (as for deepseek_v3 / llama4). The + ``shared_experts`` (a dense always-on MLP) are absent from the plan, so they stay dense/replicated. + +Load with ``trust_remote_code=False`` so the transformers-native ``MiniMaxM3VLConfig`` + modeling are +used (consistent field set); the repo's remote ``auto_map`` config is a thinner ``PreTrainedConfig`` that +omits native defaults (``rope_parameters`` / vision ``temporal_patch_size``) and desyncs from the native +modeling. + +Version floor: transformers >= 5.13 (first release carrying ``minimax_m3_vl``). +""" + +from modelopt.torch.utils.moe import require_transformers + +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", +} + +_PATCHED = False + + +def _patch_minimax_m3_experts() -> None: + """Route M3's experts through the ExpertsInterface so the generic transport takes over. + + Idempotently route ``MiniMaxM3VLExperts.forward`` through the ExpertsInterface so ``torch_a2a`` + can take over (its per-expert loop is otherwise hardcoded and blind to EP sharding). + """ + global _PATCHED + if _PATCHED: + return + from transformers.integrations.moe import use_experts_implementation + from transformers.models.minimax_m3_vl import modeling_minimax_m3_vl as m + + # is_transposed=False: gate_up_proj/down_proj are [E, out, in] (like qwen3/deepseek, unlike llama4). + # has_gate=True: gate_up is FUSED (2*inter) and split by the module's SwiGLU ``_apply_gate``. + use_experts_implementation( + m.MiniMaxM3VLExperts, + is_transposed=False, + has_bias=False, + has_gate=True, + ) + _PATCHED = True + + +def prepare(hf_config, **_): + """Adapt M3's experts to ExpertsInterface + set the EP plan on ``text_config``. Returns the config.""" + require_transformers("5.13", "minimax_m3_vl expert-parallel PTQ") + _patch_minimax_m3_experts() + text_cfg = getattr(hf_config, "text_config", None) or hf_config + if not getattr(text_cfg, "base_model_ep_plan", None): + text_cfg.base_model_ep_plan = dict(BASE_MODEL_EP_PLAN) + text_cfg._experts_implementation = "torch_a2a" + return hf_config diff --git a/modelopt/torch/models/minimax_m3_vl/specs.py b/modelopt/torch/models/minimax_m3_vl/specs.py new file mode 100644 index 00000000000..da4c0b50ed0 --- /dev/null +++ b/modelopt/torch/models/minimax_m3_vl/specs.py @@ -0,0 +1,39 @@ +# 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``). + +M3's MoE block already follows the sparse-dispatch convention; the only gap is that its experts +hardcode a per-expert loop instead of dispatching through the ExpertsInterface, so ``ep.py`` +decorates the class and the generic transport takes over unchanged. +""" + +from ..specs import EPSpec, ModelSpec, register +from .ep import BASE_MODEL_EP_PLAN, prepare + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="minimax_m3_vl", + min_transformers_version="5.13", + ep_spec=EPSpec( + base_model_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="text_config", + requires_transformers="5.13", + prepare=prepare, + ), + ) +) From 4022d749d4a26c2c887d8b05fa12983b90afa820 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 08:26:21 +0000 Subject: [PATCH 06/23] Refuse Llama-4 expert parallelism and drop the MiniMax-M3 adapter The Llama-4 adapter rewrote Llama4TextMoe.forward to dispatch its top-k through the experts interface, but it scaled each expert's output by the router score where Llama-4 scales the input, and SwiGLU does not commute with that: on a tiny config the outputs differ by up to 6.9 on values near 11. Its experts are also stored transposed, which the expert-parallel export cannot write, so the run failed at quantization anyway. Drop the adapter and refuse llama4 and llama4_text the way gpt_oss is refused: transformers ships a plan for both, which would otherwise route them into the expert-parallel path. Every transformers release with minimax_m3_vl already decorates MiniMaxM3VLExperts exactly as the M3 adapter did and ships the text config's base_model_ep_plan, and the loader selects the all-to-all experts forward after loading, so the adapter changed nothing at runtime. Keep only its transformers floor. Signed-off-by: Shengliang Xu --- modelopt/torch/models/llama4/ep.py | 128 ------------------- modelopt/torch/models/llama4/specs.py | 20 +-- modelopt/torch/models/llama4_text/specs.py | 10 +- modelopt/torch/models/minimax_m3_vl/ep.py | 90 ------------- modelopt/torch/models/minimax_m3_vl/specs.py | 14 +- tests/unit/torch/models/test_ep_specs.py | 5 +- 6 files changed, 22 insertions(+), 245 deletions(-) delete mode 100644 modelopt/torch/models/llama4/ep.py delete mode 100644 modelopt/torch/models/minimax_m3_vl/ep.py diff --git a/modelopt/torch/models/llama4/ep.py b/modelopt/torch/models/llama4/ep.py deleted file mode 100644 index b3264d1d98e..00000000000 --- a/modelopt/torch/models/llama4/ep.py +++ /dev/null @@ -1,128 +0,0 @@ -# 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. - -"""DP=EP a2a support for transformers-native ``llama4`` (Llama-4 Maverick/Scout). - -Unlike every other MoE we run on a2a, Llama-4's ``Llama4TextMoe`` does NOT use the transformers -``ExpertsInterface`` sparse-dispatch convention. Its reference forward: - - router_scores, router_logits = self.router(hidden) # router returns TWO values - routed_in = hidden.repeat(num_experts, 1) * router_scores.T # compute ALL experts, DENSELY - routed_out = self.experts(routed_in) # experts() takes ONE pre-routed tensor - out = shared_expert(hidden) + routed_out.reshape(E, -1, H).sum(0) - -i.e. routing is done inline (not via a wrappable ``RouterParallel``), the experts are computed densely -(every token through every expert, masked by the sigmoid router scores), and the experts module is -called with a single already-routed tensor -- so a2a has no ``router -> experts(hidden, top_k_index, -top_k_weights)`` seam to wrap. (This is why Llama-4 was the one MoE that didn't fit the generic path.) - -This plugin monkey-patches transformers' llama4 modeling to fit the ExpertsInterface contract so the -generic ``torch_a2a`` transport takes over unchanged: - -1. **Decorate ``Llama4TextExperts`` with ``use_experts_implementation``.** Llama-4 stores the expert - weights as ``gate_up_proj`` ``[E, hidden, 2*inter]`` and ``down_proj`` ``[E, inter, hidden]`` -- the - ``(num_experts, input_dim, output_dim)`` layout, i.e. ``is_transposed=True`` for ``_grouped_linear`` - (the opposite of qwen3/deepseek's ``[E, out, in]``). The decorator makes ``forward`` dispatch through - ``ExpertsInterface`` and exposes ``has_gate``/``has_bias``/``is_transposed``/``_apply_gate`` -- exactly - what ``moe.torch_a2a_experts_forward`` reads. The ``grouped_gemm`` shard updates ``num_experts`` to the - per-rank local count, so ``num_global = num_local * ep`` is correct. - -2. **Replace ``Llama4TextMoe.forward`` with the sparse convention:** extract the real top-k from the - router (``topk(router_logits)`` -> global expert ids; ``sigmoid`` of the top values as weights, matching - Llama-4's gating), call ``self.experts(hidden, top_k_index, top_k_weights)`` (the a2a dispatch/combine, - which applies the weights and sums over experts -> ``[T, H]``), then add the replicated shared expert. - Doing the routing here also bypasses ``RouterParallel`` entirely -- Llama4's router returning 2 values - (not the 3 ``RouterParallel._prepare_output_fn`` expects) is what broke the stock-EP attempt. - -3. **Drop the ``ep_router`` entry from ``base_model_ep_plan``** so the router stays a plain replicated - ``Linear`` (we route inline). The grouped_gemm expert-param entries stay; the a2a swap derives the - ``...experts`` container as ``torch_a2a_experts`` and drops the attention / shared-FFN plain-TP styles - (replicated, no TP dim in the DP=EP mesh). - -Because llama4 stores its experts transposed, this also needs ModelOpt's ``_QuantFusedExperts`` grouped -quantization to accept the ``is_transposed`` layout -- supported for per-tensor weight quant (FP8); NVFP4 -(per-block) transposed experts are not yet supported (the block axis would be wrong) and will raise. - -Version floor: transformers >= 5.6 (the ExpertsInterface / EP framework; enforced also by -``moe.load_and_prepare_ep``). Validated FP8 end-to-end on a tiny llama4 (Llama4ForCausalLM, EP4/1 node, -``validation/make_tiny_llama4.py``). Not yet run on real Llama-4-Maverick (4 nodes) or for NVFP4. -""" - -from modelopt.torch.utils.moe import require_transformers - -_PATCHED = False - - -def _patch_llama4_modeling() -> None: - """Idempotently adapt transformers' llama4 MoE to the ExpertsInterface sparse-dispatch convention.""" - global _PATCHED - if _PATCHED: - return - import torch - from transformers.integrations.moe import use_experts_implementation - from transformers.models.llama4 import modeling_llama4 - - # (1) Make Llama4TextExperts an ExpertsInterface experts class: dispatch forward through - # ALL_EXPERTS_FUNCTIONS (so config._experts_implementation="torch_a2a" swaps in the a2a forward) and - # expose has_gate/has_bias/is_transposed/_apply_gate. is_transposed=True: gate_up_proj/down_proj are - # stored [E, in, out] (Llama-4 does bmm(hidden, gate_up_proj)), unlike qwen3/deepseek's [E, out, in]. - use_experts_implementation( - modeling_llama4.Llama4TextExperts, - is_transposed=True, - has_bias=False, - has_gate=True, - ) - - # (2) Sparse-convention MoE forward: route inline, dispatch the real top-k into the experts (a2a - # applies the weights + sums), add the replicated shared expert. Bypasses RouterParallel. - def _moe_forward(self, hidden_states): - hidden_states = hidden_states.reshape(-1, self.hidden_dim) - router_out = self.router(hidden_states) - # Llama4Router returns (router_scores, router_logits); use the raw logits for the top-k. - router_logits = router_out[-1] if isinstance(router_out, tuple) else router_out - top_values, top_index = torch.topk(router_logits, self.top_k, dim=-1) - # Llama-4 gates with sigmoid of the selected logits (== the non-zero entries of router_scores). - top_weights = torch.sigmoid(top_values.float()).to(hidden_states.dtype) - routed_out = self.experts( - hidden_states, top_index, top_weights - ) # a2a dispatch/combine -> [T, H] - out = self.shared_expert(hidden_states) - out = out + routed_out.view(out.shape) - return out, router_logits - - modeling_llama4.Llama4TextMoe.forward = _moe_forward - _PATCHED = True - - -def prepare(hf_config, **_): - """Patch llama4 modeling for a2a and drop the router ep_plan entry. Returns the config.""" - require_transformers("5.6", "llama4 expert-parallel PTQ (ExpertsInterface a2a adapter)") - _patch_llama4_modeling() - # base_model_ep_plan lives on the text config for the Llama-4 VLM (Llama4TextConfig). - text_cfg = getattr(hf_config, "text_config", None) or hf_config - plan = dict(getattr(text_cfg, "base_model_ep_plan", None) or {}) - # We route inline in the patched Llama4TextMoe.forward, so keep the router a plain replicated Linear - # (no torch_a2a_router / RouterParallel wrapping). The a2a swap keeps the grouped_gemm expert params - # and derives the experts container; the attention / shared-FFN plain-TP styles are dropped (replicated). - plan = {k: v for k, v in plan.items() if v != "ep_router"} - text_cfg.base_model_ep_plan = plan - # The decorated Llama4TextExperts.forward dispatches on ``config._experts_implementation``. Because our - # decoration is applied at runtime (not in the source file), transformers' ``_can_set_experts_ - # implementation`` source heuristic may not see it and so from_pretrained's ``experts_implementation`` - # kwarg wouldn't stick; set it on the config now so it's present pre-load (load_and_prepare_ep also - # force-sets it post-load, but only where the attr already exists). ``torch_a2a`` is registered into - # ExpertsInterface by ``register_torch_a2a()`` (called from load_and_prepare_ep) before any forward runs. - text_cfg._experts_implementation = "torch_a2a" - return hf_config diff --git a/modelopt/torch/models/llama4/specs.py b/modelopt/torch/models/llama4/specs.py index 8657a0fdade..ace9c0b8d69 100644 --- a/modelopt/torch/models/llama4/specs.py +++ b/modelopt/torch/models/llama4/specs.py @@ -13,18 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Llama4 specs (HF model type ``llama4``). - -Llama-4 is the one MoE that does not follow the transformers sparse-dispatch convention -- its MoE -block routes inline and computes every expert densely -- so making it expert-parallel is an adapter -rather than a plan, and lives in ``ep.py``. -""" +"""Llama4 specs (HF model type ``llama4``).""" from ..specs import EPSpec, ModelSpec, MoESpec, register -from .ep import prepare __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 @@ -33,9 +35,7 @@ ModelSpec( model_type="llama4", min_transformers_version="5.5", - ep_spec=EPSpec( - plan_config_attr="text_config", requires_transformers="5.6", prepare=prepare - ), + ep_spec=EPSpec(unsupported=EP_UNSUPPORTED), moe_spec=MoESpec( block_names=("Llama4TextMoe",), expert_linear_names=("gate_up_proj", "down_proj"), diff --git a/modelopt/torch/models/llama4_text/specs.py b/modelopt/torch/models/llama4_text/specs.py index 63299712f86..7440e632063 100644 --- a/modelopt/torch/models/llama4_text/specs.py +++ b/modelopt/torch/models/llama4_text/specs.py @@ -15,11 +15,11 @@ """Llama-4 text-only specs (HF model type ``llama4_text``). -The same architecture as ``llama4`` reached through ``Llama4ForCausalLM``, so it shares the adapter. -The plan sits on the config itself rather than a nested text config -- there is no VLM wrapper here. +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.ep import prepare +from ..llama4.specs import EP_UNSUPPORTED from ..specs import EPSpec, ModelSpec, register __all__: list[str] = [] @@ -27,7 +27,7 @@ register( ModelSpec( model_type="llama4_text", - min_transformers_version="5.6", - ep_spec=EPSpec(requires_transformers="5.6", prepare=prepare), + min_transformers_version="5.5", + ep_spec=EPSpec(unsupported=EP_UNSUPPORTED), ) ) diff --git a/modelopt/torch/models/minimax_m3_vl/ep.py b/modelopt/torch/models/minimax_m3_vl/ep.py deleted file mode 100644 index 7716ddbfcca..00000000000 --- a/modelopt/torch/models/minimax_m3_vl/ep.py +++ /dev/null @@ -1,90 +0,0 @@ -# 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. - -"""DP=EP all-to-all support for transformers-native ``minimax_m3_vl`` (MiniMax-M3, a VLM MoE). - -M3's text MoE (``MiniMaxM3VLSparseMoeBlock``) already follows the ExpertsInterface **sparse-dispatch** -convention -- its forward runs the router inline and calls ``self.experts(hidden, top_k_index, -top_k_weights)`` -- and its experts store the FUSED grouped params ``gate_up_proj [E, 2*inter, hidden]`` -/ ``down_proj [E, hidden, inter]`` (the standard ``[E, out, in]`` layout, i.e. ``is_transposed=False``, -so NVFP4 per-block quant works). The only gap is that ``MiniMaxM3VLExperts.forward`` is HARDCODED (a -per-expert Python loop) instead of dispatching through the transformers ``ExpertsInterface``, so the -generic ``torch_a2a`` transport (``modelopt.torch.utils.moe``) has no seam to take over. This plugin: - -1. Decorates ``MiniMaxM3VLExperts`` with ``use_experts_implementation`` (``is_transposed=False``, - ``has_gate=True``, ``has_bias=False``) so its forward dispatches through ExpertsInterface and honors - ``config._experts_implementation`` at forward time. The generic ``torch_a2a`` experts forward then - takes over -- NCCL all-to-all dispatch -> grouped-GEMM ``_grouped_linear`` -> combine -- preserving - M3's custom SwiGLU via the module's own ``_apply_gate``. Unlike llama4, M3's MoE block already calls - ``experts(hidden, index, weights)``, so NO block-forward rewrite is needed. - -2. Sets ``base_model_ep_plan`` on ``text_config`` (M3 is a VLM; the plan bubbles up under - ``language_model.*``): the fused expert params -> ``grouped_gemm``; the experts container -> - ``moe_tp_experts`` (the a2a swap in ``load_and_prepare_ep`` derives ``torch_a2a_experts``). NO - ``ep_router`` entry -- M3's ``TopKRouter`` runs inline over the REPLICATED full router weight and - emits GLOBAL expert ids, so it stays a plain replicated module (as for deepseek_v3 / llama4). The - ``shared_experts`` (a dense always-on MLP) are absent from the plan, so they stay dense/replicated. - -Load with ``trust_remote_code=False`` so the transformers-native ``MiniMaxM3VLConfig`` + modeling are -used (consistent field set); the repo's remote ``auto_map`` config is a thinner ``PreTrainedConfig`` that -omits native defaults (``rope_parameters`` / vision ``temporal_patch_size``) and desyncs from the native -modeling. - -Version floor: transformers >= 5.13 (first release carrying ``minimax_m3_vl``). -""" - -from modelopt.torch.utils.moe import require_transformers - -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", -} - -_PATCHED = False - - -def _patch_minimax_m3_experts() -> None: - """Route M3's experts through the ExpertsInterface so the generic transport takes over. - - Idempotently route ``MiniMaxM3VLExperts.forward`` through the ExpertsInterface so ``torch_a2a`` - can take over (its per-expert loop is otherwise hardcoded and blind to EP sharding). - """ - global _PATCHED - if _PATCHED: - return - from transformers.integrations.moe import use_experts_implementation - from transformers.models.minimax_m3_vl import modeling_minimax_m3_vl as m - - # is_transposed=False: gate_up_proj/down_proj are [E, out, in] (like qwen3/deepseek, unlike llama4). - # has_gate=True: gate_up is FUSED (2*inter) and split by the module's SwiGLU ``_apply_gate``. - use_experts_implementation( - m.MiniMaxM3VLExperts, - is_transposed=False, - has_bias=False, - has_gate=True, - ) - _PATCHED = True - - -def prepare(hf_config, **_): - """Adapt M3's experts to ExpertsInterface + set the EP plan on ``text_config``. Returns the config.""" - require_transformers("5.13", "minimax_m3_vl expert-parallel PTQ") - _patch_minimax_m3_experts() - text_cfg = getattr(hf_config, "text_config", None) or hf_config - if not getattr(text_cfg, "base_model_ep_plan", None): - text_cfg.base_model_ep_plan = dict(BASE_MODEL_EP_PLAN) - text_cfg._experts_implementation = "torch_a2a" - return hf_config diff --git a/modelopt/torch/models/minimax_m3_vl/specs.py b/modelopt/torch/models/minimax_m3_vl/specs.py index da4c0b50ed0..025d2f7b22a 100644 --- a/modelopt/torch/models/minimax_m3_vl/specs.py +++ b/modelopt/torch/models/minimax_m3_vl/specs.py @@ -15,13 +15,12 @@ """MiniMax-M3 specs (HF model type ``minimax_m3_vl``). -M3's MoE block already follows the sparse-dispatch convention; the only gap is that its experts -hardcode a per-expert loop instead of dispatching through the ExpertsInterface, so ``ep.py`` -decorates the class and the generic transport takes over unchanged. +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 -from .ep import BASE_MODEL_EP_PLAN, prepare __all__: list[str] = [] @@ -29,11 +28,6 @@ ModelSpec( model_type="minimax_m3_vl", min_transformers_version="5.13", - ep_spec=EPSpec( - base_model_ep_plan=BASE_MODEL_EP_PLAN, - plan_config_attr="text_config", - requires_transformers="5.13", - prepare=prepare, - ), + ep_spec=EPSpec(requires_transformers="5.13"), ) ) diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index c52553097c9..4a409bacade 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -62,8 +62,9 @@ def test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal(model_type): def test_models_whose_own_plan_would_misroute_them_are_refused(): - """GPT-OSS ships a plan of its own, so only its refusal keeps it off the expert-parallel path.""" - assert "gpt_oss" in EP_REFUSED_MODEL_TYPES + """GPT-OSS and Llama-4 ship plans of their own, so only their refusal keeps them off the + expert-parallel path.""" + assert {"gpt_oss", "llama4", "llama4_text"} <= set(EP_REFUSED_MODEL_TYPES) def test_the_omni_vlm_reuses_the_text_models_plan(): From 6d206a15ea5011df45174e70fa42fd83df914810 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 12:10:15 +0000 Subject: [PATCH 07/23] Drop Nemotron-H's expert-parallel style, which the loader never used nemotron_h/ep.py registered a NemotronHExpertsEP parallel style for the plan's nemotron_experts_ep entry, but load_and_prepare_ep maps that entry to torch_a2a_experts before loading, so the style never ran. Registering it imported transformers' all_reduce_forward, which 5.18 no longer has, so applying Nemotron-H's spec failed on CI's latest transformers; its _prepare_* hooks are not called from 5.16 on either. The plan moves into the spec that both spellings of the architecture share, with no prepare hook. Signed-off-by: Shengliang Xu --- modelopt/torch/models/nemotron_h/ep.py | 84 ------------------- modelopt/torch/models/nemotron_h/specs.py | 13 ++- .../torch/models/nemotron_h_omni/specs.py | 3 +- 3 files changed, 12 insertions(+), 88 deletions(-) delete mode 100644 modelopt/torch/models/nemotron_h/ep.py diff --git a/modelopt/torch/models/nemotron_h/ep.py b/modelopt/torch/models/nemotron_h/ep.py deleted file mode 100644 index b10d58cfee9..00000000000 --- a/modelopt/torch/models/nemotron_h/ep.py +++ /dev/null @@ -1,84 +0,0 @@ -# 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. - -"""The expert-parallel style NemotronH's block-routed MoE needs. - -Everything else in this package is data; this is the one piece that has to be code, because -transformers resolves parallel styles through a global registry rather than from a config. -""" - -_REGISTERED = False - - -def _register_style() -> None: - """Register ``NemotronHExpertsEP`` under the name the spec's plan refers to. - - Imported lazily so this package stays importable without torch or transformers, and - idempotent because a PTQ run may prepare several models in one process. - """ - global _REGISTERED - if _REGISTERED: - return - - import torch - from transformers.integrations.tensor_parallel import ( - ALL_PARALLEL_STYLES, - MoeTensorParalellExperts, - all_reduce_forward, - ) - - class NemotronHExpertsEP(MoeTensorParalellExperts): - """Expert parallelism for a gate that emits logits only. - - The usual ``ep_router`` style remaps expert indices on the ROUTER's output, which - requires the router to return ``(scores, indices)``. NemotronH's gate returns logits - and the group-limited top-k happens inside the MoE block, so there is nothing for - ``ep_router`` to remap and the experts would receive global indices they cannot use. - - So the remap happens at the EXPERTS' input instead: a token routed to one of this - rank's local experts gets a local index, and a token belonging elsewhere is pointed at - local index 0 with weight zero, which makes the experts' ``index_add`` a no-op for it. - ``grouped_gemm`` has already set ``num_experts`` to the local count, and the inherited - ``all_reduce_forward`` sums each rank's partial expert output. - """ - - def _prepare_input_fn(self, mod, inputs, device_mesh): - hidden, top_k_index, top_k_weights = inputs[0], inputs[1], inputs[2] - local = mod.num_experts # local count, after grouped_gemm shards dim 0 - start = device_mesh.get_local_rank() * local - is_local = (top_k_index >= start) & (top_k_index < start + local) - local_index = torch.where(is_local, top_k_index - start, torch.zeros_like(top_k_index)) - local_weights = torch.where(is_local, top_k_weights, torch.zeros_like(top_k_weights)) - return (hidden, local_index, local_weights) - - def _prepare_output_fn(self, mod, outputs, device_mesh): - return all_reduce_forward(outputs, device_mesh) - - ALL_PARALLEL_STYLES.register("nemotron_experts_ep", NemotronHExpertsEP()) - _REGISTERED = True - - -def prepare(hf_config, **_): - """Register the style the plan names, then hand the config back unchanged.""" - _register_style() - return hf_config - - -BASE_MODEL_EP_PLAN = { - "layers.*.mixer.experts.up_proj": "grouped_gemm", - "layers.*.mixer.experts.down_proj": "grouped_gemm", - "layers.*.mixer.experts": "nemotron_experts_ep", -} -"""Ungated experts, so ``up_proj``/``down_proj`` rather than a gate/up pair.""" diff --git a/modelopt/torch/models/nemotron_h/specs.py b/modelopt/torch/models/nemotron_h/specs.py index dd7db03a5b6..8356310948c 100644 --- a/modelopt/torch/models/nemotron_h/specs.py +++ b/modelopt/torch/models/nemotron_h/specs.py @@ -16,16 +16,25 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register -from .ep import BASE_MODEL_EP_PLAN, prepare __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, prepare=prepare), + 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",), diff --git a/modelopt/torch/models/nemotron_h_omni/specs.py b/modelopt/torch/models/nemotron_h_omni/specs.py index e022f5b8430..6fc14007fcb 100644 --- a/modelopt/torch/models/nemotron_h_omni/specs.py +++ b/modelopt/torch/models/nemotron_h_omni/specs.py @@ -23,7 +23,7 @@ real parameter paths -- while a plan on the top config reaches no expert at all. """ -from ..nemotron_h.ep import BASE_MODEL_EP_PLAN, prepare +from ..nemotron_h.specs import BASE_MODEL_EP_PLAN from ..specs import EPSpec, ModelSpec, register __all__: list[str] = [] @@ -37,7 +37,6 @@ ep_spec=EPSpec( base_model_ep_plan=BASE_MODEL_EP_PLAN, plan_config_attr="llm_config", - prepare=prepare, ), ) ) From faaaa04f37b9a4a9a1f974fba831b7a9ea270546 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Mon, 12 Oct 2026 01:28:43 +0000 Subject: [PATCH 08/23] Record that transformers ships NemotronH-Omni from 5.18 transformers 5.18 added nemotron_h_omni, so its spec can no longer claim remote code; test_specs_vs_transformers fails on the latest transformers until it names the release. Expert parallelism still needs the checkpoint's code, whose llm_config the plan targets. Signed-off-by: Shengliang Xu --- modelopt/torch/models/nemotron_h_omni/specs.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/models/nemotron_h_omni/specs.py b/modelopt/torch/models/nemotron_h_omni/specs.py index 6fc14007fcb..3d9b6dce4cb 100644 --- a/modelopt/torch/models/nemotron_h_omni/specs.py +++ b/modelopt/torch/models/nemotron_h_omni/specs.py @@ -31,9 +31,10 @@ register( ModelSpec( model_type="nemotron_h_omni", - # The Omni VLM ships its own modeling code with the checkpoint; only its language - # sub-model is the transformers-native NemotronH. - modeling_source="remote_code", + # 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", From 461c92c4f5fa13a325b6a56b0c3e9e3db8b1f070 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 23:42:03 +0000 Subject: [PATCH 09/23] Cite the expert-parallel loader where it lives, in model_load_utils Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index c8b3e155852..346b3dbdbf9 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -162,7 +162,8 @@ def apply_distributed_support( Returns the config to load, unchanged for a model type with no spec or no EP section. The generic all-to-all transport is not wired here -- it self-registers in - :func:`load_and_prepare_ep`, so nothing about it is per-model. + :func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep`, so nothing about it + is per-model. """ from modelopt.torch.models import get_spec From 9f3d96af57bad2e0ee0f107bc1ca53736586799f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 01:55:43 +0000 Subject: [PATCH 10/23] Apply the model spec before deciding and loading expert-parallel For several models the spec is what supplies the expert-parallel plan: DeepSeek-V3 ships none of its own, so deciding without the spec would read its MoE as dense. Unless the run rules expert parallelism out (a single process, or an EP degree of 1), the model's spec is applied to the config before rank 0 builds the model to decide, and the returned config carries it. A load prepares that config once, and not at all when handed the one resolve_distribution_plan returned, since a model's preparation can patch its modeling code. A run that rules expert parallelism out applies no spec, so a model whose spec refuses expert parallelism, like GPT-OSS, still loads FSDP2-sharded. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 24 +++++++- .../test_load_pretrained_distribution.py | 50 ++++++++++++++--- .../utils/test_resolve_distribution_plan.py | 56 +++++++++++++++++-- 3 files changed, 116 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1d28d24b13d..5c39ff8a5de 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -461,6 +461,15 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: + from modelopt.torch.utils.moe import apply_distributed_support + + if hf_config is None: + hf_config = apply_distributed_support( + model_type, + AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code), + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + ) return load_and_prepare_ep( ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs ) @@ -1289,9 +1298,14 @@ def resolve_distribution_plan( ``config``, the config the load will use when the caller has one, is decided from instead of the checkpoint's, which is then not read: each rank takes a copy, as deciding may add to it. - A plan that leaves an expert-parallel MoE whole on every rank raises. + Unless the run rules expert parallelism out (a single process, or an EP degree of 1), the + model's spec is applied to the config first + (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what + supplies the expert-parallel plan, and the returned config carries it. A plan that leaves an + expert-parallel MoE whole on every rank raises. """ from modelopt.torch.utils import distributed as dist + from modelopt.torch.utils.moe import apply_distributed_support # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. @@ -1317,6 +1331,14 @@ def checkpoint_read() -> bool: dist.all_reduce("checkpoint read", checkpoint_read, all) ckpt_path, hf_config = read[0] world_size = dist.size() + if world_size > 1 and str(env.ep_degree() if ep_size is None else ep_size) != "1": + # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. + hf_config = apply_distributed_support( + hf_config.model_type, + hf_config, + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + ) plan = dist.broadcast_from( "plan", lambda: _decide_distribution( diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 8429c4cc58e..289788b332c 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -33,6 +33,8 @@ import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist +import modelopt.torch.utils.moe as moe +from modelopt.torch.models import get_spec from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils @@ -273,30 +275,60 @@ def test_hf_config_is_refused_without_the_plan_it_came_with( def test_an_expert_parallel_plan_loads_through_the_ep_loader( - run_as, serve, loaders, auto_class, tmp_path + run_as, serve, loaders, auto_class, monkeypatch, tmp_path ): - """An EP plan loads through load_and_prepare_ep, with the plan, and through no other load.""" + """Decided automatically, an EP plan loads through load_and_prepare_ep -- with the plan and the + spec-prepared config -- and through no other load.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(4) - serve(_deepseek_v3(n_routed_experts=8)) - plan = DistributionPlan(4, ep_size=4, dense="fsdp") + served = serve(_deepseek_v3(n_routed_experts=8)) model = load_pretrained( - str(tmp_path), - distribution=plan, - trust_remote_code=True, - attn_implementation="eager", - freeze=False, + str(tmp_path), trust_remote_code=True, attn_implementation="eager", freeze=False ) call = loaders.pop("ep") + plan = DistributionPlan(4, ep_size=4, dense="fsdp") assert (call["ckpt_path"], call["plan"]) == (str(tmp_path), plan) assert (call["trust_remote_code"], call["attn_implementation"]) == (True, "eager") assert call["freeze"] is False + assert call["hf_config"] is served + assert served.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan assert distribution_plan(model) == plan assert loaders == {} assert auto_class.calls == [] +@pytest.mark.parametrize( + ("planned", "given_config", "preparations"), + [(False, False, 1), (True, False, 1), (True, True, 0)], + ids=["decided_by_the_load", "planned", "planned_with_its_config"], +) +def test_a_load_prepares_its_config_once( + run_as, serve, loaders, monkeypatch, tmp_path, planned, given_config, preparations +): + """A model's EP preparation can patch its modeling code, so a load runs it once, and not at all + when handed the config resolve_distribution_plan prepared.""" + run_as(4) + config = serve(_deepseek_v3(n_routed_experts=8)) + prepared = [] + + def apply_distributed_support(model_type, hf_config, **kwargs): + prepared.append(model_type) + hf_config.base_model_ep_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan + return hf_config + + monkeypatch.setattr(moe, "apply_distributed_support", apply_distributed_support) + kwargs = {"distribution": DistributionPlan(4, ep_size=4)} if planned else {} + if given_config: + kwargs["hf_config"] = config + + load_pretrained(str(tmp_path), **kwargs) + + assert len(prepared) == preparations + assert loaders["ep"]["hf_config"] is config + + def _every_rank_loads_rank0s_plan_job(rank, size, ckpt): handed = DistributionPlan(size) if rank == 0 else DistributionPlan(size, dense="fsdp") with bounded_collectives(): diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 177334cef62..251373bcb9f 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -38,12 +38,15 @@ AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, + DeepseekV3Config, + GptOssConfig, LlamaConfig, Qwen2_5_VLConfig, Qwen3MoeConfig, ) import modelopt.torch.utils.distributed as dist +from modelopt.torch.models import get_spec from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.moe import ( _routed_expert_modules, @@ -73,6 +76,24 @@ def _qwen3_moe(num_experts=8): return Qwen3MoeConfig(num_experts=num_experts, moe_intermediate_size=8, **_TINY) +def _deepseek_v3(n_routed_experts=8): + """DeepSeek-V3 without a plan of its own, so only its model spec can make it EP-capable. + + Its one layer is an MoE layer: DeepSeek-V3's first layers are dense by default. + """ + config = DeepseekV3Config( + n_routed_experts=n_routed_experts, + first_k_dense_replace=0, + moe_intermediate_size=8, + n_group=1, + topk_group=1, + num_experts_per_tok=1, + **_TINY, + ) + config.base_model_ep_plan = None # holds even if a later transformers ships a plan + return config + + def _bytes(module): return sum(p.numel() * p.element_size() for p in module.parameters()) @@ -382,6 +403,7 @@ def test_an_ep_capable_moe_is_planned_expert_parallel(run_as, serve, tmp_path, c assert ("Distribution:" in capsys.readouterr().out) == (rank == 0) +@needs_qwen3_moe_ep_plan # its spec asks for the release that ships its plan def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_path): """The config counts 8 experts, but with no base_model_ep_plan nothing could shard them.""" run_as(4) @@ -392,11 +414,36 @@ def test_an_moe_without_an_ep_plan_is_not_expert_parallel(run_as, serve, tmp_pat assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4) -@needs_qwen3_moe_ep_plan +def test_the_spec_is_applied_before_the_decision(run_as, serve, tmp_path): + """DeepSeek-V3's plan comes only from its spec. Deciding before applying it would read the MoE + as dense; applied first, it shards, and the config handed on carries the spec's plan.""" + run_as(4) + serve(_deepseek_v3(n_routed_experts=8)) + + plan, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert plan == DistributionPlan(4, 4) + assert hf_config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + + +def test_a_run_without_expert_parallelism_needs_no_spec(run_as, serve, monkeypatch, tmp_path): + """GPT-OSS cannot be expert-parallel, which its spec says by refusing. A run that rules expert + parallelism out, as FSDP2 does, does not ask its spec, and its spec does not refuse it.""" + run_as(2) + serve(GptOssConfig(num_local_experts=4, **_TINY)) + + with pytest.raises(NotImplementedError, match="cannot be quantized expert-parallel yet"): + resolve_distribution_plan(str(tmp_path)) + plan, _ = resolve_distribution_plan(str(tmp_path), ep_size=1, dense="fsdp") + assert plan == DistributionPlan(2, dense="fsdp") + monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") + assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == plan + + @pytest.mark.parametrize( ("config", "ep_degree"), [ - pytest.param(_qwen3_moe(num_experts=3), None, id="auto_finds_no_degree"), + pytest.param(_deepseek_v3(n_routed_experts=3), None, id="auto_finds_no_degree"), pytest.param(_qwen3_moe(num_experts=8), "1", id="explicit_degree_one"), ], ) @@ -406,6 +453,8 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With nothing sharded either, every rank would hold the whole MoE: that must fail, not run.""" if ep_degree is not None: + if Version(transformers.__version__) < Version("5.7"): + pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) run_as(2) serve(config) @@ -414,14 +463,13 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( resolve_distribution_plan(str(tmp_path)) -@needs_qwen3_moe_ep_plan def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a replica on every rank.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", "1") monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") run_as(2) - serve(_qwen3_moe(num_experts=8)) + serve(_deepseek_v3(n_routed_experts=8)) assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(2, dense="fsdp") From eb87fe03875d44283f508385768ca2dd530b525e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:56:00 +0000 Subject: [PATCH 11/23] Apply the model spec on every rank as rank 0 decided Each rank read its own MODELOPT_EP_DEGREE to decide whether to apply the model spec, though the knobs are rank 0's. A rank whose environment ruled expert parallelism out skipped the spec rank 0 applied, loaded without the expert-parallel plan the others shard by, and raised once the mesh was up, leaving them waiting. Rank 0 decides it and sends it, as it does the plan. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 12 ++++-- .../utils/test_resolve_distribution_plan.py | 39 +++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 5c39ff8a5de..93e6ca3efe6 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1298,8 +1298,8 @@ def resolve_distribution_plan( ``config``, the config the load will use when the caller has one, is decided from instead of the checkpoint's, which is then not read: each rank takes a copy, as deciding may add to it. - Unless the run rules expert parallelism out (a single process, or an EP degree of 1), the - model's spec is applied to the config first + Unless rank 0's knobs rule expert parallelism out (a single process, or an EP degree of 1), + every rank applies the model's spec to the config first (:func:`~modelopt.torch.utils.moe.apply_distributed_support`): for several models it is what supplies the expert-parallel plan, and the returned config carries it. A plan that leaves an expert-parallel MoE whole on every rank raises. @@ -1331,7 +1331,13 @@ def checkpoint_read() -> bool: dist.all_reduce("checkpoint read", checkpoint_read, all) ckpt_path, hf_config = read[0] world_size = dist.size() - if world_size > 1 and str(env.ep_degree() if ep_size is None else ep_size) != "1": + # Whether the run rules expert parallelism out is rank 0's call, as the plan is: a rank that + # skipped the spec while the others applied it would load without the expert-parallel plan + # they shard by, and fail once the mesh is up, leaving them waiting in a collective. + if dist.broadcast_from( + "spec", + lambda: world_size > 1 and str(env.ep_degree() if ep_size is None else ep_size) != "1", + ): # Before deciding: deciding without the spec would read an expert-parallel MoE as dense. hf_config = apply_distributed_support( hf_config.model_type, diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 251373bcb9f..12676c5b451 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -662,6 +662,45 @@ def test_every_rank_takes_the_plan_rank0_decided(cpu_dist_workers, tmp_path): ) +def _as_rank(monkeypatch, rank, **rank0_messages): + """Run as ``rank`` of a 4-rank run, answering its broadcasts; returns what this rank handed them. + + A non-zero rank receives, for each broadcast, ``rank0_messages[label]``, as rank 0 would have + sent it. Every rank reads the checkpoint as this one does. + """ + sent = [] + + def broadcast_object_list(message, group=None, group_src=None): + sent.append(message[0]) + if rank != 0: + message[0] = rank0_messages[message[0][0]] + + monkeypatch.setattr(dist, "is_distributed_run", lambda: True) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) + monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 4) + monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) + monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike) + return sent + + +@pytest.mark.parametrize("rank0_applies", [True, False]) +def test_other_ranks_apply_the_spec_as_rank0_decided(monkeypatch, serve, tmp_path, rank0_applies): + """A rank whose own EP knob says otherwise follows rank 0: had it skipped the spec rank 0 + applied, it would load without the expert-parallel plan the others shard by.""" + monkeypatch.setenv("MODELOPT_EP_DEGREE", "auto" if not rank0_applies else "1") + rank0_plan = DistributionPlan(4, 4) if rank0_applies else DistributionPlan(4) + _as_rank(monkeypatch, 1, spec=("spec", rank0_applies, None), plan=("plan", rank0_plan, None)) + serve(_deepseek_v3(n_routed_experts=8)) + + plan, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert plan == rank0_plan + spec_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan + assert (getattr(hf_config, "base_model_ep_plan", None) == spec_plan) is rank0_applies + + def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt): isdir = os.path.isdir with ( From de6afcd2a716ceaf89764e1d77bd986e602828b8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 01:21:01 +0000 Subject: [PATCH 12/23] Refuse a whole-on-every-rank MoE whose plan only its spec supplies An EP degree of 1 skips the model's spec, so the guard against leaving an expert-parallel MoE whole on every rank found routed experts only through the config's own base_model_ep_plan. DeepSeek-V3, whose plan comes only from its spec, therefore loaded whole on every rank at MODELOPT_EP_DEGREE=1 with a replicated dense part, where Qwen3-MoE, which carries its plan, refused. The guard now also looks for the experts through the plan the spec would supply. The spec itself stays skipped: it refuses models that cannot be expert-parallel, which an EP degree of 1 is how to load. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 15 +++++++++++ .../torch/utils/plugins/model_load_utils.py | 10 ++++++- .../utils/test_resolve_distribution_plan.py | 26 ++++++++++++++++--- 3 files changed, 47 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 346b3dbdbf9..5432a95878b 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -197,6 +197,21 @@ def apply_distributed_support( return hf_config +def _spec_ep_plan(model_type: str) -> dict[str, str]: + """The ``base_model_ep_plan`` :func:`apply_distributed_support` would attach for ``model_type``. + + Empty for a model type with no spec, no expert-parallel section, or one that refuses expert + parallelism. Only reads the spec: nothing is applied, checked or patched. + """ + from modelopt.torch.models import get_spec + + spec = get_spec(model_type) + ep = getattr(spec, "ep_spec", None) if spec is not None else None + if ep is None or ep.unsupported: + return {} + return dict(ep.base_model_ep_plan or {}) + + # --------------------------------------------------------------------------------------------------- # DP=EP all-to-all expert transport (`torch_a2a`) # diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 93e6ca3efe6..49ecb88d795 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1185,6 +1185,7 @@ def _decide_distribution( from modelopt.torch.utils.moe import ( _ep_config, _routed_expert_modules, + _spec_ep_plan, choose_ep_dp_degrees, num_experts_from_config, ) @@ -1245,7 +1246,14 @@ def _decide_distribution( f"The dense part of a block-fp8 {hf_config.model_type} checkpoint cannot be sharded " "expert-parallel: Distributed Checkpoint cannot dequantize its F8_E8M0 scales." ) - if experts and ep_size == 1 and dense == "replicate": + # An EP degree of 1 skips the model's spec (resolve_distribution_plan), so a model whose plan only + # its spec supplies shows no routed experts above: look for them through the plan the spec would + # supply, or such an MoE would load whole on every rank where one carrying its own plan refuses. + if ( + ep_size == 1 + and dense == "replicate" + and (experts or _routed_expert_modules(model, _spec_ep_plan(hf_config.model_type))) + ): raise RuntimeError( f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at " f"world_size={world_size}, so every rank would hold the whole MoE. Shard its experts " diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 12676c5b451..47915b4226d 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -445,15 +445,18 @@ def test_a_run_without_expert_parallelism_needs_no_spec(run_as, serve, monkeypat [ pytest.param(_deepseek_v3(n_routed_experts=3), None, id="auto_finds_no_degree"), pytest.param(_qwen3_moe(num_experts=8), "1", id="explicit_degree_one"), + pytest.param(_deepseek_v3(n_routed_experts=8), "1", id="explicit_degree_one_spec_plan"), ], ) def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( run_as, serve, monkeypatch, tmp_path, config, ep_degree ): - """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none. With - nothing sharded either, every rank would hold the whole MoE: that must fail, not run.""" + """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none, + whether the plan comes with the config (Qwen3-MoE) or only from the model's spec, which an EP + degree of 1 skips (DeepSeek-V3). With nothing sharded either, every rank would hold the whole + MoE: that must fail, not run.""" if ep_degree is not None: - if Version(transformers.__version__) < Version("5.7"): + if config.model_type == "qwen3_moe" and Version(transformers.__version__) < Version("5.7"): pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) run_as(2) @@ -463,6 +466,23 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( resolve_distribution_plan(str(tmp_path)) +def test_the_plan_a_spec_would_supply(monkeypatch): + """What that guard reads at an EP degree of 1: the spec's plan, nothing for a model with no + spec or no expert-parallel section, and nothing for a spec that refuses expert parallelism even + when it names a plan.""" + from modelopt.torch import models + from modelopt.torch.utils import moe + + assert moe._spec_ep_plan("deepseek_v3") == get_spec("deepseek_v3").ep_spec.base_model_ep_plan + assert moe._spec_ep_plan("llama") == {} + assert moe._spec_ep_plan("no_such_model_type") == {} + refusing = SimpleNamespace( + ep_spec=SimpleNamespace(unsupported="not yet", base_model_ep_plan={"layers.*.mlp": "x"}) + ) + monkeypatch.setattr(models, "get_spec", lambda model_type: refusing) + assert moe._spec_ep_plan("anything") == {} + + def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): """Asked for no expert parallelism but for dense sharding, the MoE is sharded whole by FSDP2, which is not the waste of a replica on every rank.""" From 828b929a3e51e1f42c87fca4b0698fd904bd85eb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 01:48:48 +0000 Subject: [PATCH 13/23] Let a spec state the expert plan its prepare attaches At an EP degree of 1 the guard against holding an expert-parallel MoE whole on every rank falls back on the plan the model's spec would supply, since that degree skips the spec. It read only the spec's base_model_ep_plan, so a remote-code model whose prepare attaches the plan, after patching the checkpoint's modeling, showed no routed experts and loaded whole on every rank. Running prepare there is not an option: it patches the model's code, for a run that is not expert-parallel. EPSpec.prepared_ep_plan states the plan prepare attaches, and the plan the guard reads includes it. prepare still attaches the plan itself; the spec only states it. Signed-off-by: Shengliang Xu --- modelopt/torch/models/specs.py | 10 +++++++ modelopt/torch/utils/moe.py | 7 +++-- .../utils/test_resolve_distribution_plan.py | 28 +++++++++++++------ 3 files changed, 33 insertions(+), 12 deletions(-) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index 2e1804095da..a901c1bc524 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -209,6 +209,16 @@ class EPSpec(SpecSection): patching classes or installing on-the-fly weight decompression. Runs after the declarative fields above have been applied.""" + prepared_ep_plan: dict[str, str] | None = None + """The ``base_model_ep_plan`` ``prepare`` attaches, for a model whose plan only ``prepare`` + supplies. + + Stated, not applied: ``prepare`` still attaches it. A run that rules expert parallelism out + skips ``prepare``, which patches the model's code, yet must still tell an expert-parallel MoE + it would hold whole on every rank from a dense model. It finds the routed experts through this + plan, in the model as the checkpoint's own code builds it: ``prepare`` replaces the experts in + place, so the plan names their containers there too.""" + unsupported: str | None = None """Why this model cannot be quantized expert-parallel yet, when its config would route it there. diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 5432a95878b..26535c996d4 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -200,8 +200,9 @@ def apply_distributed_support( def _spec_ep_plan(model_type: str) -> dict[str, str]: """The ``base_model_ep_plan`` :func:`apply_distributed_support` would attach for ``model_type``. - Empty for a model type with no spec, no expert-parallel section, or one that refuses expert - parallelism. Only reads the spec: nothing is applied, checked or patched. + Itself, or through the spec's ``prepare``. Empty for a model type with no spec, no + expert-parallel section, or one that refuses expert parallelism. Only reads the spec: nothing + is applied, checked or patched, and ``prepare`` does not run. """ from modelopt.torch.models import get_spec @@ -209,7 +210,7 @@ def _spec_ep_plan(model_type: str) -> dict[str, str]: ep = getattr(spec, "ep_spec", None) if spec is not None else None if ep is None or ep.unsupported: return {} - return dict(ep.base_model_ep_plan or {}) + return dict(ep.base_model_ep_plan or ep.prepared_ep_plan or {}) # --------------------------------------------------------------------------------------------------- diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 47915b4226d..f98e74a3d57 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -46,7 +46,7 @@ ) import modelopt.torch.utils.distributed as dist -from modelopt.torch.models import get_spec +from modelopt.torch.models import EPSpec, get_spec from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.moe import ( _routed_expert_modules, @@ -467,20 +467,30 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( def test_the_plan_a_spec_would_supply(monkeypatch): - """What that guard reads at an EP degree of 1: the spec's plan, nothing for a model with no - spec or no expert-parallel section, and nothing for a spec that refuses expert parallelism even - when it names a plan.""" + """What that guard reads at an EP degree of 1: the spec's plan, or the plan it states its + ``prepare`` attaches, read without running ``prepare``; nothing for a model with no spec or no + expert-parallel section, and nothing for a spec that refuses expert parallelism even when it + names a plan.""" from modelopt.torch import models from modelopt.torch.utils import moe assert moe._spec_ep_plan("deepseek_v3") == get_spec("deepseek_v3").ep_spec.base_model_ep_plan assert moe._spec_ep_plan("llama") == {} assert moe._spec_ep_plan("no_such_model_type") == {} - refusing = SimpleNamespace( - ep_spec=SimpleNamespace(unsupported="not yet", base_model_ep_plan={"layers.*.mlp": "x"}) - ) - monkeypatch.setattr(models, "get_spec", lambda model_type: refusing) - assert moe._spec_ep_plan("anything") == {} + prepared = {"layers.*.mlp.experts": "moe_tp_experts"} + specs = { + "preparing": SimpleNamespace( + ep_spec=EPSpec( + prepare=lambda *a, **k: pytest.fail("ran prepare"), prepared_ep_plan=prepared + ) + ), + "refusing": SimpleNamespace( + ep_spec=SimpleNamespace(unsupported="not yet", base_model_ep_plan={"layers.*.mlp": "x"}) + ), + } + monkeypatch.setattr(models, "get_spec", specs.get) + assert moe._spec_ep_plan("preparing") == prepared + assert moe._spec_ep_plan("refusing") == {} def test_an_ep_capable_moe_may_be_fsdp2_sharded_instead(run_as, serve, monkeypatch, tmp_path): From 467afb6518e5c9c9f994c420ef1f38cdf8ca14f9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 02:11:22 +0000 Subject: [PATCH 14/23] Point an MoE whose spec refuses expert parallelism at FSDP2 alone The refusal of an MoE held whole on every rank told every model to raise MODELOPT_EP_DEGREE above 1. GPT-OSS and Llama-4 carry plans of their own, so they reach that refusal too, but their specs refuse expert parallelism: following the advice would only have hit that refusal next. For them the error now says they cannot be quantized expert-parallel yet and points at FSDP2 alone, as their specs' own refusal does. Both messages name the ways to ask for FSDP2: --use_fsdp2 in hf_ptq, or MODELOPT_FSDP_SIZE=world. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 17 ++++++------ .../torch/utils/plugins/model_load_utils.py | 12 ++++++++- .../utils/test_resolve_distribution_plan.py | 26 ++++++++++++++----- 3 files changed, 39 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 26535c996d4..07a5e11a743 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -141,6 +141,13 @@ def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict } +def _ep_spec(model_type: str): + """``model_type``'s :class:`~modelopt.torch.models.EPSpec`, or ``None``.""" + from modelopt.torch.models import get_spec + + return getattr(get_spec(model_type), "ep_spec", None) + + def apply_distributed_support( model_type: str, hf_config, @@ -165,10 +172,7 @@ def apply_distributed_support( :func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep`, so nothing about it is per-model. """ - from modelopt.torch.models import get_spec - - spec = get_spec(model_type) - ep = getattr(spec, "ep_spec", None) if spec is not None else None + ep = _ep_spec(model_type) if ep is None: return hf_config @@ -204,10 +208,7 @@ def _spec_ep_plan(model_type: str) -> dict[str, str]: expert-parallel section, or one that refuses expert parallelism. Only reads the spec: nothing is applied, checked or patched, and ``prepare`` does not run. """ - from modelopt.torch.models import get_spec - - spec = get_spec(model_type) - ep = getattr(spec, "ep_spec", None) if spec is not None else None + ep = _ep_spec(model_type) if ep is None or ep.unsupported: return {} return dict(ep.base_model_ep_plan or ep.prepared_ep_plan or {}) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 49ecb88d795..bc614becac4 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -1184,6 +1184,7 @@ def _decide_distribution( """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone.""" from modelopt.torch.utils.moe import ( _ep_config, + _ep_spec, _routed_expert_modules, _spec_ep_plan, choose_ep_dp_degrees, @@ -1254,11 +1255,20 @@ def _decide_distribution( and dense == "replicate" and (experts or _routed_expert_modules(model, _spec_ep_plan(hf_config.model_type))) ): + fsdp2 = "(--use_fsdp2 in examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world)" + ep = _ep_spec(hf_config.model_type) + if ep is not None and ep.unsupported: + # Its spec refuses an EP degree above 1, so FSDP2 is the one way left to shard it. + raise RuntimeError( + f"{hf_config.model_type} cannot be quantized expert-parallel yet, and the plan " + f"replicates it at world_size={world_size}, so every rank would hold the whole " + f"MoE. Shard the whole model with FSDP2 instead {fsdp2}." + ) raise RuntimeError( f"{hf_config.model_type} can be expert-parallel, but the plan resolved to ep_size=1 at " f"world_size={world_size}, so every rank would hold the whole MoE. Shard its experts " "(MODELOPT_EP_DEGREE above 1, dividing both the expert count and the world size), or " - "shard the whole model with FSDP2." + f"shard the whole model with FSDP2 {fsdp2}." ) return DistributionPlan(world_size, ep_size, dense) diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index f98e74a3d57..19b6834fc04 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -440,21 +440,33 @@ def test_a_run_without_expert_parallelism_needs_no_spec(run_as, serve, monkeypat assert resolve_distribution_plan(str(tmp_path), dense="fsdp")[0] == plan +# The refusals of an MoE left whole on every rank, by whether its spec allows expert parallelism. +_CAN_BE_EP = "can be expert-parallel, but the plan resolved" +_NO_EP_YET = "cannot be quantized expert-parallel yet, and the plan replicates" + + @pytest.mark.parametrize( - ("config", "ep_degree"), + ("config", "ep_degree", "refusal"), [ - pytest.param(_deepseek_v3(n_routed_experts=3), None, id="auto_finds_no_degree"), - pytest.param(_qwen3_moe(num_experts=8), "1", id="explicit_degree_one"), - pytest.param(_deepseek_v3(n_routed_experts=8), "1", id="explicit_degree_one_spec_plan"), + pytest.param(_deepseek_v3(n_routed_experts=3), None, _CAN_BE_EP, id="auto_finds_no_degree"), + pytest.param(_qwen3_moe(num_experts=8), "1", _CAN_BE_EP, id="explicit_degree_one"), + pytest.param( + _deepseek_v3(n_routed_experts=8), "1", _CAN_BE_EP, id="explicit_degree_one_spec_plan" + ), + pytest.param( + GptOssConfig(num_local_experts=4, **_TINY), "1", _NO_EP_YET, id="spec_refuses_ep" + ), ], ) def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( - run_as, serve, monkeypatch, tmp_path, config, ep_degree + run_as, serve, monkeypatch, tmp_path, config, ep_degree, refusal ): """No EP degree above 1 divides 3 experts over 2 ranks, and an explicit 1 asks for none, whether the plan comes with the config (Qwen3-MoE) or only from the model's spec, which an EP degree of 1 skips (DeepSeek-V3). With nothing sharded either, every rank would hold the whole - MoE: that must fail, not run.""" + MoE: that must fail, not run. GPT-OSS carries a plan too, but its spec refuses expert + parallelism, so this check refuses it as an MoE that cannot be expert-parallel yet; the spec's + own refusal, a NotImplementedError and so a RuntimeError too, words it differently.""" if ep_degree is not None: if config.model_type == "qwen3_moe" and Version(transformers.__version__) < Version("5.7"): pytest.skip("qwen3_moe's base_model_ep_plan ships in transformers 5.7") @@ -462,7 +474,7 @@ def test_an_ep_capable_moe_left_whole_on_every_rank_is_an_error( run_as(2) serve(config) - with pytest.raises(RuntimeError, match="can be expert-parallel, but the plan resolved"): + with pytest.raises(RuntimeError, match=refusal): resolve_distribution_plan(str(tmp_path)) From 26b338a13b05db97431999dc30e559cd8ebc9a69 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:19:45 +0000 Subject: [PATCH 15/23] Test on two real ranks that every rank applies the spec as rank 0 decided The test played rank 1 against a stand-in broadcast that answered with rank 0's decision. The pooled job that checks every rank takes rank 0's plan now covers the spec too: rank 1's own EP knob says the opposite of rank 0's, and rank 1 must apply DeepSeek-V3's spec exactly when rank 0 did, over the real broadcast, or it would load without the plan the others shard by. Signed-off-by: Shengliang Xu --- .../utils/test_resolve_distribution_plan.py | 89 +++++++++---------- 1 file changed, 42 insertions(+), 47 deletions(-) diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 19b6834fc04..e14d3deae50 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -686,63 +686,58 @@ def _knobs(**values): yield -def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected): +def _every_rank_takes_rank0s_plan_job(rank, size, ckpt, knobs, expected, spec_applied): with bounded_collectives(), _knobs(**knobs[rank]): - plan, _ = resolve_distribution_plan(ckpt) + plan, hf_config = resolve_distribution_plan(ckpt) assert plan == expected, (rank, plan) + if spec_applied is not None: + spec_plan = get_spec(hf_config.model_type).ep_spec.base_model_ep_plan + applied = getattr(hf_config, "base_model_ep_plan", None) == spec_plan + assert applied is spec_applied, (rank, applied) -def test_every_rank_takes_the_plan_rank0_decided(cpu_dist_workers, tmp_path): - """Over a real broadcast. Rank 1's own knobs, FSDP2 over the world, would decide another plan, - and ranks loading different plans would build mismatched meshes and hang.""" - LlamaConfig(**_TINY).save_pretrained(tmp_path) +@pytest.mark.parametrize( + ("config", "knobs", "expected", "spec_applied"), + [ + pytest.param( + LlamaConfig(**_TINY), + [{}, {"MODELOPT_FSDP_SIZE": "world"}], + DistributionPlan(2), + None, + id="rank1-would-shard", + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), + [{}, {"MODELOPT_EP_DEGREE": "1"}], + DistributionPlan(2, 2), + True, + id="rank1-would-skip-the-spec", + ), + pytest.param( + _deepseek_v3(n_routed_experts=8), + [{"MODELOPT_EP_DEGREE": "1", "MODELOPT_FSDP_SIZE": "world"}, {}], + DistributionPlan(2, dense="fsdp"), + False, + id="rank1-would-apply-the-spec", + ), + ], +) +def test_every_rank_takes_the_plan_rank0_decided( + cpu_dist_workers, tmp_path, config, knobs, expected, spec_applied +): + """Over real broadcasts. Rank 1's own knobs would decide another plan, or apply the model's spec + where rank 0 skipped it, or skip it where rank 0 applied it: ranks loading different plans, or + one without the expert-parallel plan the others shard by, build mismatched meshes and hang.""" + config.save_pretrained(tmp_path) cpu_dist_workers.run( _every_rank_takes_rank0s_plan_job, ckpt=str(tmp_path), - knobs=[{}, {"MODELOPT_FSDP_SIZE": "world"}], - expected=DistributionPlan(2), + knobs=knobs, + expected=expected, + spec_applied=spec_applied, ) -def _as_rank(monkeypatch, rank, **rank0_messages): - """Run as ``rank`` of a 4-rank run, answering its broadcasts; returns what this rank handed them. - - A non-zero rank receives, for each broadcast, ``rank0_messages[label]``, as rank 0 would have - sent it. Every rank reads the checkpoint as this one does. - """ - sent = [] - - def broadcast_object_list(message, group=None, group_src=None): - sent.append(message[0]) - if rank != 0: - message[0] = rank0_messages[message[0][0]] - - monkeypatch.setattr(dist, "is_distributed_run", lambda: True) - monkeypatch.setattr(dist, "size", lambda group=None: 4) - monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) - monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: rank) - monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 4) - monkeypatch.setattr(torch.distributed, "broadcast_object_list", broadcast_object_list) - monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike) - return sent - - -@pytest.mark.parametrize("rank0_applies", [True, False]) -def test_other_ranks_apply_the_spec_as_rank0_decided(monkeypatch, serve, tmp_path, rank0_applies): - """A rank whose own EP knob says otherwise follows rank 0: had it skipped the spec rank 0 - applied, it would load without the expert-parallel plan the others shard by.""" - monkeypatch.setenv("MODELOPT_EP_DEGREE", "auto" if not rank0_applies else "1") - rank0_plan = DistributionPlan(4, 4) if rank0_applies else DistributionPlan(4) - _as_rank(monkeypatch, 1, spec=("spec", rank0_applies, None), plan=("plan", rank0_plan, None)) - serve(_deepseek_v3(n_routed_experts=8)) - - plan, hf_config = resolve_distribution_plan(str(tmp_path)) - - assert plan == rank0_plan - spec_plan = get_spec("deepseek_v3").ep_spec.base_model_ep_plan - assert (getattr(hf_config, "base_model_ep_plan", None) == spec_plan) is rank0_applies - - def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt): isdir = os.path.isdir with ( From 5c8c91ac50d3a0b564ef4d54b9e8e2df6dc96dba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 16:30:41 +0000 Subject: [PATCH 16/23] Import the spec-plan test's modules at module scope test_the_plan_a_spec_would_supply imported modelopt.torch.models and modelopt.torch.utils.moe inside the test, for no circular-import or optional-dependency reason: the module's transformers guard already precedes its other modelopt imports. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_resolve_distribution_plan.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index e14d3deae50..7592dc8da7f 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -46,7 +46,9 @@ ) import modelopt.torch.utils.distributed as dist +from modelopt.torch import models from modelopt.torch.models import EPSpec, get_spec +from modelopt.torch.utils import moe from modelopt.torch.utils.distributed import DistributionPlan from modelopt.torch.utils.moe import ( _routed_expert_modules, @@ -483,9 +485,6 @@ def test_the_plan_a_spec_would_supply(monkeypatch): ``prepare`` attaches, read without running ``prepare``; nothing for a model with no spec or no expert-parallel section, and nothing for a spec that refuses expert parallelism even when it names a plan.""" - from modelopt.torch import models - from modelopt.torch.utils import moe - assert moe._spec_ep_plan("deepseek_v3") == get_spec("deepseek_v3").ep_spec.base_model_ep_plan assert moe._spec_ep_plan("llama") == {} assert moe._spec_ep_plan("no_such_model_type") == {} From fc276dba43819c9dc66272da36b98b93525a5f6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:58:44 +0000 Subject: [PATCH 17/23] Trim the model-spec EP tests: drop what the applier tests cover - test_some_model_declares_expert_parallel_support goes: with no spec registered, the named-model applier and decision tests fail first. - test_a_model_type_with_no_spec_is_left_alone goes: every multi-rank dense decision test runs the same early return. - test_the_plan_is_attached_when_the_config_has_none goes: test_the_spec_is_applied_before_the_decision attaches the real deepseek_v3 plan. - The refusal test is parametrized over gpt_oss, llama4 and llama4_text by name rather than over whichever specs refuse, so a lost refusal fails it instead of shrinking it. It replaces test_models_whose_own_plan_would_misroute_them_are_refused. - The two prepare probes are one test. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_ep_specs.py | 67 ++++++------------------ 1 file changed, 16 insertions(+), 51 deletions(-) diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 4a409bacade..d0d98a5aa1a 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -25,7 +25,6 @@ from modelopt.torch.utils.moe import apply_distributed_support EP_MODEL_TYPES = [s.model_type for s in get_specs() if s.ep_spec is not None] -EP_REFUSED_MODEL_TYPES = [mt for mt in EP_MODEL_TYPES if get_spec(mt).ep_spec.unsupported] def cfg(**kw): @@ -35,11 +34,6 @@ def cfg(**kw): # --------------------------------------------------------------------------- the registered specs -def test_some_model_declares_expert_parallel_support(): - """A guard on the wiring itself: an empty list here means no spec reached the registry.""" - assert EP_MODEL_TYPES - - @pytest.mark.parametrize("model_type", EP_MODEL_TYPES) def test_every_ep_plan_is_well_formed(model_type): """Globs map to style names; a malformed entry matches nothing and fails silently at load.""" @@ -61,12 +55,6 @@ def test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal(model_type): ) -def test_models_whose_own_plan_would_misroute_them_are_refused(): - """GPT-OSS and Llama-4 ship plans of their own, so only their refusal keeps them off the - expert-parallel path.""" - assert {"gpt_oss", "llama4", "llama4_text"} <= set(EP_REFUSED_MODEL_TYPES) - - def test_the_omni_vlm_reuses_the_text_models_plan(): """Two spellings of one architecture must not drift into two different plans.""" assert ( @@ -84,18 +72,6 @@ def test_the_omni_vlm_puts_its_plan_on_the_nested_config(): # ----------------------------------------------------------------------------------- the applier -def test_a_model_type_with_no_spec_is_left_alone(): - config = cfg() - assert apply_distributed_support("not_a_real_model", config) is config - assert not hasattr(config, "base_model_ep_plan") - - -def test_the_plan_is_attached_when_the_config_has_none(): - config = cfg() - apply_distributed_support("deepseek_v3", config) - assert config.base_model_ep_plan == get_spec("deepseek_v3").ep_spec.base_model_ep_plan - - def test_the_models_own_plan_always_wins(): """Stock transformers ships a plan for some models; overwriting it would substitute our idea of the layout for the model's own.""" @@ -127,8 +103,10 @@ def test_a_missing_nested_config_is_an_error_not_a_silent_skip(): apply_distributed_support("nemotron_h_omni", cfg()) -@pytest.mark.parametrize("model_type", EP_REFUSED_MODEL_TYPES) +@pytest.mark.parametrize("model_type", ["gpt_oss", "llama4", "llama4_text"]) def test_an_unsupported_model_is_refused_with_its_config_untouched(model_type): + """These ship plans of their own, so only their spec's refusal keeps them off the + expert-parallel path.""" config = cfg() with pytest.raises(NotImplementedError, match="cannot be quantized expert-parallel yet"): @@ -158,35 +136,14 @@ def test_a_refusal_comes_before_the_rest_of_the_section(monkeypatch): assert vars(config) == {} -def test_prepare_runs_and_its_config_is_returned(monkeypatch): +def test_prepare_runs_after_the_plan_is_attached_and_its_config_is_returned(monkeypatch): + """A prepare hook may inspect or adjust the plan, so it must not run before there is one.""" seen = {} def prepare(hf_config, **kwargs): - seen.update(kwargs) - seen["config"] = hf_config + seen.update(kwargs, config=hf_config, plan_present=hasattr(hf_config, "base_model_ep_plan")) return "replaced" - spec = ModelSpec(model_type="_probe", ep_spec=EPSpec(prepare=prepare)) - monkeypatch.setattr( - "modelopt.torch.models.get_spec", lambda mt: spec if mt == "_probe" else None - ) - - config = cfg() - assert ( - apply_distributed_support("_probe", config, ckpt_path="/c", trust_remote_code=True) - == "replaced" - ) - assert seen == {"config": config, "ckpt_path": "/c", "trust_remote_code": True} - - -def test_prepare_runs_after_the_plan_is_attached(monkeypatch): - """A prepare hook may inspect or adjust the plan, so it must not run before there is one.""" - observed = {} - - def prepare(hf_config, **_): - observed["plan_present"] = hasattr(hf_config, "base_model_ep_plan") - return hf_config - spec = ModelSpec( model_type="_probe", ep_spec=EPSpec(base_model_ep_plan={"layers.*.experts": "grouped_gemm"}, prepare=prepare), @@ -194,6 +151,14 @@ def prepare(hf_config, **_): monkeypatch.setattr( "modelopt.torch.models.get_spec", lambda mt: spec if mt == "_probe" else None ) + config = cfg() + + returned = apply_distributed_support("_probe", config, ckpt_path="/c", trust_remote_code=True) - apply_distributed_support("_probe", cfg()) - assert observed["plan_present"] + assert returned == "replaced" + assert seen == { + "config": config, + "ckpt_path": "/c", + "trust_remote_code": True, + "plan_present": True, + } From ee095f7c05e32c07af4e38231bac66040756a68c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:47:09 +0000 Subject: [PATCH 18/23] List EPSpec in the model specs' public API modelopt.torch.models exports EPSpec, but specs.py, where it is defined, left it out of its __all__. Signed-off-by: Shengliang Xu --- modelopt/torch/models/specs.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/torch/models/specs.py b/modelopt/torch/models/specs.py index a901c1bc524..77e26366d4b 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -50,6 +50,7 @@ import torch.nn as nn __all__ = [ + "EPSpec", "ExportSpec", "MoESpec", "ModelSpec", From f48b37e30fd7993250d0971a5783c896c755ac14 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 20 Oct 2026 01:51:31 +0000 Subject: [PATCH 19/23] Read an expert-parallel load's config with the Hub options too The expert-parallel route of a sharded load handed a plan but no config read the config without token, cache_dir or local_files_only, as the FSDP2 route did before. It now reads it with the Hub options the load was given. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/plugins/model_load_utils.py | 4 +++- .../utils/test_load_pretrained_distribution.py | 13 +++++++++---- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index bc614becac4..38e40834c54 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -466,7 +466,9 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op if hf_config is None: hf_config = apply_distributed_support( model_type, - AutoConfig.from_pretrained(ckpt_path, trust_remote_code=trust_remote_code), + AutoConfig.from_pretrained( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ), ckpt_path=ckpt_path, trust_remote_code=trust_remote_code, ) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 289788b332c..a15ba5fde10 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -305,12 +305,16 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( ids=["decided_by_the_load", "planned", "planned_with_its_config"], ) def test_a_load_prepares_its_config_once( - run_as, serve, loaders, monkeypatch, tmp_path, planned, given_config, preparations + run_as, loaders, monkeypatch, tmp_path, planned, given_config, preparations ): """A model's EP preparation can patch its modeling code, so a load runs it once, and not at all - when handed the config resolve_distribution_plan prepared.""" + when handed the config resolve_distribution_plan prepared. Every config it reads takes the Hub + options the load was given (custom code an ``auto_map`` names in another repo).""" run_as(4) - config = serve(_deepseek_v3(n_routed_experts=8)) + config, reads = _deepseek_v3(n_routed_experts=8), [] + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(k) or config + ) prepared = [] def apply_distributed_support(model_type, hf_config, **kwargs): @@ -323,10 +327,11 @@ def apply_distributed_support(model_type, hf_config, **kwargs): if given_config: kwargs["hf_config"] = config - load_pretrained(str(tmp_path), **kwargs) + load_pretrained(str(tmp_path), local_files_only=True, **kwargs) assert len(prepared) == preparations assert loaders["ep"]["hf_config"] is config + assert all(read.get("local_files_only") is True for read in reads) def _every_rank_loads_rank0s_plan_job(rank, size, ckpt): From 0a7357b0bf49da442916de792e27451a7bb8458e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 19:05:35 +0000 Subject: [PATCH 20/23] Date the nemotron_h_omni spec package 2026, as LICENSE_HEADER does Signed-off-by: Shengliang Xu --- modelopt/torch/models/nemotron_h_omni/__init__.py | 2 +- modelopt/torch/models/nemotron_h_omni/specs.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/models/nemotron_h_omni/__init__.py b/modelopt/torch/models/nemotron_h_omni/__init__.py index 5b9b820672c..8e852a66f9d 100644 --- a/modelopt/torch/models/nemotron_h_omni/__init__.py +++ b/modelopt/torch/models/nemotron_h_omni/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); diff --git a/modelopt/torch/models/nemotron_h_omni/specs.py b/modelopt/torch/models/nemotron_h_omni/specs.py index 3d9b6dce4cb..7e9f4bcddc6 100644 --- a/modelopt/torch/models/nemotron_h_omni/specs.py +++ b/modelopt/torch/models/nemotron_h_omni/specs.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# 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"); From 3687b9a628a99bc337975a68997f7c0a85753cdf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:20:51 +0000 Subject: [PATCH 21/23] Prepare a config passed with a handed-in expert-parallel plan A config passed for the load is the caller's own. With the plan decided by the load, resolve_distribution_plan prepares a copy of it through the model spec. With a plan handed in, the expert-parallel load took the config raw, so the spec's preparation never ran: DeepSeek-V3 failed for want of the expert-parallel plan the spec adds, and a model the spec refuses was not refused. The load now prepares a copy, once, as deciding does. Signed-off-by: Shengliang Xu --- .../torch/utils/plugins/model_load_utils.py | 22 ++++++++++---- .../test_load_pretrained_distribution.py | 30 +++++++++++++------ 2 files changed, 38 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 38e40834c54..7a94b8182e5 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -535,7 +535,8 @@ def load_pretrained( rather than ignore it. ``hf_config`` is the config :func:`resolve_distribution_plan` returned with the plan handed in as ``distribution``: the load builds from it as given, so without that plan it is refused. A ``config`` passed for the load is the one the plan is decided from, and a - sharded load takes it as its ``hf_config``. + sharded load takes it as its ``hf_config``: an expert-parallel load a copy of it the model spec + prepares (:func:`~modelopt.torch.utils.moe.apply_distributed_support`), as deciding does. Whatever the placement: @@ -613,10 +614,6 @@ def checked() -> DistributionPlan: else: raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") route = _route(plan) - if route is not None and config is not None: - # The sharded load builds the model from hf_config: the plan's copy of this one, or it. - del kwargs["config"] - kwargs.setdefault("hf_config", config) if model_type is None: # The load uses the config passed for it, if any, so the loading hook follows that one. model_type = getattr(hf_config or config, "model_type", None) or ( @@ -624,6 +621,21 @@ def checked() -> DistributionPlan: ckpt_path, trust_remote_code=trust_remote_code, **hub_options ).model_type ) + if route is not None and config is not None: + # The sharded load builds the model from hf_config: the plan's copy of this one, or it. With + # a plan handed in, an expert-parallel load prepares a copy of it, as deciding would have. + del kwargs["config"] + if "hf_config" not in kwargs: + if route == _EXPERT_PARALLEL: + from modelopt.torch.utils.moe import apply_distributed_support + + config = apply_distributed_support( + model_type, + copy.deepcopy(config), + ckpt_path=ckpt_path, + trust_remote_code=trust_remote_code, + ) + kwargs["hf_config"] = config if route is not None: model = _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_options) # A sharded load finds the unplaced keys while it places the weights. diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index a15ba5fde10..6788cc25a7a 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -300,18 +300,25 @@ def test_an_expert_parallel_plan_loads_through_the_ep_loader( @pytest.mark.parametrize( - ("planned", "given_config", "preparations"), - [(False, False, 1), (True, False, 1), (True, True, 0)], - ids=["decided_by_the_load", "planned", "planned_with_its_config"], + ("planned", "passed_as", "preparations"), + [(False, None, 1), (True, None, 1), (True, "hf_config", 0), (True, "config", 1)], + ids=[ + "decided_by_the_load", + "planned", + "planned_with_its_config", + "planned_with_a_config_for_the_load", + ], ) def test_a_load_prepares_its_config_once( - run_as, loaders, monkeypatch, tmp_path, planned, given_config, preparations + run_as, loaders, monkeypatch, tmp_path, planned, passed_as, preparations ): """A model's EP preparation can patch its modeling code, so a load runs it once, and not at all - when handed the config resolve_distribution_plan prepared. Every config it reads takes the Hub - options the load was given (custom code an ``auto_map`` names in another repo).""" + when handed the config resolve_distribution_plan prepared. A config passed for the load is the + caller's own, so the load prepares a copy of it. Every config it reads takes the Hub options the + load was given (custom code an ``auto_map`` names in another repo).""" run_as(4) config, reads = _deepseek_v3(n_routed_experts=8), [] + as_passed = config.to_dict() monkeypatch.setattr( transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(k) or config ) @@ -324,13 +331,18 @@ def apply_distributed_support(model_type, hf_config, **kwargs): monkeypatch.setattr(moe, "apply_distributed_support", apply_distributed_support) kwargs = {"distribution": DistributionPlan(4, ep_size=4)} if planned else {} - if given_config: - kwargs["hf_config"] = config + if passed_as: + kwargs[passed_as] = config load_pretrained(str(tmp_path), local_files_only=True, **kwargs) assert len(prepared) == preparations - assert loaders["ep"]["hf_config"] is config + loaded = loaders["ep"]["hf_config"] + if passed_as == "config": + assert loaded is not config and loaded.base_model_ep_plan + assert config.to_dict() == as_passed + else: + assert loaded is config assert all(read.get("local_files_only") is True for read in reads) From f0768166df95f6a17689eacd07a71eb687c218a4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:20:55 +0000 Subject: [PATCH 22/23] Refuse what a model's spec refuses for a config handed in with its plan The config resolve_distribution_plan prepared comes back to the load with its plan, and the load does not prepare it again: preparing can patch modeling code. It skipped the spec's checks with the preparation, so a hand-made plan and config loaded GPT-OSS or Llama-4 expert-parallel although their specs refuse it. check_distributed_support now runs the spec's refusal and transformers floor on their own, and apply_distributed_support starts with it. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 27 +++++++++++++------ .../torch/utils/plugins/model_load_utils.py | 6 ++++- .../test_load_pretrained_distribution.py | 15 ++++++++++- 3 files changed, 38 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 07a5e11a743..ac77951582c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -148,6 +148,24 @@ def _ep_spec(model_type: str): return getattr(get_spec(model_type), "ep_spec", None) +def check_distributed_support(model_type: str) -> None: + """Refuse ``model_type``'s expert parallelism if its spec does; enforce its transformers floor. + + :func:`apply_distributed_support` checks this before it prepares a config; a config already + prepared with its plan needs only the check. + """ + ep = _ep_spec(model_type) + if ep is None: + return + if ep.unsupported: + raise NotImplementedError( + f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " + "without expert parallelism, e.g. with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py)." + ) + if ep.requires_transformers: + require_transformers(ep.requires_transformers, f"{model_type} expert-parallel PTQ") + + def apply_distributed_support( model_type: str, hf_config, @@ -172,18 +190,11 @@ def apply_distributed_support( :func:`~modelopt.torch.utils.plugins.model_load_utils.load_and_prepare_ep`, so nothing about it is per-model. """ + check_distributed_support(model_type) ep = _ep_spec(model_type) if ep is None: return hf_config - if ep.unsupported: - raise NotImplementedError( - f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " - "without expert parallelism, e.g. with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py)." - ) - if ep.requires_transformers: - require_transformers(ep.requires_transformers, f"{model_type} expert-parallel PTQ") - if ep.base_model_ep_plan: target = hf_config if ep.plan_config_attr: diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 7a94b8182e5..50645c496fb 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -461,7 +461,7 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op trust_remote_code = kwargs.pop("trust_remote_code", False) hf_config = kwargs.pop("hf_config", None) if route == _EXPERT_PARALLEL: - from modelopt.torch.utils.moe import apply_distributed_support + from modelopt.torch.utils.moe import apply_distributed_support, check_distributed_support if hf_config is None: hf_config = apply_distributed_support( @@ -472,6 +472,10 @@ def _load_sharded(route, plan, ckpt_path, auto_class, model_type, kwargs, hub_op ckpt_path=ckpt_path, trust_remote_code=trust_remote_code, ) + else: + # Prepared with its plan, and preparing can patch modeling code, so it is not prepared + # again; a model its spec refuses is refused all the same. + check_distributed_support(model_type) return load_and_prepare_ep( ckpt_path, plan, trust_remote_code=trust_remote_code, hf_config=hf_config, **kwargs ) diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 6788cc25a7a..8d27af338b0 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -29,7 +29,7 @@ transformers = pytest.importorskip("transformers") -from transformers import AutoModel, DeepseekV3Config, LlamaConfig, LlamaForCausalLM +from transformers import AutoModel, DeepseekV3Config, GptOssConfig, LlamaConfig, LlamaForCausalLM import modelopt.torch.models.hf as models_hf import modelopt.torch.utils.distributed as dist @@ -346,6 +346,19 @@ def apply_distributed_support(model_type, hf_config, **kwargs): assert all(read.get("local_files_only") is True for read in reads) +def test_a_model_its_spec_refuses_is_refused_with_its_config_handed_in(run_as, loaders, tmp_path): + """Handed the config prepared with its plan, the load does not prepare it again, but a model + whose spec refuses expert parallelism is refused all the same, before anything loads.""" + run_as(4) + config = GptOssConfig(**_TINY, num_local_experts=4) + + with pytest.raises(NotImplementedError, match="gpt_oss cannot be quantized expert-parallel"): + load_pretrained( + str(tmp_path), distribution=DistributionPlan(4, ep_size=4), hf_config=config + ) + assert loaders == {} + + def _every_rank_loads_rank0s_plan_job(rank, size, ckpt): handed = DistributionPlan(size) if rank == 0 else DistributionPlan(size, dense="fsdp") with bounded_collectives(): From 233cb7ebf3a265a78f4fe1e45ffabfe0f517db7b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 23 Oct 2026 07:08:23 +0000 Subject: [PATCH 23/23] Name the way around expert parallelism when a model is below its transformers floor A multi-rank run that leaves the EP degree to auto checks the model's spec before its plan is decided, so a model whose expert parallelism needs a newer transformers (Qwen3-MoE below 5.7) stops there. Like the refusal of a model that cannot be expert-parallel yet, the error now says it can also be quantized without expert parallelism. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index ac77951582c..cc1b074d4ab 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -152,18 +152,26 @@ def check_distributed_support(model_type: str) -> None: """Refuse ``model_type``'s expert parallelism if its spec does; enforce its transformers floor. :func:`apply_distributed_support` checks this before it prepares a config; a config already - prepared with its plan needs only the check. + prepared with its plan needs only the check. A multi-rank run that leaves the EP degree to + ``auto`` gets here before its plan is decided, so both errors name the way around expert + parallelism. """ ep = _ep_spec(model_type) if ep is None: return + without_ep = ( + "without expert parallelism, e.g. with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py)" + ) if ep.unsupported: raise NotImplementedError( f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " - "without expert parallelism, e.g. with FSDP2 (--use_fsdp2 in examples/hf_ptq/hf_ptq.py)." + f"{without_ep}." ) if ep.requires_transformers: - require_transformers(ep.requires_transformers, f"{model_type} expert-parallel PTQ") + try: + require_transformers(ep.requires_transformers, f"{model_type} expert-parallel PTQ") + except RuntimeError as e: + raise RuntimeError(f"{e} Or quantize it {without_ep}.") from None def apply_distributed_support(