diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index 98aaf9e57da..7d71f9e106f 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,9 @@ phimoe, nemotron, nemotron_h, + minimax_m3_vl, + llama4_text, + 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..a4b95c790e9 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,13 @@ ModelSpec( model_type="gpt_oss", min_transformers_version="5.5", + ep_spec=EPSpec( + unsupported=( + "GPT-OSS stores its experts transposed, (num_experts, in, out), with gate and up " + "interleaved in gate_up_proj, and the expert-parallel export cannot write that " + "layout." + ), + ), moe_spec=MoESpec( # GPT-OSS fuses gate and up into a single gate_up_proj. # transformers names the block GptOssMLP; GptOssMoE is kept for the diff --git a/modelopt/torch/models/llama4/specs.py b/modelopt/torch/models/llama4/specs.py index dc733398d87..ace9c0b8d69 100644 --- a/modelopt/torch/models/llama4/specs.py +++ b/modelopt/torch/models/llama4/specs.py @@ -15,10 +15,18 @@ """Llama4 specs (HF model type ``llama4``).""" -from ..specs import ModelSpec, MoESpec, register +from ..specs import EPSpec, ModelSpec, MoESpec, register __all__: list[str] = [] +# transformers ships a base_model_ep_plan for Llama-4, which alone would route it into the +# expert-parallel path. +EP_UNSUPPORTED = ( + "Llama-4's MoE block routes tokens itself and runs every expert densely, rather than " + "dispatching its top-k through the experts interface, and its experts are stored transposed, " + "(num_experts, in, out), which the expert-parallel export cannot write." +) + # Llama4TextExperts is fused: one module holding 3-D gate_up_proj and down_proj # parameters, run through torch.bmm. There is no (gate, up) pair left to fuse, and the # grouped-export path (get_experts_list) does not apply. The block also carries a @@ -27,6 +35,7 @@ ModelSpec( model_type="llama4", min_transformers_version="5.5", + ep_spec=EPSpec(unsupported=EP_UNSUPPORTED), moe_spec=MoESpec( block_names=("Llama4TextMoe",), expert_linear_names=("gate_up_proj", "down_proj"), 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..7440e632063 --- /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``, with its own +``base_model_ep_plan`` on the root config, so it is refused expert parallelism the same way. +""" + +from ..llama4.specs import EP_UNSUPPORTED +from ..specs import EPSpec, ModelSpec, register + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="llama4_text", + min_transformers_version="5.5", + ep_spec=EPSpec(unsupported=EP_UNSUPPORTED), + ) +) 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/specs.py b/modelopt/torch/models/minimax_m3_vl/specs.py new file mode 100644 index 00000000000..025d2f7b22a --- /dev/null +++ b/modelopt/torch/models/minimax_m3_vl/specs.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""MiniMax-M3 specs (HF model type ``minimax_m3_vl``). + +transformers ships everything M3 needs to run expert-parallel: its MoE block calls the experts with +the router's top-k, its experts dispatch through the experts interface, and its text config carries a +``base_model_ep_plan``. The generic all-to-all transport takes over unchanged. +""" + +from ..specs import EPSpec, ModelSpec, register + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="minimax_m3_vl", + min_transformers_version="5.13", + ep_spec=EPSpec(requires_transformers="5.13"), + ) +) diff --git a/modelopt/torch/models/nemotron_h/specs.py b/modelopt/torch/models/nemotron_h/specs.py index 7198c2b31b3..8356310948c 100644 --- a/modelopt/torch/models/nemotron_h/specs.py +++ b/modelopt/torch/models/nemotron_h/specs.py @@ -15,15 +15,26 @@ """Nemotron-H specs (HF model type ``nemotron_h``).""" -from ..specs import ExportSpec, ModelSpec, MoESpec, register +from ..specs import EPSpec, ExportSpec, ModelSpec, MoESpec, register __all__: list[str] = [] +# Ungated experts, so ``up_proj``/``down_proj`` rather than a gate/up pair. The gate emits logits and +# the MoE block runs its own group-limited top-k, so the plan has no router entry and the experts +# receive global expert ids -- which the expert-parallel all-to-all experts forward routes on. The +# loader maps ``nemotron_experts_ep``, a name no transformers style carries, to that forward. +BASE_MODEL_EP_PLAN = { + "layers.*.mixer.experts.up_proj": "grouped_gemm", + "layers.*.mixer.experts.down_proj": "grouped_gemm", + "layers.*.mixer.experts": "nemotron_experts_ep", +} + register( ModelSpec( model_type="nemotron_h", min_transformers_version="5.5", export_spec=ExportSpec(grouped_expert_export=True), + ep_spec=EPSpec(base_model_ep_plan=BASE_MODEL_EP_PLAN), moe_spec=MoESpec( # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). block_names=("NemotronHMOE",), 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..8e852a66f9d --- /dev/null +++ b/modelopt/torch/models/nemotron_h_omni/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""NemotronH-Omni specs (HF model type ``nemotron_h_omni``).""" + +from . import specs 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..7e9f4bcddc6 --- /dev/null +++ b/modelopt/torch/models/nemotron_h_omni/specs.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""NemotronH-Omni specs (HF model type ``nemotron_h_omni``). + +The Omni VLM's language sub-model is a native NemotronH, so its MoE facts are NemotronH's. +What differs is where the expert-parallel plan has to be attached: on the nested ``llm_config``, +not the VLM root. transformers aggregates each sub-model's plan up to the root WITH its module +prefix, so a plan on ``llm_config`` surfaces as ``language_model..layers.*`` and matches the +real parameter paths -- while a plan on the top config reaches no expert at all. +""" + +from ..nemotron_h.specs import BASE_MODEL_EP_PLAN +from ..specs import EPSpec, ModelSpec, register + +__all__: list[str] = [] + +register( + ModelSpec( + model_type="nemotron_h_omni", + # transformers ships the Omni VLM itself from 5.18. Expert parallelism still runs on the + # checkpoint's own code (``trust_remote_code=True``), whose sub-model config is the + # ``llm_config`` the plan below targets; transformers' port names it ``text_config``. + min_transformers_version="5.18", + ep_spec=EPSpec( + base_model_ep_plan=BASE_MODEL_EP_PLAN, + plan_config_attr="llm_config", + ), + ) +) 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..77e26366d4b 100644 --- a/modelopt/torch/models/specs.py +++ b/modelopt/torch/models/specs.py @@ -44,9 +44,13 @@ 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__ = [ + "EPSpec", "ExportSpec", "MoESpec", "ModelSpec", @@ -155,6 +159,75 @@ 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=...)`` 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.""" + + 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. + + 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): """Subsystem section: per-model data of the unified HF export path. @@ -233,6 +306,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..cc1b074d4ab 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -141,6 +141,98 @@ 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 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. 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 " + f"{without_ep}." + ) + if ep.requires_transformers: + 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( + model_type: str, + hf_config, + *, + ckpt_path: "str | None" = None, + trust_remote_code: 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:`~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.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) + 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``. + + 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. + """ + 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 {}) + + # --------------------------------------------------------------------------------------------------- # 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 1d28d24b13d..50645c496fb 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -461,6 +461,21 @@ 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, check_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, **hub_options + ), + 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 ) @@ -524,7 +539,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: @@ -602,10 +618,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 ( @@ -613,6 +625,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. @@ -1175,7 +1202,9 @@ 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, num_experts_from_config, ) @@ -1236,12 +1265,28 @@ 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))) + ): + 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) @@ -1289,9 +1334,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 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. """ 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 +1367,20 @@ def checkpoint_read() -> bool: dist.all_reduce("checkpoint read", checkpoint_read, all) ckpt_path, hf_config = read[0] world_size = dist.size() + # 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, + 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/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py new file mode 100644 index 00000000000..d0d98a5aa1a --- /dev/null +++ b/tests/unit/torch/models/test_ep_specs.py @@ -0,0 +1,164 @@ +# 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 + +pytest.importorskip("transformers") + +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 + + +@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_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 or ep.unsupported, ( + 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_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()) + + +@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"): + 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_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, config=hf_config, plan_present=hasattr(hf_config, "base_model_ep_plan")) + return "replaced" + + 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 + ) + config = cfg() + + returned = apply_distributed_support("_probe", config, ckpt_path="/c", trust_remote_code=True) + + assert returned == "replaced" + assert seen == { + "config": config, + "ckpt_path": "/c", + "trust_remote_code": True, + "plan_present": True, + } diff --git a/tests/unit/torch/utils/test_load_pretrained_distribution.py b/tests/unit/torch/utils/test_load_pretrained_distribution.py index 8429c4cc58e..8d27af338b0 100644 --- a/tests/unit/torch/utils/test_load_pretrained_distribution.py +++ b/tests/unit/torch/utils/test_load_pretrained_distribution.py @@ -29,10 +29,12 @@ 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 +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,90 @@ 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", "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, 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. 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 + ) + 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 passed_as: + kwargs[passed_as] = config + + load_pretrained(str(tmp_path), local_files_only=True, **kwargs) + + assert len(prepared) == preparations + 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) + + +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(): diff --git a/tests/unit/torch/utils/test_resolve_distribution_plan.py b/tests/unit/torch/utils/test_resolve_distribution_plan.py index 177334cef62..7592dc8da7f 100644 --- a/tests/unit/torch/utils/test_resolve_distribution_plan.py +++ b/tests/unit/torch/utils/test_resolve_distribution_plan.py @@ -38,12 +38,17 @@ AutoModel, AutoModelForCausalLM, AutoModelForImageTextToText, + DeepseekV3Config, + GptOssConfig, LlamaConfig, Qwen2_5_VLConfig, Qwen3MoeConfig, ) 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, @@ -73,6 +78,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 +405,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,36 +416,101 @@ 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 + + +# 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(_qwen3_moe(num_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=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. 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. 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") monkeypatch.setenv("MODELOPT_EP_DEGREE", ep_degree) 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)) -@needs_qwen3_moe_ep_plan +def test_the_plan_a_spec_would_supply(monkeypatch): + """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.""" + 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") == {} + 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): """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") @@ -596,21 +685,55 @@ 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, )