diff --git a/modelopt/torch/utils/moe.py b/modelopt/torch/utils/moe.py index 2aff08e6891..314cf8810fe 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,140 @@ 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``.""" + # 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) + + +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, 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 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( + f"{model_type} cannot be quantized expert-parallel yet: {ep.unsupported} Quantize it " + f"{without_ep}." + ) + # 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, 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,) + 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)}." + ) + existing = getattr(_ep_config(hf_config), "base_model_ep_plan", None) + target = hf_config + 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, nor a plan " + "of its own." + ) + 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: + 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..7a4bd9ef182 100644 --- a/tests/unit/torch/models/test_ep_specs.py +++ b/tests/unit/torch/models/test_ep_specs.py @@ -13,28 +13,69 @@ # 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 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, +) -from modelopt.torch.models import get_spec, get_specs +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( + transformers is None or 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 +# 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) @@ -46,6 +87,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.""" @@ -59,3 +108,230 @@ 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") + + +@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.""" + 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 +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", + [ + {"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. 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")) + replacement.base_model_ep_plan = hf_config.base_model_ep_plan + return replacement + + 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 is replacement + assert returned.base_model_ep_plan == spec.ep_spec.base_model_ep_plan + assert seen == { + "config": config, + "ckpt_path": "/c", + "trust_remote_code": True, + "plan_present": True, + } + + +@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 + ``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