From 8f86e3212a147f8d97dded08ccdfb07233224ded Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:02:25 +0000 Subject: [PATCH 1/6] Prepare a model's config for expert parallelism from its spec apply_distributed_support reads a model's EPSpec and prepares the config it is given for the expert-parallel load. It checks the transformers floor: the load's own 5.7, or the spec's later release, for a model with no spec too. It attaches the spec's base_model_ep_plan to the config the load reads it from, unless the model already carries a plan wherever the load reads one, and refuses a placement the load would not read. Then it runs the spec's prepare. check_distributed_support is the floor and refusal check alone, _spec_refuses_ep its predicate, and _spec_ep_plan the plan a spec would supply, read without running prepare. The nested configs the load reads a plan and an expert count from are one tuple, _SUB_CONFIG_ATTRS. Nothing calls them yet; the expert-parallel decision and load do in a later change. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 146 +++++++++++++- tests/unit/torch/models/test_ep_specs.py | 237 ++++++++++++++++++++++- 2 files changed, 375 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 2aff08e6891..a12bb87dee9 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -35,6 +35,10 @@ if TYPE_CHECKING: from torch import nn +# The transformers release the expert-parallel load needs, for every model: 5.6 builds a model only +# with the experts implementations it ships, never torch_a2a. +_EP_MIN_TRANSFORMERS = "5.7" + def require_transformers(min_version: str, feature: str) -> None: """Fail loudly if the installed transformers is older than ``min_version`` (the floor for ``feature``). @@ -55,6 +59,11 @@ def require_transformers(min_version: str, feature: str) -> None: ) +# The nested configs a VLM's expert-parallel plan and expert count are read from. A spec attaching +# its plan anywhere else would be read as having none, and the MoE planned as a dense model. +_SUB_CONFIG_ATTRS = ("text_config", "llm_config") + + def _ep_config(hf_config): """Return the config that carries the EP plan / expert count. @@ -64,7 +73,7 @@ def _ep_config(hf_config): """ if getattr(hf_config, "base_model_ep_plan", None): return hf_config - for _attr in ("text_config", "llm_config"): + for _attr in _SUB_CONFIG_ATTRS: sub = getattr(hf_config, _attr, None) if sub is not None and getattr(sub, "base_model_ep_plan", None): return sub @@ -78,11 +87,7 @@ def _num_experts_from_config(hf_config) -> int: Recurses into ``.text_config`` / ``.llm_config`` so VLMs whose nested text submodel is the MoE (Kimi-K2.7, Llama-4 multimodal, NemotronH-Omni, ...) report the LM's expert count. """ - for cfg in ( - hf_config, - getattr(hf_config, "text_config", None), - getattr(hf_config, "llm_config", None), - ): + for cfg in (hf_config, *(getattr(hf_config, attr, None) for attr in _SUB_CONFIG_ATTRS)): if cfg is None: continue for attr in ("num_experts", "num_local_experts", "n_routed_experts", "num_routed_experts"): @@ -142,6 +147,135 @@ def _routed_expert_modules(model: "nn.Module", ep_plan: dict[str, str]) -> "dict return found +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, or below its transformers floor. + + The floor is the expert-parallel load's own, or its spec's release when that is later; a model + with no spec or no EP section still needs the load's. + :func:`apply_distributed_support` checks this before it prepares a config; a config already + prepared with its plan needs only the check. An EP degree left to ``auto`` picks none for a + model this refuses (:func:`_spec_refuses_ep`), so only an EP degree asked for explicitly, or an + expert-parallel plan handed in, meets these errors; both name the way around expert parallelism. + """ + ep = _ep_spec(model_type) + # The FSDP2 load builds a causal LM and reads its tensors as stored, so a quantized or + # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process. + without_ep = ( + "without expert parallelism: in a single process, or FSDP2-sharded (--use_fsdp2 in " + "examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world) if the checkpoint is an " + "unquantized causal LM" + ) + if ep is not None and ep.unsupported: + raise NotImplementedError( + f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " + f"{without_ep}." + ) + from packaging.version import parse + + # A model whose spec states no release, or that has no spec, still needs the load's floor: auto + # must pick no EP degree below it, rather than reach the load and fail there. + stated = ep.requires_transformers if ep is not None else None + floor = max(_EP_MIN_TRANSFORMERS, stated or _EP_MIN_TRANSFORMERS, key=parse) + try: + require_transformers(floor, f"{model_type} expert-parallel PTQ") + except RuntimeError as e: + raise RuntimeError(f"{e} Or quantize it {without_ep}.") from None + + +def _spec_refuses_ep(model_type: str) -> bool: + """Whether :func:`check_distributed_support` refuses ``model_type``'s expert parallelism here.""" + try: + check_distributed_support(model_type) + except RuntimeError: # NotImplementedError is one too + return True + return False + + +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. It wins wherever the load reads + one (:func:`_ep_config`), not only on the config the spec would take: on the root, or on + another nested config. + + Returns the config to load, unchanged for a model type with no spec or no EP section. The plan + is written into the config it is given, so a caller that keeps its own config passes a copy, as + every caller here does. 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 + + # The spellings of the config the plan goes on; the first this config has takes it. Checked for + # every spec, before anything is attached or prepared, whatever supplies the plan. + attrs = ep.plan_config_attr or () + if isinstance(attrs, str): + attrs = (attrs,) + unread = [attr for attr in attrs if attr not in _SUB_CONFIG_ATTRS] + if unread: + raise ValueError( + f"{model_type}'s spec attaches its expert-parallel plan to " + f"config.{' or config.'.join(unread)}, where the expert-parallel load does not read " + f"one; it reads config.{' and config.'.join(_SUB_CONFIG_ATTRS)}." + ) + target = hf_config + if attrs: + target = next( + (sub for attr in attrs if (sub := getattr(hf_config, attr, None)) is not None), None + ) + if target is None: + raise ValueError( + f"{model_type} expects its expert-parallel plan on " + f"config.{' or config.'.join(attrs)}, which this config does not have." + ) + if ep.base_model_ep_plan and not getattr(_ep_config(hf_config), "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 whose spec refuses expert parallelism outright + (``unsupported``). A model below its spec's transformers floor keeps its plan: its experts are + routed all the same, so a run holding them whole on every rank is refused rather than + replicated. 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/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index d58f4145c42..74dbcade82d 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -13,15 +13,40 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The expert-parallel section of the model specs.""" +"""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 get_spec, get_specs +transformers = pytest.importorskip("transformers") + +from packaging.version import Version + +from modelopt.torch import models +from modelopt.torch.models import EPSpec, ModelSpec, get_spec, get_specs +from modelopt.torch.utils import moe +from modelopt.torch.utils.moe import ( + _EP_MIN_TRANSFORMERS, + _ep_config, + apply_distributed_support, + check_distributed_support, +) EP_MODEL_TYPES = [s.model_type for s in get_specs() if s.ep_spec is not None] +# Applying a spec enforces the transformers floor the expert-parallel load needs. +needs_ep_transformers = pytest.mark.skipif( + Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS), + reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}", +) + + +def cfg(**kw): + return SimpleNamespace(**kw) + + # --------------------------------------------------------------------------- the registered specs @@ -59,3 +84,211 @@ 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 == ("text_config", "llm_config") assert get_spec("nemotron_h").ep_spec.plan_config_attr is None + + +# ----------------------------------------------------------------------------------- the applier + + +@needs_ep_transformers +@pytest.mark.parametrize( + ("model_type", "config"), + [ + ("deepseek_v3", lambda own: cfg(base_model_ep_plan=own)), + ( + "nemotron_h_omni", + lambda own: cfg(text_config=cfg(), llm_config=cfg(base_model_ep_plan=own)), + ), + ("nemotron_h_omni", lambda own: cfg(base_model_ep_plan=own, text_config=cfg())), + ], + ids=["on-the-root", "on-another-nested-config", "on-a-vlms-root"], +) +def test_the_models_own_plan_always_wins(model_type, config): + """Stock transformers ships a plan for some models; overwriting it would substitute our idea + of the layout for the model's own. It wins wherever the load reads it, not only on the config + the spec would take: ours on text_config would shadow the model's on llm_config, which the + load reads after it, or lie unused beside a plan on the root.""" + own = {"layers.*.mlp.experts": "stock_style"} + config = config(own) + apply_distributed_support(model_type, config) + assert _ep_config(config).base_model_ep_plan == own + assert not hasattr(getattr(config, "text_config", cfg()), "base_model_ep_plan") + + +def test_a_model_without_a_spec_still_needs_the_loads_transformers(monkeypatch): + """A plan that ships with transformers needs no spec, but its load needs the release every + expert-parallel load does: auto picks no EP degree below it, rather than reach the load.""" + monkeypatch.setattr(transformers, "__version__", "5.6.2") + with pytest.raises(RuntimeError, match=r"requires transformers >= 5\.7.*without expert"): + check_distributed_support("no_such_model_type") + monkeypatch.setattr(transformers, "__version__", _EP_MIN_TRANSFORMERS) + check_distributed_support("no_such_model_type") + + +@needs_ep_transformers +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 + + +@needs_ep_transformers +@pytest.mark.parametrize( + "attrs", [("llm_config",), ("text_config",), ("text_config", "llm_config")] +) +def test_a_vlm_plan_lands_on_the_nested_config_only(attrs): + """On the root config the globs match no parameter, and the failure is silence, not an error. + The Omni checkpoint's own code names that config llm_config, and transformers' port + text_config: a multi-rank run of either applies the spec, whether it wants experts sharded or + not. A config with both takes it on text_config, which the load reads the expert count from + first.""" + config = cfg(**{attr: cfg() for attr in attrs}) + apply_distributed_support("nemotron_h_omni", config) + assert getattr(config, attrs[0]).base_model_ep_plan + others = [config, *(getattr(config, attr) for attr in attrs[1:])] + assert not any(hasattr(other, "base_model_ep_plan") for other in others) + + +@needs_ep_transformers +@pytest.mark.parametrize("by_prepare", [False, True], ids=["by-the-spec", "by-prepare"]) +def test_a_missing_nested_config_is_an_error_not_a_silent_skip(monkeypatch, by_prepare): + model_type = "nemotron_h_omni" + if by_prepare: + # A plan only the spec's prepare supplies is refused the same way, before prepare runs. + model_type = "_probe" + spec = ModelSpec( + model_type=model_type, + ep_spec=EPSpec( + plan_config_attr=("text_config", "llm_config"), + prepare=lambda *args, **kwargs: pytest.fail("ran prepare"), + prepared_ep_plan={"layers.*.experts": "grouped_gemm"}, + ), + ) + monkeypatch.setattr( + "modelopt.torch.models.get_spec", lambda mt: spec if mt == model_type else None + ) + with pytest.raises(ValueError, match=r"config\.text_config or config\.llm_config"): + apply_distributed_support(model_type, cfg()) + + +@needs_ep_transformers +@pytest.mark.parametrize( + "supplied", + [ + {"base_model_ep_plan": {"layers.*.experts": "grouped_gemm"}}, + {"prepare": lambda *args, **kwargs: pytest.fail("ran prepare")}, + ], + ids=["by-the-spec", "by-prepare"], +) +def test_a_plan_attached_where_the_load_does_not_read_one_is_an_error(monkeypatch, supplied): + """The load reads a VLM's plan off text_config or llm_config: attached anywhere else, it would + be read as no plan at all, and the MoE planned as a dense model. Refused whichever supplies the + plan, the spec or its prepare, before either touches the config.""" + spec = ModelSpec( + model_type="_probe", ep_spec=EPSpec(plan_config_attr="language_config", **supplied) + ) + monkeypatch.setattr( + "modelopt.torch.models.get_spec", lambda mt: spec if mt == "_probe" else None + ) + config = cfg(language_config=cfg()) + + with pytest.raises(ValueError, match="language_config, where the expert-parallel load"): + apply_distributed_support("_probe", config) + assert vars(config.language_config) == {} + + +@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. The refusal names a single process first: FSDP2 loads only an + unquantized causal LM.""" + config = cfg() + + with pytest.raises(NotImplementedError, match=r"yet: .* in a single process, or FSDP2-sharded"): + 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) == {} + + +@needs_ep_transformers +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, + } + + +def test_the_plan_a_spec_would_supply(monkeypatch): + """What the plan decision reads at an EP degree of 1, to refuse an MoE left whole on every rank: + 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. A model below its spec's + transformers floor keeps its plan: its experts are no less routed, and whole on every rank they + are refused like any other's.""" + 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"}) + ), + "below_its_floor": SimpleNamespace( + ep_spec=EPSpec(base_model_ep_plan=prepared, requires_transformers="999") + ), + } + monkeypatch.setattr(models, "get_spec", specs.get) + assert moe._spec_ep_plan("preparing") == prepared + assert moe._spec_ep_plan("refusing") == {} + assert moe._spec_refuses_ep("below_its_floor") + assert moe._spec_ep_plan("below_its_floor") == prepared From 35012df0b4b05b45ec07bef2d3afeedd994a6f8f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:26:05 +0000 Subject: [PATCH 2/6] Check every registered plan's styles, containers and config placement The well-formedness test only checked for non-empty strings: it now checks every entry's style against the ones the expert-parallel load implements, that the plan names a routed-experts container, and does so for the plan a spec's prepare attaches too. A new registry test checks that every spec puts its plan on a config the load reads. check_distributed_support drops its local import of parse, which the module imports at the top, and _ep_spec's local import says why it is local. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 4 +-- tests/unit/torch/models/test_ep_specs.py | 37 +++++++++++++++++++----- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index a12bb87dee9..0bc3ab8317c 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -149,6 +149,8 @@ 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``.""" + # Imported here: importing modelopt.torch.models registers every model spec, which importing + # this module should not do. from modelopt.torch.models import get_spec return getattr(get_spec(model_type), "ep_spec", None) @@ -177,8 +179,6 @@ def check_distributed_support(model_type: str) -> None: f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " f"{without_ep}." ) - from packaging.version import parse - # A model whose spec states no release, or that has no spec, still needs the load's floor: auto # must pick no EP degree below it, rather than reach the load and fail there. stated = ep.requires_transformers if ep is not None else None diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 74dbcade82d..e450e8a8bc2 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -50,16 +50,29 @@ def cfg(**kw): # --------------------------------------------------------------------------- the registered specs +# The styles the expert-parallel load implements for a plan's entries. +_KNOWN_STYLES = frozenset(moe._EXPERTS_STYLES) | {"grouped_gemm"} + + @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}" + """Every entry uses a style the expert-parallel load implements, and the plan names a + routed-experts container it can count and shard: a malformed plan fails silently at load.""" + ep = get_spec(model_type).ep_spec + for plan in (ep.base_model_ep_plan, ep.prepared_ep_plan): + if plan is None: + continue + 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 style in _KNOWN_STYLES, ( + f"{model_type}: {glob!r} uses style {style!r}, which the expert-parallel load does " + f"not implement (known: {sorted(_KNOWN_STYLES)})" + ) + assert moe._routed_expert_keys(plan), ( + f"{model_type}'s plan names no routed-experts container, so the load can neither count " + "nor shard its experts" + ) @pytest.mark.parametrize("model_type", EP_MODEL_TYPES) @@ -71,6 +84,14 @@ def test_a_spec_declares_a_plan_a_floor_a_prepare_or_a_refusal(model_type): ) +@pytest.mark.parametrize("model_type", EP_MODEL_TYPES) +def test_a_plan_goes_where_the_load_reads_one(model_type): + """A spelling the load does not read is read as no plan at all, and the MoE planned as dense.""" + attrs = get_spec(model_type).ep_spec.plan_config_attr or () + attrs = (attrs,) if isinstance(attrs, str) else attrs + assert not set(attrs) - set(moe._SUB_CONFIG_ATTRS), f"{model_type}: {attrs}" + + def test_the_omni_vlm_reuses_the_text_models_plan(): """Two spellings of one architecture must not drift into two different plans, nor share one dict, which a change through either spec would carry into the other.""" From 71fea4c4d4b466df551edf33339ef39e4eb024fa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:57:17 +0000 Subject: [PATCH 3/6] Run the spec registry checks without transformers again The applier's tests skipped the whole module when transformers was missing, which took the registry and plan checks, which need none of it, out of the partial-install unit run. Skip only the tests that read the installed transformers. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_ep_specs.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index e450e8a8bc2..be268ed7542 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -18,9 +18,6 @@ from types import SimpleNamespace import pytest - -transformers = pytest.importorskip("transformers") - from packaging.version import Version from modelopt.torch import models @@ -33,12 +30,18 @@ check_distributed_support, ) +try: + import transformers +except ImportError: # the optional [hf] extra: only the checks of its release need it + transformers = None + EP_MODEL_TYPES = [s.model_type for s in get_specs() if s.ep_spec is not None] +needs_transformers = pytest.mark.skipif(transformers is None, reason="needs transformers") # Applying a spec enforces the transformers floor the expert-parallel load needs. needs_ep_transformers = pytest.mark.skipif( - Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS), + transformers is None or Version(transformers.__version__) < Version(_EP_MIN_TRANSFORMERS), reason=f"expert parallelism needs transformers >= {_EP_MIN_TRANSFORMERS}", ) @@ -135,6 +138,7 @@ def test_the_models_own_plan_always_wins(model_type, config): assert not hasattr(getattr(config, "text_config", cfg()), "base_model_ep_plan") +@needs_transformers def test_a_model_without_a_spec_still_needs_the_loads_transformers(monkeypatch): """A plan that ships with transformers needs no spec, but its load needs the release every expert-parallel load does: auto picks no EP degree below it, rather than reach the load.""" @@ -284,6 +288,7 @@ def prepare(hf_config, **kwargs): } +@needs_transformers def test_the_plan_a_spec_would_supply(monkeypatch): """What the plan decision reads at an EP degree of 1, to refuse an MoE left whole on every rank: the spec's plan, or the plan it states its ``prepare`` attaches, read without running From bc362206803e59ac1502c157f2e479f6608917a6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 01:57:20 +0000 Subject: [PATCH 4/6] Say an MTP head the FSDP2 load cannot build rules it out too The way around expert parallelism offered FSDP2 for any unquantized causal LM, but --use_fsdp2 refuses a Nemotron-H checkpoint carrying MTP weights. Name that as well. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 0bc3ab8317c..cc53e203588 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -168,11 +168,13 @@ def check_distributed_support(model_type: str) -> None: """ ep = _ep_spec(model_type) # The FSDP2 load builds a causal LM and reads its tensors as stored, so a quantized or - # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process. + # multimodal checkpoint (GPT-OSS's MXFP4, Llama-4's) needs a single process, as does one with + # an MTP head the load cannot build (--use_fsdp2 refuses Nemotron-H's). without_ep = ( "without expert parallelism: in a single process, or FSDP2-sharded (--use_fsdp2 in " - "examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world) if the checkpoint is an " - "unquantized causal LM" + "examples/hf_ptq/hf_ptq.py, or MODELOPT_FSDP_SIZE=world) if the FSDP2 load supports the " + "checkpoint, which takes an unquantized causal LM with no MTP head it cannot build (such " + "as Nemotron-H's)" ) if ep is not None and ep.unsupported: raise NotImplementedError( From 3a86e4adc352a223e9ca85cb432cc85244830186 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:02:06 +0000 Subject: [PATCH 5/6] Check that a config prepare returns in place of its own keeps the plan The test of the spec's prepare hook had it return a sentinel, so it held only that apply_distributed_support returns what prepare returned. A replacement config must also carry over the plan the spec attached, as EPSpec.prepare's contract asks: the stand-in now returns one that does, and the test checks the returned config is that one and holds the plan. Signed-off-by: Shengliang Xu --- tests/unit/torch/models/test_ep_specs.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index be268ed7542..0ca7d90d673 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -261,12 +261,16 @@ def test_a_refusal_comes_before_the_rest_of_the_section(monkeypatch): @needs_ep_transformers 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.""" + """A prepare hook may inspect or adjust the plan, so it must not run before there is one. A + config it returns in place of the one it was handed is what the load gets, the plan carried + over, as the spec's contract asks.""" seen = {} + replacement = cfg() def prepare(hf_config, **kwargs): seen.update(kwargs, config=hf_config, plan_present=hasattr(hf_config, "base_model_ep_plan")) - return "replaced" + replacement.base_model_ep_plan = hf_config.base_model_ep_plan + return replacement spec = ModelSpec( model_type="_probe", @@ -279,7 +283,8 @@ def prepare(hf_config, **kwargs): returned = apply_distributed_support("_probe", config, ckpt_path="/c", trust_remote_code=True) - assert returned == "replaced" + assert returned is replacement + assert returned.base_model_ep_plan == spec.ep_spec.base_model_ep_plan assert seen == { "config": config, "ckpt_path": "/c", From 2b3dfc668dc9059d173610b1c4c4831015f1b97f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:27:13 +0000 Subject: [PATCH 6/6] Honor a config's own plan before requiring the spec's nested config apply_distributed_support required the nested config a spec names for its plan before it looked for a plan the config already carries, so a config with a plan of its own at its root, which wins and which the expert-parallel load reads there, was refused for lacking a nested config nothing needed. A plan the config carries is now found first; with none, a config without the nested config is still refused, naming both. Signed-off-by: Shengliang Xu --- modelopt/torch/utils/moe.py | 11 +++++++---- tests/unit/torch/models/test_ep_specs.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index cc53e203588..314cf8810fe 100644 --- a/modelopt/torch/utils/moe.py +++ b/modelopt/torch/utils/moe.py @@ -233,7 +233,8 @@ def apply_distributed_support( return hf_config # The spellings of the config the plan goes on; the first this config has takes it. Checked for - # every spec, before anything is attached or prepared, whatever supplies the plan. + # every spec, before anything is attached or prepared, whatever supplies the plan, but for a + # config carrying a plan of its own, which wins wherever it sits and needs no nested config. attrs = ep.plan_config_attr or () if isinstance(attrs, str): attrs = (attrs,) @@ -244,17 +245,19 @@ def apply_distributed_support( f"config.{' or config.'.join(unread)}, where the expert-parallel load does not read " f"one; it reads config.{' and config.'.join(_SUB_CONFIG_ATTRS)}." ) + existing = getattr(_ep_config(hf_config), "base_model_ep_plan", None) target = hf_config - if attrs: + if attrs and not existing: target = next( (sub for attr in attrs if (sub := getattr(hf_config, attr, None)) is not None), None ) if target is None: raise ValueError( f"{model_type} expects its expert-parallel plan on " - f"config.{' or config.'.join(attrs)}, which this config does not have." + f"config.{' or config.'.join(attrs)}, which this config does not have, nor a plan " + "of its own." ) - if ep.base_model_ep_plan and not getattr(_ep_config(hf_config), "base_model_ep_plan", None): + if ep.base_model_ep_plan and not existing: target.base_model_ep_plan = dict(ep.base_model_ep_plan) if ep.prepare is not None: diff --git a/tests/unit/torch/models/test_ep_specs.py b/tests/unit/torch/models/test_ep_specs.py index 0ca7d90d673..7a4bd9ef182 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -199,6 +199,18 @@ def test_a_missing_nested_config_is_an_error_not_a_silent_skip(monkeypatch, by_p apply_distributed_support(model_type, cfg()) +@needs_ep_transformers +def test_a_plan_at_the_root_needs_no_nested_config(): + """A config carrying a plan of its own at its root keeps it, as an existing plan does wherever it + sits, and the expert-parallel load reads it there: the nested config the spec's own plan would + go on is not required.""" + plan = {"layers.*.experts": "grouped_gemm"} + config = cfg(base_model_ep_plan=dict(plan)) + + assert apply_distributed_support("nemotron_h_omni", config) is config + assert config.base_model_ep_plan == plan + + @needs_ep_transformers @pytest.mark.parametrize( "supplied",