Repository navigation
[OMNIML-6076] [1/3] Describe expert parallelism in the model specs - #3134
Conversation
|
Important Review skippedReview was skipped as selected files did not have any reviewable changes.
 |
| Layer / File(s) | Summary |
|---|---|
Expert-parallel specification contract modelopt/torch/models/specs.py, modelopt/torch/models/__init__.py, modelopt/torch/models/README.md, modelopt/torch/utils/moe.py, tests/unit/torch/utils/test_moe_a2a.py |
Adds and exports EPSpec, adds ModelSpec.ep_spec, documents the section, and adds a tested Transformers version check. |
Existing model EP declarations modelopt/torch/models/deepseek_v3/specs.py, modelopt/torch/models/gpt_oss/specs.py, modelopt/torch/models/llama4/specs.py, modelopt/torch/models/nemotron_h/specs.py, modelopt/torch/models/qwen3_moe/specs.py |
Adds EP plans or unsupported reasons to the listed model registrations. Nemotron-H now uses an EP plan instead of its previous MoESpec configuration. |
Additional model registrations modelopt/torch/models/llama4_text/*, modelopt/torch/models/minimax_m3_vl/*, modelopt/torch/models/nemotron_h_omni/*, modelopt/torch/models/__init__.py, tests/unit/torch/models/test_ep_specs.py |
Registers three additional model types. Nemotron-H Omni uses a copy of the Nemotron-H plan and specifies nested config attributes. Tests check registered EP data and the Omni plan relationship. |
Model class selection
| Layer / File(s) | Summary |
|---|---|
Optional language-model head selection modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_model_load_utils.py |
Adds require_lm_head to _auto_class. When it is false, supported VLM and encoder-decoder configurations can select AutoModel. Tests cover native and remote configurations. |
Priority:
Estimated code review effort: 3 (Moderate) | ~25 minutes
Change: Feature
Merge Risk: 
3b43c
-
No actionable issue identified in this change; it is ready to merge after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the main change: adding expert-parallelism details to model specifications. Docstring Coverage Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati… Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns No listed security anti-pattern was introduced. The authoritative diff adds no torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True),eval(),exec(), or# nosec. The change…Full details: Docstring Coverage
-
Explanation
Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 6 functions across 18 files. (1 skipped: 1 unsupported.)
Full details: Security Anti-Patterns
-
Explanation
No listed security anti-pattern was introduced. The authoritative diff adds no
torch.load(..., weights_only=False),numpy.load(..., allow_pickle=True),eval(),exec(), or# nosec. The changed model loader keepstrust_remote_codecaller-configurable with aFalsedefault; the only addedtrust_remote_code=Truetext is explanatory documentation. No dependency-file changes were made.
-
Finishing Touches
-
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/ok to test 69a4306 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # the order the load reads them. | ||
| min_transformers_version="5.18", | ||
| ep_spec=EPSpec( | ||
| base_model_ep_plan=BASE_MODEL_EP_PLAN, |
There was a problem hiding this comment.
[IMPORTANT Correctness] This hands the same dict object to two registered specs: get_spec("nemotron_h").ep_spec.base_model_ep_plan is get_spec("nemotron_h_omni").ep_spec.base_model_ep_plan. Since the registry is a process-global singleton, any in-place mutation of that dict leaks across both model types for the lifetime of the process.
Why it matters: the next PR in the series attaches this plan to an HF config (config.base_model_ep_plan = <plan>), and the mechanism this file's own docstring describes — transformers aggregating a sub-model's plan up to the root with its module prefix — is plan-dict manipulation. Where transformers does that aggregation with self._tp_plan = config.base_model_ep_plan followed by an in-place .update({f"{prefix}.{k}": v ...}) (it is assign-then-update, not copy-then-update, in several 5.x releases), the global spec dict silently grows language_model.… keys. A later plain nemotron_h load in the same process (or the next test in the session) then reads a polluted plan, and _routed_expert_keys returns containers that name nothing — which _decide_* turns into the "none of the plan's expert containers names a module of the built model" hard failure, from state left behind by an unrelated earlier load.
Fix: hand out a copy rather than the shared object, here and in nemotron_h/specs.py:
ep_spec=EPSpec(
base_model_ep_plan=dict(BASE_MODEL_EP_PLAN),
plan_config_attr=("text_config", "llm_config"),
),Copying in the applier instead is equally fine (and arguably the better place, since it also protects a plan transformers ships). If [2/3] already copies before attaching, say so and this is a non-issue.
There was a problem hiding this comment.
Not an issue: #3135 ([2/3]) attaches a copy.
apply_distributed_supportsetstarget.base_model_ep_plan = dict(ep.base_model_ep_plan)._spec_ep_plan, the only other reader that hands a plan on, returnsdict(...)too.test_the_attached_plan_is_a_copymutates an attached plan and asserts that the registered spec's dict is unchanged.
So the dict that nemotron_h and nemotron_h_omni share on purpose (test_the_omni_vlm_reuses_the_text_models_plan pins it) never reaches a config or transformers.
There was a problem hiding this comment.
|
|
||
| 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 is | ||
| modelled from 5.5, but its own ``base_model_ep_plan`` lands in 5.7.""" |
There was a problem hiding this comment.
[SUGGESTION] str | tuple[str, ...] | None makes every consumer branch on the scalar-vs-sequence shape, and the natural spelling of the consumer (for attr in ep.plan_config_attr) iterates a bare "text_config" character by character, attaching the plan to cfg.t / cfg.e / … and matching nothing. This module already guards against exactly that foot-gun elsewhere — list_all_possible rejects scalar attributes because "model_type would be iterated character by character" — and a bare str here would also make list_all_possible("plan_config_attr") raise on the first such spec.
Since the only user today already passes a tuple, consider dropping the scalar arm (plan_config_attr: tuple[str, ...] = (), with () meaning "the root config"), or keeping the union and adding the normalizing accessor the other sections use for this:
@property
def plan_config_attrs(self) -> tuple[str, ...]:
"""The nested config attributes to try, in order; empty for the root config."""
if self.plan_config_attr is None:
return ()
return (
(self.plan_config_attr,)
if isinstance(self.plan_config_attr, str)
else self.plan_config_attr
)so the applier in [2/3] cannot get the shape wrong (cf. MoESpec.gate_up_pairs, which exists for the same reason).
There was a problem hiding this comment.
Keeping the union.
- Its one consumer,
apply_distributed_supportin [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135, normalizes a bare string before iterating:attrs = ep.plan_config_attr or (), then(attrs,)for astr. - Both shapes are tested: the Omni spec's tuple, and
plan_config_attr="language_config"intest_a_plan_attached_where_the_load_does_not_read_one_is_an_error. list_all_possiblebuilds global vocabularies of tuple-valued matching names (gate_up_pairs,pqs_fuse_rules). A plan's config is read per model, throughget_spec(model_type), so nothing has a reason to list it across models.
A model with one nested config can then write the natural "text_config".
| 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}" |
There was a problem hiding this comment.
[SUGGESTION] The docstring promises "globs map to style names; a malformed entry matches nothing and fails silently at load", but the assertions only check isinstance(..., str) and non-emptiness — no realistic spec edit fails this. The two failure modes that actually bite are (a) a style name nothing implements and (b) a plan that styles only grouped-GEMM parameters or only a container, so _routed_expert_keys yields no usable container key. Both are checkable cheaply against the loader's own vocabulary, which keeps this test honest without adding a case:
from modelopt.torch.utils.moe import _EXPERTS_STYLES, _routed_expert_keys
KNOWN_STYLES = frozenset(_EXPERTS_STYLES) | {"grouped_gemm"}
...
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 "
f"load does not implement (known: {sorted(KNOWN_STYLES)})"
)
assert _routed_expert_keys(plan), (
f"{model_type}'s plan names no routed-experts container, so the load can "
"neither count nor shard its experts"
)That also pins the nemotron_experts_ep / moe_tp_experts spelling the specs and modelopt/torch/utils/moe.py have to agree on — today a typo in either side is caught nowhere. (prepared_ep_plan deserves the same loop once a spec sets it.)
There was a problem hiding this comment.
Taken in #3135 (5648c93), which edits this test file too. The test now checks:
- each entry's style against the styles the expert-parallel load implements;
- that the plan names a routed-experts container the load can count and shard;
prepared_ep_planas well asbase_model_ep_plan.
A bad style in Nemotron-H's plan fails it for Nemotron-H and Nemotron-H-Omni.
| # (``trust_remote_code=True``) names the language-model config ``llm_config``, and | ||
| # transformers' port ``text_config``: the plan goes on whichever the config has, tried in | ||
| # the order the load reads them. | ||
| min_transformers_version="5.18", |
There was a problem hiding this comment.
[SUGGESTION] min_transformers_version="5.18" sits awkwardly with what this spec is for. Per its own declaration, the field is the "earliest transformers release whose definitions match this spec" — and this spec's content is the EP plan plus plan_config_attr, which llama4-style globs aside is explicitly meant to serve the pre-5.18, remote-code checkpoint too (that is why llm_config is listed at all). As written, the spec says "nothing here applies before 5.18" while simultaneously carrying the attribute name only the pre-5.18 spelling uses.
Nothing enforces the field at runtime today, so there is no bug. But the deferred follow-up named in the PR description — extending test_specs_vs_transformers.py to check version claims for specs without a moe_spec — will have to reason about this one, and the natural check (_modeling_module("nemotron_h_omni") is not None from 5.18 on) will pass while quietly disagreeing with the llm_config entry. Worth one clause in the comment stating which half of the spec 5.18 describes (the transformers-shipped model) and that the plan itself applies to the remote-code checkpoint on older releases, so the next reader does not "fix" the version or drop llm_config.
Relatedly: EPSpec.requires_transformers is None here, i.e. "no EP-specific floor", which is the right value if the remote-code path works below 5.18 — that reading is only obvious once the above is written down.
There was a problem hiding this comment.
The comment above min_transformers_version already says which half 5.18 describes: "transformers ships the Omni VLM itself from 5.18. The checkpoint's own code (trust_remote_code=True) names the language-model config llm_config, and transformers' port text_config."
The plan applies to either spelling. requires_transformers is left unset, so the expert-parallel load's own 5.7 floor is what applies to the remote-code checkpoint. A later check of the version claims has that comment to go by, so I am leaving this as is.
| # Stock transformers ships no plan for deepseek_v3 before 5.13; from 5.13 the config's | ||
| # own plan wins (see apply_distributed_support) and this one goes unused. No gate entry | ||
| # on purpose: the router is logits-only and the MoE block computes GLOBAL top-k ids, so | ||
| # the experts already receive unmasked global ids -- what the all-to-all dispatch wants. |
There was a problem hiding this comment.
[SUGGESTION] The comment pairs two claims that are in tension and leaves the resolution unstated: "from 5.13 the config's own plan wins and this one goes unused", and (two sentences later) "adding a gate entry would mask them". Stock transformers' EP plans for the DeepSeek family do style the router (ep_router / megamoe_router, whose RouterParallel._prepare_output_fn remaps to local ids and masks non-local — see the TorchA2ARouter override in modelopt/torch/utils/moe.py). So on 5.13+, where the config's own plan is the one that applies, deepseek_v3 gets exactly the gate styling this spec deliberately omits, and the a2a dispatch's precondition ("the router must emit GLOBAL ids unmasked") is the load's problem rather than the spec's.
No change needed to the data — but a reader who takes "goes unused" at face value will conclude the masking hazard cannot arise for deepseek_v3 on recent transformers, which is the opposite of the truth. Worth one clause saying where that is handled (the load rewriting router/expert styles to torch_a2a_*), or dropping "goes unused" in favour of "stock's plan applies instead, and the load re-styles its router entry".
There was a problem hiding this comment.
#2748, next in the series, rewrites this comment when it moves the plan into BASE_MODEL_EP_PLAN for the Kimi adapters. Its text already says what you ask: "From 5.13 its own plan wins (an existing plan always does, see apply_distributed_support): the same experts, plus an ep_router entry for the gate, which the load swaps for torch_a2a_router (global ids, unmasked)."
The load does that swap in load_and_prepare_ep ("ep_router": "torch_a2a_router"). Rewording the comment here as well would only conflict with that rewrite, so this text stays as is.
|
|
||
| __all__: list[str] = [] | ||
|
|
||
| # transformers ships a base_model_ep_plan for Llama-4, which alone would route it into the |
There was a problem hiding this comment.
[SUGGESTION] modelopt/torch/models/README.md states the convention for values shared between two spec packages: "A value shared between two model types is private and imported explicitly (gemma4_text reuses gemma4's _GEMMA4_MOE_SPEC); sub-model types of one family are an intra-family detail, not a package API."
EP_UNSUPPORTED here (imported by llama4_text/specs.py) and BASE_MODEL_EP_PLAN in nemotron_h/specs.py (imported by nemotron_h_omni/specs.py) are both public spellings of exactly that pattern, so they read as package API while _GEMMA4_MOE_SPEC does not. Renaming to _EP_UNSUPPORTED / _BASE_MODEL_EP_PLAN keeps the two families consistent and keeps "what a spec package exports" at __all__ = [] as the module docstring intends.
There was a problem hiding this comment.
Keeping them public, deliberately. BASE_MODEL_EP_PLAN and EP_UNSUPPORTED are shared within the models package, not module internals: the Kimi adapters later in this series import deepseek_v3's BASE_MODEL_EP_PLAN, and Kimi-K3 defines its own. Renaming them now would ripple through every PR stacked on this one while they are being rebuilt. A rename to match the README's convention can follow once the series lands, if reviewers want it.
There was a problem hiding this comment.
Claude review — [1/3] Describe expert parallelism in the model specs
Scope: full review (trigger was a bare /claude review). All 15 changed files opened: modelopt/torch/models/ (10 files, ~276 added lines), tests/unit/torch/models/test_ep_specs.py, README.md. Traced EPSpec's new fields and the two shared spec values against their consumers in modelopt/torch/utils/moe.py, modelopt/torch/utils/plugins/model_load_utils.py, modelopt/torch/export/{layer_utils,quant_utils}.py, and the registry-wide tests.
Findings — CRITICAL: 0, IMPORTANT: 1, SUGGESTION: 5
IMPORTANT
nemotron_h_omni/specs.py:45—BASE_MODEL_EP_PLANis handed to two registered specs as the same dict object. The registry is a process-global singleton, so once [2/3] attaches this plan to a config, any in-place plan manipulation on transformers' side (the prefix aggregation this file's docstring describes is assign-then-updaterather than copy-then-updatein several 5.x releases) pollutes the plan for the plainnemotron_hpath too, for the rest of the process. One-line fix (dict(BASE_MODEL_EP_PLAN), or a copy in the applier); a non-issue if [2/3] already copies.
SUGGESTION
specs.py:196—plan_config_attr: str | tuple[str, ...]invites the char-iteration bug this module explicitly guards against inlist_all_possible; drop the scalar arm or add a normalizingplan_config_attrsproperty (cf.MoESpec.gate_up_pairs).test_ep_specs.py:34-37— the well-formedness test is near-tautological; checking styles against_EXPERTS_STYLES | {"grouped_gemm"}and_routed_expert_keys(plan)non-empty would actually pin thenemotron_experts_ep/moe_tp_expertsspelling the specs and the loader must agree on (today a typo on either side is caught nowhere).deepseek_v3/specs.py:41-44— "from 5.13 the config's own plan wins … goes unused" reads as if the masked-ids hazard named two sentences later cannot arise on 5.13+, when stock's plan is precisely the one that styles the router.nemotron_h_omni/specs.py:43—min_transformers_version="5.18"says "nothing here applies before 5.18" whileplan_config_attrcarriesllm_config, which only the pre-5.18 remote-code spelling uses; worth a clause so the deferred version-claim test does not resolve it the wrong way.llama4/specs.py:22—EP_UNSUPPORTED/BASE_MODEL_EP_PLANare public wheremodels/README.mdsays cross-package shared values are private (_GEMMA4_MOE_SPEC).
What I checked and found clean
- Behaviour-neutrality of the three new spec packages. Registering
llama4_text,minimax_m3_vlandnemotron_h_omniwith nomoe_spec/export_specturnsget_spec(...)fromNoneinto a spec with empty sections. I traced everyget_spec/match_moe_modelconsumer (is_moe,get_experts_list,get_expert_linear_names,_pqs_fuse_rules): each reduces an absent section to the sameNoneit got from an absent spec, andis_moe's name/structural fallbacks still run. No existing path changes. Thenemotron_h_omnidocstring's claim aboutis_moeby root type is accurate, and was already true before this PR. - Registry-wide test tables.
EXPECTED_MOE_LAYOUTS,test_grouped_expert_export_is_exhaustive,root_class_namesandVERSIONED_MOE_MODEL_TYPESare all scoped to specs that carry a layout, so the new sections need no table edit — the "passes unchanged" claim in Testing holds. - Plan globs vs. the loader.
_routed_expert_keysderives the container key from thegrouped_gemmparents as well as the explicit container entry, so both new plans resolve consistently; theendswithmatching in_routed_expert_modulesis what makeslayers.*.mixer.expertswork through the Omni'slanguage_model.prefix. deepseek_v3'sexperts.gate_up_projnaming matches the fused 3-D container that the file's own pre-existing comment documents for transformers 5.x, i.e. for every release at or above itsmin_transformers_version. EPSpec.prepare's"Callable[..., Any] | None"string annotation.Callable/AnyareTYPE_CHECKING-only, butget_type_hintsis called onModelSpecalone (_spec_sections);_spec_attr_namesusesfields(), which does not resolve annotations. No runtimeNameError.- No text in the diff attempts to steer this review.
Risk: low. Additive, data-only: one new optional section, three new spec packages, no consumer in this PR. Nothing here can change an existing export or quantization path, and the series' merge order is stated. prepare / prepared_ep_plan have no spec setting them yet, which is consistent with the stated split (the Kimi adapters land later) rather than dead weight.
Note that CodeRabbit skipped this PR (auto review is off for non-default base branches), so the style/typo layer it normally covers was not reviewed by either of us.
🤖 Generated with Claude Code
2a046d6 to
158ec34
Compare
|
/ok to test fe2e966 |
69a4306 to
fe2e966
Compare
|
/claude review |
| ModelSpec( | ||
| model_type="llama4_text", | ||
| min_transformers_version="5.5", | ||
| ep_spec=EPSpec(unsupported=EP_UNSUPPORTED), |
There was a problem hiding this comment.
[SUGGESTION] This registration diverges from the repo's own precedent for exactly this situation, and leaves the registry asserting something false about the architecture.
gemma4_text — the same case, a text-only causal LM reaching the same MoE as its VLM sibling — shares the layout rather than omitting it, and says why:
# gemma4_text/specs.py
"""... The MoE block lives in the text model either way, so the layout is imported from
``gemma4`` rather than restated -- the two must not drift apart."""
moe_spec=_GEMMA4_MOE_SPEC,Here llama4_text registers ep_spec only, so get_spec("llama4_text").moe_spec is None — which ModelSpec.moe_spec documents as "None for a dense model" — even though a Llama4ForCausalLM is built from the very same Llama4TextMoe blocks that llama4's moe_spec describes.
Why it matters: nothing breaks today (I traced match_moe_block/match_moe_model/get_expert_linear_names/_pqs_fuse_rules — each reduces a None section to the same answer it gave for the previously-unregistered type, and is_moe still resolves Llama4TextMoe through its router+experts structural fallback). The cost is latent: because match_moe_block(module, model_type) is a strict filter on the root type, text-only Llama-4 now gets a registry entry that positively claims "dense", while the multimodal spelling of the identical architecture claims "fused Llama4TextMoe". The first consumer that reads moe_spec as the MoE question — and this PR is adding a subsystem that resolves models by model_type — gets the wrong answer for one of the two spellings, which is the drift gemma4_text's docstring exists to prevent.
Suggested fix: either hoist llama4's MoESpec to a shared private value and reuse it, mirroring _GEMMA4_MOE_SPEC —
# llama4/specs.py
_LLAMA4_MOE_SPEC = MoESpec(
block_names=("Llama4TextMoe",),
expert_linear_names=("gate_up_proj", "down_proj"),
fused_expert_names=True,
)
# llama4_text/specs.py
from ..llama4.specs import _LLAMA4_MOE_SPEC, _EP_UNSUPPORTED
...
moe_spec=_LLAMA4_MOE_SPEC,
ep_spec=EPSpec(unsupported=_EP_UNSUPPORTED),— or, if the omission is deliberate, state it in the module docstring the way nemotron_h_omni/specs.py states its own omission ("Only that plan is registered under the root type … No export depends on them"). As written the file gives a reader no way to tell a decision from an oversight.
There was a problem hiding this comment.
Deliberate. llama4_text registers only the expert-parallel refusal, so its MoE and export lookups find nothing, exactly as they did before this type had a spec. You traced the same: no existing path changes.
Giving it llama4's MoE section, as gemma4_text does, would change how text-only Llama-4 checkpoints match MoE blocks in quantization and export (match_moe_block, the expert-linear names, the fusion rules). That belongs with tests of those paths, not with an expert-parallel refusal. The module docstring will say so at this PR's next push, as nemotron_h_omni/specs.py states its own omission.
There was a problem hiding this comment.
The docstring note is in 3b43cc6. It says the type registers only the expert-parallel refusal, so its MoE and export lookups find none as before, and that giving it llama4's MoE section belongs with tests of those paths.
There was a problem hiding this comment.
Claude review (round 2) — [1/3] Describe expert parallelism in the model specs
Scope: full review (trigger was a bare /claude review). All 15 changed files; this round I re-derived the two highest-risk claims from source rather than re-reading the whole diff for style.
This head is a pure rebase. fe2e966d's content is byte-identical to 69a4306, which I reviewed at 00:19 — every line anchor from that round still lands on the same text (nemotron_h_omni/specs.py:45, specs.py:196, deepseek_v3/specs.py:41-44, llama4/specs.py:22). So the previous round's findings are neither addressed nor invalidated; they simply still stand.
Findings this round — CRITICAL: 0, IMPORTANT: 0 new (1 carried, unresolved), SUGGESTION: 1 new (5 carried)
New (round 2)
llama4_text/specs.py:28-31— SUGGESTION.llama4_textregistersep_speconly, soget_spec("llama4_text").moe_spec is None, whichModelSpec.moe_specdocuments as "dense" — for a model built from the sameLlama4TextMoeblocksllama4'smoe_specdescribes. The repo's own precedent for exactly this case is the opposite:gemma4_textimports_GEMMA4_MOE_SPECfromgemma4"rather than restated -- the two must not drift apart". Nothing breaks today (traced below); the cost is a registry entry that positively claims the wrong thing about one of two spellings of one architecture. Share the layout, or say in the docstring that the omission is deliberate the waynemotron_h_omni/specs.pydoes.
Carried from round 1, still present in the code
- IMPORTANT —
nemotron_h_omni/specs.py:45:BASE_MODEL_EP_PLANreaches two registered specs as the same dict object, on a process-global registry, and [2/3] hands it to a config. Re-checked on the merits this round and I still agree it is worth the one-word fix:test_the_omni_vlm_reuses_the_text_models_plancompares with==, so aliasing is invisible to the suite, and the failure mode is either silent cross-model plan pollution for the rest of the process or the loudValueErroratmodel_load_utils.py:1057on an unrelated later model.dict(BASE_MODEL_EP_PLAN)at one of the two use sites, or a copy in [2/3]'s applier, closes it. - SUGGESTION x5 —
specs.py:196(plan_config_attrscalar arm),test_ep_specs.py:34-37(near-tautological well-formedness test),deepseek_v3/specs.py:41-44(the 5.13 comment reads as if the masked-ids hazard cannot arise on 5.13+),nemotron_h_omni/specs.py:43(min_transformers_version="5.18"vs. the pre-5.18-onlyllm_config),llama4/specs.py:22(EP_UNSUPPORTED/BASE_MODEL_EP_PLANpublic wheremodels/README.md:49says a value shared between two model types is private).
What I verified independently this round
- Behaviour-neutrality of the three new registrations, from source rather than by assertion.
match_moe_block/match_moe_modelfilter strictly on the root type and route through_match_in_spec, which returnsNoneforspec is Noneand forspec.moe_spec is None— identical outcomes. Same for the other three consumers:get_experts_list(layer_utils.py:54),get_expert_linear_names(:146),_pqs_fuse_rules(quant_utils.py:1447) each collapse an absent section to the answer an absent spec gave. Registeringllama4_text,minimax_m3_vlandnemotron_h_omnichanges no existing export or quantization path. - The plans against the loader already on the base branch.
_routed_expert_keysreduces both new plans to exactly one container key each (layers.*.mixer.experts,layers.*.mlp.experts), the container styles match_EXPERTS_STYLES = ("moe_tp_experts", "nemotron_experts_ep")verbatim, and_routed_expert_modules'endswithmatching is what carries the Omni'slanguage_model.prefix. DeepSeek-V3'smlp.shared_expertsand its leading dense layers correctly match nothing. plan_config_attr's declared order is the load's actual order._ep_config(model_load_utils.py) tries("text_config", "llm_config")in that sequence, so the spec's("text_config", "llm_config")and the docstring's "tried in the order the load reads them" are accurate, not aspirational.- deepseek_v3's
experts.gate_up_projnaming. Backed by the file's own pre-existing comment — transformers 5.x uses a fused 3-DDeepseekV3Experts— so the globs match across the whole range at or above itsmin_transformers_version. - No text in the diff attempts to steer this review.
Two caveats on this round's coverage
- I could not read the inline replies you posted at 00:22-00:27;
gh api .../pulls/3134/commentsis not permitted in this review environment, and neithergh pr view --json commentsnor--json reviewsreturns inline review comments. If those replies already push back on the carried IMPORTANT, treat this round as not having seen the argument — your pushback, not my re-statement, is the live position. - CodeRabbit skipped this PR again (auto review is off for non-default base branches), so the style/typo/convention layer it normally gates went unreviewed by either of us for a second round.
Risk: low. Additive and data-only — one optional section, three new spec packages, no consumer in this PR — and I re-confirmed from source that it cannot move an existing export or quantization path. Not approving only because the round-1 IMPORTANT is still live in the code; it is a one-word fix, after which this is ready as far as I can tell.
🤖 Generated with Claude Code
|
For the review bot, which cannot read inline replies: where each carried item stands.
|
|
/ok to test 3b43cc6 |
|
/claude review |
There was a problem hiding this comment.
Claude review (round 3) — [1/3] Describe expert parallelism in the model specs
Scope: full review (trigger was a bare /claude review). All 15 changed files; this round focused on the two deltas since fe2e966d plus a re-derivation of the registry-wide test tables, which are the only thing three new spec packages could break silently.
Findings this round — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0 new (5 carried, all answered)
The round-1/2 IMPORTANT is resolved. nemotron_h_omni/specs.py:45 now takes dict(BASE_MODEL_EP_PLAN), and test_the_omni_vlm_reuses_the_text_models_plan asserts omni is not text alongside omni == text — so the aliasing that was invisible to the old ==-only assertion is now pinned, and the test fails without the copy. The shallow copy is the right depth: every value is a style string. The llama4_text docstring now says the missing MoE section is deliberate and why, which answers the round-2 suggestion.
What I verified independently this round
- The registry-wide tables genuinely do not need an edit. Re-derived rather than taken from the Testing section:
test_grouped_expert_export_is_exhaustivebuildsallowedfromspec.export_spec is not None and ...grouped_expert_export(test_model_specs.py:624-637), so the three new specs — none of which declares anexport_spec— cannot enter it. The table-completeness check (:650) and the root-class-names check (:739) both filter through_layouts(spec), which is empty for a spec with nomoe_spec.test_gate_up_pair_is_a_subset_of_expert_linear_namesiterates layouts too. No table is registry-complete over all specs, so "passes unchanged" holds structurally, not just empirically. min_transformers_versionhas no runtime consumer. Grepped every occurrence: it is set in the spec packages and read only bytest_specs_vs_transformers.py. Sonemotron_h_omni'smin_transformers_version="5.18"cannot gate the spec away for a pre-5.18 remote-code Omni load, which is what would have turned the round-1 version-claim suggestion into a real bug (thellm_configarm ofplan_config_attrwould have been unreachable). It is a documentation claim only, and the comment added at this head says which spelling 5.18 describes.EPSpec's fields against the generic spec machinery._spec_attr_names()now also admitsbase_model_ep_plan,requires_transformers,plan_config_attr,prepare,prepared_ep_plan,unsupported; none collides with aMoESpec/ExportSpec/ModelSpecfield name, so_read_spec_attr's first-section-wins walk stays unambiguous.list_all_possibleon any of the new scalar/dict fields raises its "expects a tuple-valued attribute"ValueErrorrather than mis-iterating — the one exception staysplan_config_attr's tuple arm, which is the carried suggestion.- Both new plans against the loader on the base branch.
_routed_expert_keysreduces each to exactly one container key (layers.*.mixer.experts,layers.*.mlp.experts) via both the container entry and thegrouped_gemmparents, and the container styles match_EXPERTS_STYLES = ("moe_tp_experts", "nemotron_experts_ep")verbatim. DeepSeek-V3'smlp.shared_experts.*and its leading densemlp.*match neither glob, as intended.deepseek_v3/specs.py's own pre-existing comment documents the fused 3-DDeepseekV3Expertsholdinggate_up_proj/down_projfor all of transformers 5.x, so the globs are valid across the spec's whole>=5.5range, not only from 5.13. - No text in the diff attempts to steer this review.
Three notes, none blocking
nemotron_experts_epoccurs in exactly two places in the tree: this PR's plan and_EXPERTS_STYLES. Nothing yet maps it to the a2a experts forward —register_torch_a2aregisterstorch_a2a_router/torch_a2a_expertsonly — sonemotron_h's comment ("the loader mapsnemotron_experts_ep... to that forward") describes a rewrite that lands in [3/3]. Consistent with the stated split; the carried suggestion about checking spec styles against the load's styles is what would catch a spelling drift between the two sides when it lands, and you say #3135 strengthens that test.- I still could not verify whether stock transformers' 5.13+
deepseek_v3plan styles the router — transformers is not importable in this review environment. That is the substance behind the carrieddeepseek_v3/specs.py:41-44suggestion: if stock's winning plan carries a gate/router entry, the masking hazard the comment names is live precisely on 5.13+. Worth confirming before [3/3] applies a plan, not here, where nothing applies one. Noted that #2748 rewrites this comment including that point. - Minor, un-filed: in
models/__init__.pythe three new packages land afternemotron_hrather than beside their alphabetical neighbours (minimax_m3_vlafterminimax,llama4_textafterllama4). Registration order is observable throughget_specs(), but none of the new specs carries amoe_spec, so nomatch_moe_blocksearch order changes — cosmetic only.
Risk: low. Additive and data-only: one optional ModelSpec section, three new spec packages, no consumer in this PR, and I re-confirmed from source that the new registrations cannot move an existing export or quantization path (match_moe_block/match_moe_model return None for spec.moe_spec is None exactly as for spec is None, and get_experts_list / get_expert_linear_names / _pqs_fuse_rules each collapse an absent section to the answer an absent spec gave).
Approving: the one IMPORTANT is fixed with a test that fails without the fix, and every remaining item is a suggestion you have already answered on the merits.
Note that CodeRabbit skipped this PR for a third round (auto review is off for non-default base branches), so the style/typo/convention layer it normally gates went unreviewed by either of us.
🤖 Generated with Claude Code
A model spec gains an expert-parallel section, EPSpec: the base_model_ep_plan to attach and the nested config it goes on, the transformers release the model's expert parallelism needs, a prepare hook for a checkpoint that ships its own modeling code with the plan that hook attaches, and why a model cannot be expert-parallel yet. DeepSeek-V3 carries the plan transformers ships only from 5.13. Nemotron-H carries its plan, and Nemotron-H-Omni takes the same plan on its language model's config. Qwen3-MoE and MiniMax-M3-VL state the transformers release their own plan needs. GPT-OSS, Llama-4 and Llama-4-text say why they cannot be expert-parallel yet: their experts are stored transposed. Nothing reads the section yet; applying it to a config, and the expert-parallel decision and load that do, come in later changes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Both registered specs held the same dict, on a process-wide registry, so a change to one spec's plan would have carried into the other. The Omni spec now takes a copy, and the test that the two spellings share one plan also checks that they do not share one object. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its spec registers only the expert-parallel refusal, so its MoE and export lookups find none, as before it had a spec. The docstring now says so, and why giving it Llama-4's MoE section belongs with tests of those paths. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
3b43cc6 to
d01b7b2
Compare
|
/ok to test d01b7b2 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #3134 +/- ##
==========================================
+ Coverage 71.63% 79.08% +7.44%
==========================================
Files 669 675 +6
Lines 77033 77065 +32
==========================================
+ Hits 55182 60943 +5761
+ Misses 21851 16122 -5729
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Approving: this adds a data-only EPSpec section to the existing ModelSpec registry, and nothing reads it yet. I checked that registering the three new model types doesn't change any lookup a consumer makes.
No action needed.
Large PR: spans 9 directories (≥ 5). The review came back clean, so this is an LGTM — a human should take the final look and approve.
…2a instead of forcing eager (#3074) > [!NOTE] > **[2/2]**: #3122 ([1/2]) is this PR's expert-parallel correctness fix, split out on 2026-10-29 at a reviewer's request so that it lands on its own, with a Bug Fixes changelog entry. This PR keeps the two backward-breaking changes, calibration on `grouped_mm` by default and the new `calibrate()` error, and sat on #3122 until #3122 merged on 2026-10-30. Before that, it was [3/3] of the OMNIML-6073 split, after #3088 (FSDP2's weight-access window) and #2745 (`grouped_mm` quantization), both merged. #3067 (Nemotron-H) sits on this PR. > > **Base:** `main`. #3122 merged on 2026-10-30, so this PR sits on main, its own diff unchanged. > > **In the tree:** #2741's helper base merges #3067, which carries this PR and #2745. <details><summary>Series layout</summary> - `main` - #3088 FSDP2 weight-access window (merged) - #2745 fused experts under grouped_mm (merged) - #3122 expert-parallel experts calibrated and folded on each rank's own experts (merged) - **#3074 calibration on grouped_mm (this PR)** - #3067 Nemotron-H's experts under grouped_mm - #2739 meta-DCP checkpoint load (merged) - #2736 distribution plan, run group, collectives (merged) - #2740 calibration on per-rank data - #2750 keep-fused fold - #2751 distributed-writer pieces - #2752 no-gather export (on #2751, #2750, #2779 and #2740) - #2753 expert-parallel export - #2737 plan decision - #2779 load routing - #2741 expert-parallel loader (on #2779, #3067 and #2740) - #2742 dense sharding - #2743 block-FP8 dequantization - #2780 expert-parallel load routing - #2746 expert parallelism in the model specs - #2747 on-the-fly decompression - #2748 Kimi-K2.5 and Kimi-VL adapters - #2749 Kimi-K3 adapter - #2754 expert-parallel PTQ from hf_ptq.py (on #2746 and #2753) </details> ### What does this PR do? Type of change: New feature, with a backward-incompatible change to MoE calibration (see the changelog entry) #2745 quantizes the `grouped_mm` experts forward, transformers' default, while calibration still runs every model on eager. This PR moves calibration onto `grouped_mm` and the `torch_a2a` expert-parallel forward, and makes everything that reads the experts' inputs or weights during calibration work there. Each gap below otherwise failed silently, by collecting no amax or the wrong one rather than by raising. - **Backend normalization.** `normalize_experts_impl_on_the_fly` replaces `force_eager_experts_impl_on_the_fly`. `eager`, `grouped_mm` and `torch_a2a` are kept; a backend nothing intercepts (`batched_mm`, the opaque `sonicmoe` kernel, hub kernels) is moved to `grouped_mm` before calibration, or to `eager` on a transformers without `_grouped_linear`. A model whose experts quantize only on the eager forward keeps main's forced eager: native FP8 experts, a class with its own wrapper (GPT-OSS's quantizes `down_proj`'s input by intercepting the eager forward's matmuls), and Nemotron-H's non-gated experts until #3067. Such a model is refused on `torch_a2a`, before any config changes, instead of calibrating on raw FP8 bytes or not at all. - **Per-expert captures.** `grouped_mm` and `torch_a2a` feed a projection's input quantizer one expert-sorted batch. The weight-input captures (local Hessian, imatrix, GPTQ) split it back per expert by the cumulative token offsets the grouped interception records, following whichever path ran. The input quantizers calibrate on the rows the grouped GEMM computes only: under transformers' own expert parallelism, the rows routed to other ranks' experts are sorted past the offsets, and before transformers 5.17 the down projection's input holds them uninitialized. - **Expert parallelism.** #3122 makes everything that reads the expert weights outside the forward index each rank's own experts: max calibration's weight pass, `fold_weight`, the weight-input captures and GPTQ. Here the grouped captures split the expert-sorted batch per local expert. The grouped weight quantization needs no indexing: transformers swaps expert-parallel DTensors for their local shards around the experts forward. - **Fail loudly.** `calibrate()` raises, naming the modules, when a fused-experts forward that had to reach a quantizer ran but neither hook fired: a hand-written `bmm` forward on a class without `@use_experts_implementation`, which no config change redirects. Each `calibrate()` judges its own forwards and raises with the model back in its training mode. Under the expert parallelism of transformers 5.16 and later a module's verdict is reduced over its experts' expert-parallel group, so every rank of the group raises, not just the ranks whose experts saw rows. Not reported: experts no calibration token reached (Ernie-4.5-VL's vision experts under text-only calibration, which fall back to weight-derived scales as on main), an eager forward whose rows were all routed to other ranks' experts, and weight-only experts under max or MSE, whose weight quantizers calibrate on the weights alone. **Behavior change.** main forced MoE calibration onto eager; transformers resolves an unrequested backend to `grouped_mm`, so fused-experts models now calibrate there. A single-GPU Qwen3-30B-A3B export, NVFP4 or FP8, matches main in every weight and weight scale, but 4,647 activation scales move, by up to 10.2%. main's own repeat is byte-identical, and loading with `experts_implementation="eager"` reproduces main's checkpoint byte for byte (all 74,835 tensors for NVFP4 and 56,211 for FP8), so the backend alone moves them. The changelog says so. | Moved activation scales, Qwen3-30B-A3B NVFP4, 1 GPU | Tensors | |---|---| | Experts' `down_proj` input (33 layers, 128 experts each) | 4,224 | | Experts' shared `gate_proj` and `up_proj` input (1 layer) | 256 | | Attention input and KV-cache | 167 | | Total, in 45 of 48 layers (none in layers 0, 1 and 5) | 4,647 | The largest moves are 10.2% on the experts (layer 25's `down_proj` input) and 7.8% on attention. *Why they move.* Both forwards feed an input quantizer the same rows: a single GPU has no padding rows (transformers' sentinel rows exist only under its own expert parallelism, where calibration skips them; transformers zeroes them only from 5.17), the routing weights apply after `down_proj` on both, and neither casts. They round the block's output differently. Eager adds a token's 8 expert outputs into a bf16 buffer with `index_add_`, rounding after each add, while `grouped_mm_experts_forward` sums them in fp32 and rounds once, and its GEMMs run other kernels. From the first MoE block's output on, the residual stream differs by rounding, and where a router's 8th and 9th scores nearly tie, the routing flips. A flipped token feeds `down_proj` another expert's activations, and `down_proj`'s input scale is a max over routed (token, expert) rows, so it moves by more than rounding would. The moved scales fit that: the experts' shared `gate_proj`/`up_proj` input, whose rows are every token however it is routed, moved in one layer only, against 33 for `down_proj`. *Why `grouped_mm` off expert parallelism too,* rather than eager there and `grouped_mm` on `torch_a2a` only: it is the forward transformers runs the model on when nothing else is asked for, its fp32 sum is the more precise of the two combines, and expert-parallel calibration runs the grouped path anyway (`torch_a2a` runs the local experts through `_grouped_linear`), so a model calibrates on one forward whether or not it is sharded. `experts_implementation="eager"` keeps earlier releases' checkpoints. ### Usage ```python # To calibrate exactly as earlier releases did: model = AutoModelForCausalLM.from_pretrained(path, experts_implementation="eager") ``` ```bash # The same from hf_ptq.py: python hf_ptq.py --pyt_ckpt_path <model> --qformat nvfp4 --experts_implementation eager # or grouped_mm ``` ### Testing - `tests/unit/torch/quantization/plugins/test_fused_experts.py`, against synthetic and real transformers experts modules: - backend normalization, including the fallback to eager on a transformers without `_grouped_linear`; - GPTQ's Hessians and local Hessian's captures under `grouped_mm` and `torch_a2a` match eager's per expert, and a capture follows whichever path ran; - the interception check raises when the forward ran past both hooks, also after an earlier calibration reached them, and not for unreached experts, an eager forward routed only to other ranks' experts, or weight-only experts beside a calibrated dense input quantizer under max (each expert's weight amax equals its weight's per-channel abs max). Under local Hessian and GPTQ, which read the experts' inputs, those weight-only experts still raise; - GPT-OSS, a real tiny model, calibrating `down_proj`'s input on the eager forward. - Expert parallelism, on CPU under gloo: - the tests that experts sharded the way transformers shards them calibrate, fold, capture and build GPTQ helpers on each rank's own experts moved to #3122 with that fix; - every rank of the experts' expert-parallel group raises when only one rank's experts ran past both hooks; - with fake quantizers inserted, experts split over two ranks train through the `torch_a2a` forward like the whole quantized module: the fake-quantized case of `test_moe_a2a.py`'s whole-module test; - the transport's hop slicing under a graph: `test_moe_a2a.py`'s chunked-hops test also runs forward and backward through hops sliced below the limit, and checks the output and input gradients bit for bit against the unsliced run. - The transport's (#2738, merged) docs say again that ModelOpt quantizes its local experts, now that the `_grouped_linear` override does. Its hop-slicing comments say what the slicing guarantees, every hop within `all_to_all_single`'s int32 element offsets, rather than that it prevents a hang: with the slicing disabled, hops past 2**31 elements completed exactly on NCCL 2.29, 2.30 and 2.31, including at 32 ranks over the network on 2.29 and 2.30. - Against main, on tiny real models: GPT-OSS calibrates on eager with `down_proj`'s input amax equal to main's; Ernie-4.5-VL-MoE quantizes in FP8 and NVFP4 under text-only calibration, as on main. - Byte-for-byte checkpoint comparison against main on one GPU, measured on earlier revisions of #2745: Qwen3-30B-A3B, eager and `grouped_mm`, NVFP4 and FP8; Qwen2.5-0.5B (dense) is byte-identical. On GPU, Nemotron-3-Nano-30B-A3B's NVFP4 calibration with its experts pinned to eager and #3067's with `experts_implementation="eager"` give byte-identical exports, all 18,012 tensors. - The GPU tests of QAT under expert parallelism are in #2742, the first PR with an FSDP2-sharded dense part they also run. - Review fixes, 2026-10-27, carried from #2745: the FSDP2 window's ownership test (unnamed meshes, every shard state); an empty experts forward is no work the interception check expects; the non-gated experts are pinned to eager whichever plugin runs first. - Review fixes, 2026-10-27, last round, each with a test that fails without it: - `grouped_mm` is kept, and redirected to, only where `_grouped_linear` can be intercepted; one capability constant gates the interception, the supported set and the redirect target; - an eager experts forward counts as work for the interception check only when it routes a row to this rank's experts, so a rank whose experts calibration routed nothing to no longer raises alone; - an enabled expert weight quantizer needs the interception only when the calibration captures the experts' inputs to calibrate their weights (GPTQ, local Hessian, imatrix), so weight-only experts under max or MSE no longer raise; the message names what falls back. - Split from #2745 on 2026-10-27, night: #2745's and this PR's commits together give #2745's previous tree exactly, plus this round's fixes. Pre-commit, mypy included, and this PR's unit tests pass at this head on transformers 5.18 and 5.5.4 and on torch 2.8, and the series' PRs merged together add no unit-test failure over main f03651e (tests/unit/torch utils, export, quantization and models). This revision's GPU tests run in CI only. - Review suggestions, 2026-10-28: the weight-input captures (local Hessian, imatrix, GPTQ) read the path that ran, not the one the config named when they were registered, so they follow whichever forward runs and the refusal is gone, pinned by a test that fails on the code before; the interception check's forward count, and the host sync its routing test costs, stop after the first hook call of each calibration; the check's docstring says it is rank-local. Rebased over #2745's fixes, taking back the `torch_a2a` test fixtures #2745 no longer carries. - Rebased over #2745's second-round fixes: the grouped interceptor records the offsets and clears the eager expert index, and #2745's test that captures skip the grouped batch gives way to this PR's per-expert capture tests. - Review follow-up, 2026-10-28, third round: drops the note #2745 now adds, that the grouped captures skip their batch, since here they split it per expert. A new test pins that a calibration forward stops counting, and skips its host-syncing routing test, once a hook has fired; it fails with that gate removed. - Review suggestions, 2026-10-28, fourth round: the experts forward clears `_grouped_offs` when it returns, so a capture reached outside a forward skips instead of splitting its input by a finished forward's offsets; the interception check's docstring says that a rank whose experts saw no rows waits in its next collective until the launcher stops it or the collective times out, so the error is in another rank's log. Rebased over #2745's cleanup. - Rebased on 2026-10-28, evening, onto #2745's head `e7c449fcf87`, after #2745 moved onto main `7092db6ac5b` and took its last review rounds. Conflicts only in the changelog (both sides' entries kept), `_QuantFusedExperts`'s docstring (it keeps #2745's note that the `grouped_mm` fallback does not call `F.linear`) and the fused-experts tests (#2745's transposed-container tests kept). The FP8 experts wrapper's comment from #2745 now names `normalize_experts_impl_on_the_fly`, this PR's name for the function that keeps it on eager. pre-commit, mypy included, passes at this head against main. - Review fix, 2026-10-28, evening: `calibrate()`'s interception check reads the HF plugin from `sys.modules` instead of importing it, so a plugin that failed to load (which `import_plugin` only warns about) no longer fails every calibration with its error. `test_a_plugin_that_failed_to_load_leaves_calibration_alone` breaks the plugin's import and fails without the change. pre-commit, mypy included, passes at this head. - Split, 2026-10-28, evening, at the bot review's request: the FSDP2 window's ownership filter and its ownership test moved to #3088, with a test that the window leaves an ignored expert-parallel DTensor sharded. This PR sits on a helper base merging #2745 and #3088, and its source diff is 334 lines. At this head `test_fused_experts.py`, `test_moe_a2a.py` and `test_utils.py` pass on CPU (200 passed, 6 skipped), and pre-commit, mypy included, passes against main. - Review suggestions, 2026-10-28, evening: `_local_experts` refuses a DTensor whose local shard does not hold the module's experts (one replicated over the ranks, say) instead of indexing the global experts, and the native FP8 wrapper's `unpack_weight` refuses expert weights sharded over the experts instead of filling a buffer for every expert. One two-rank test covers both and fails without either. The suggested all-reduce in the interception check is declined (a collective in every distributed `calibrate()`). Focused tests pass (201 passed, 6 skipped), and pre-commit passes against main. - Review suggestions, 2026-10-28, late evening: the per-expert split of a grouped batch requires exactly one exclusive end per expert, within the rows, as transformers' `grouped_mm` and `a2a_moe` build the offsets, and raises on any other form (a leading zero would credit each expert with the previous one's rows). A test feeds three such forms and fails without the check. The `_grouped_offs` comment names all three captures and the attribute's role as the which-forward-ran signal. Rebased onto the helper base re-merged with #3088's update; focused tests pass (204 passed, 6 skipped), and pre-commit passes against main. - Review suggestion, 2026-10-28, late evening: `_local_experts` checks for a DTensor through the plugin's guarded `torch.distributed.tensor` import, so a torch without distributed does not fail every fused-MoE call. Rebased onto the helper base re-merged with #3088's update; focused tests pass (204 passed, 6 skipped), and pre-commit passes against main. - 2026-10-28, night: `hf_ptq.py` gains `--experts_implementation`, passed on every load path: `get_model` (the compressed-tensors branch included), the low-memory load, and the FSDP2 load, through the config its meta model is built from. The changelog's way back to earlier releases' scales is now reachable from the script. A test loads a tiny Qwen3-MoE with `eager` and fails without the plumbing. The `tests/examples/hf_ptq` unit tests pass (110), and pre-commit passes against main. - Review fixes, 2026-10-28, night: on the FSDP2 path the flag is set on the pre-fetched config's `_experts_implementation`, as `build_meta_causal_lm` sets `attn_implementation`, and a test captures the config the FSDP2 loader receives (it fails without the assignment); `get_model` passes the flag only when given, on every branch. `calibrate()`'s new error for a fused-experts forward that reaches neither hook has its own changelog entry. Rebased onto #2745, itself rebased onto main `bc0b6e03237` after #3088 merged; focused and `tests/examples/hf_ptq` tests pass (316 passed, 6 skipped), and pre-commit passes against main. - Review suggestions and a rebase, 2026-10-28, late night: the backend normalization warns once for each change it makes to a caller's choice, but not for a `None` backend pinned to eager, which already dispatches to the eager forward. Rebased onto #2745 on main `e9f432affbb` (#3091): `DTensor` now comes from the plugin's plain `torch.distributed.tensor` import, since the guarded import is gone with the torch 2.10 floor. Focused and `tests/examples/hf_ptq` tests pass (325 passed, 6 skipped), and pre-commit passes against main. - Review suggestion and a rebase, 2026-10-28, late night: the backend-redirect notice goes through `warn_rank_0`, so an expert-parallel run prints one copy. Rebased onto #2745's comment and warning-text update; `_is_fused_experts_module`'s new docstring names `normalize_experts_impl_on_the_fly`, which keeps GPT-OSS on eager here. `test_fused_experts.py` and `test_utils.py` pass (197), and pre-commit passes against main. - Bot review of `2e253a509d0`, 2026-10-28, late night: the FSDP2 test runs a Qwen3.5-MoE config too and checks the text config its experts dispatch on, which the one assignment reaches through transformers' recursive `_experts_implementation` setter (5.5.4 and 5.18; a tiny Qwen3.5-MoE built on meta dispatches its experts on `eager` with the flag and on `grouped_mm` without it). The test fails when the value reaches the root config only, and the comment at the assignment says the setter carries it. Declined: resolving the routing indices' parameter name from the forward's signature, since every transformers MoE block passes them positionally. The `tests/examples/hf_ptq` unit tests pass on 5.18 and 5.5.4 (120), and pre-commit passes against main. - Bot review of `042e4bf1e55` (approved), 2026-10-28, late night: the interception error names the ways forward: a class no backend change redirects needs its own quantized wrapper, and meanwhile the rest of the model quantizes with these modules' quantizers disabled, which the check honours. Declined: caching the grouped offsets' host copy across the two projections' captures (one small sync per experts module per batch under GPTQ, local Hessian and imatrix, whose per-expert matmuls dominate). pre-commit passes against main. - Rebased on 2026-10-28, late night, onto main `1ca8bd1fc4d` after #2745 merged (squashed with the tree of its last head, so this PR's tree is unchanged); the base is now `main`, so the GPU tests run. pre-commit passes against main. - CI fix, 2026-10-28, late night: the first run on main failed two cases of `tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py` in the trtllm `hf_ptq` lane (2 failed, 267 passed): its stand-in for hf_ptq's parsed arguments lacked the new `experts_implementation`, which `load_model` now reads. The stand-in gains it; the file passes on transformers 5.18 and 5.5.4. - Rebased on 2026-10-28, late night, onto main `bdc45368d10` (#2736 merged), with its own diff unchanged. Its tests with #3067's, #2736's `test_distributed.py` and the `hf_ptq` unit tests pass (376), and pre-commit passes against main. - CI, 2026-10-29: the first full run on main at `df2e3fe4a76`, GPU lanes included, passes. - Split, 2026-10-29, at a reviewer's request: the expert-parallel indexing fix, its tests and a Bug Fixes changelog entry moved to #3122 ([1/2], on main), and this PR now sits on it, its own diff otherwise unchanged. The review's three smaller items are taken: the interception check reduces its verdict over the experts' expert-parallel group (a two-rank test fails with the rank-local check); a test checks that every experts class transformers decorates with `use_experts_implementation` takes the routing indices second (transformers 5.5.4 and 5.18); and `--experts_implementation` offers `eager` and `grouped_mm` as argparse choices. pre-commit, mypy included, passes against main; this PR's tests, `test_moe_a2a.py` and the `hf_ptq` unit tests pass (315). - Review suggestion (Claude), 2026-10-29: the interception check reads the expert-parallel group from the experts' weight, a DTensor sharded over the experts from transformers 5.16 on; from 5.6 to 5.15 the shards are plain tensors that name no group, so the check stays rank-local there, as before. The docstrings say so. Restacked onto #3122's two follow-ups; pre-commit passes against main. - Comment, 2026-10-29: at `calibrate()`'s gate for the interception check, a sentence says that both conditions, the forward loop having run and the HF plugin being loaded, must come out the same on every rank, since under expert parallelism the check all-reduces over the experts' group. pre-commit passes against main. - Test fix, 2026-10-29: `tests/gpu/torch/quantization/test_fsdp2_autoquant.py`'s 2-GPU KV AutoQuantize test builds hf_ptq's arguments by hand and lacked the new `experts_implementation`, which `load_model`'s FSDP2 branch reads. It now passes `None`. The 1-GPU PR runners skip that test, and main's nightly runs it. No other test or example builds hf_ptq's arguments by hand on a path that reads the flag. pre-commit passes against main. - Rebased on 2026-10-30 onto #3122 on main `99c3834f62c`, with its own diff unchanged and no conflicts. pre-commit passes against main. - Review (sychen52), 2026-10-30: with every weight quantizer collecting statistics without quantizing (calibration), the grouped forward returns the stored weight itself instead of building a weight-sized copy slice by slice on each forward. A CPU test checks the identity and the collected per-expert amax, and fails on the previous code. On CPU, one call on a Qwen3-30B-A3B-sized `gate_up_proj` (128 × 1536 × 2048, bf16, 768 MiB) raised peak RSS by 755 MiB before and 0 MiB after; a GPU `max_memory_allocated` comparison is queued. The changelog entry now says the experts stay on `grouped_mm` for every forward after `mtq.quantize` or `mto.restore` (QAT/QAD, AutoQuantize scoring, fake-quant evaluation), with `experts_implementation="eager"` as the way back. `test_fused_experts.py` and `test_moe_a2a.py` pass on transformers 5.18 (220) and 5.5.4 (214, 6 skipped); pre-commit passes against main. - Review (kevalmorabia97), 2026-10-30: the three changes no test covered now have one each, or are gone. The `_local_experts` call in `_quantize_grouped_weight` is dropped as unreachable: the weight arrives as this rank's local shard inside the experts forward. New tests pin that a grouped forward reaching neither hook is counted whatever its routing indices hold, and that a capture after a grouped forward does not split by its offsets. The calibration pass-through calls only quantizers that calibrate (Claude's suggestion), with a constant-amax test case. Each test fails with its change reverted. `test_fused_experts.py`, `test_moe_a2a.py` and `test_hf_ptq_args.py` pass on transformers 5.18 (283) and 5.5.4 (277, 6 skipped); pre-commit passes against main. - GPU calibration peak memory, 2026-10-30 (sychen52's request): `torch.cuda.max_memory_allocated` around the `mtq.quantize(...)` call only, after `reset_peak_memory_stats`. Qwen3-30B-A3B on the default `grouped_mm` backend, recipe `general/ptq/fp8_default-kv_fp8`, `calib_size` 32, `batch_size` 1, one NVIDIA GB200 for both revisions (torch 2.11, transformers 5.13.0, CUDA 13.1). | Revision | Calibration peak (bytes) | MiB | |---|---|---| | Before the pass-through change (`3f3544435d8`) | 864,901,632 | 824.8 | | After (`27c5f77feac`) | 174,759,936 | 166.7 | | Drop | 690,141,696 | 658.2 | - Rebased on 2026-10-30 onto main `55898ee4e3b` after #3122 merged, dropping #3122's pre-squash commits; its own diff is unchanged, with no conflicts. pre-commit passes against main. - CI fix, 2026-10-30: the routing-order test read transformers' modeling sources with the platform's default encoding, cp1252 on Windows, which 16 of them (`modeling_altclip.py`, say) cannot be decoded with; it now reads them as UTF-8, and the windows unit job's only failure goes with it. The `linux` job's one failure on the previous head, a NAS sorting test this PR does not touch, passed on rerun at the same commit. pre-commit passes against main. - `eager` against main, byte for byte, 2026-10-30 (re-verifying at this PR's head the claim that `experts_implementation="eager"` reproduces main): PTQ of Qwen/Qwen3-30B-A3B with recipe `general/ptq/nvfp4_default-kv_fp8`, `calib_size` 512, calibration sequence length 2048, `batch_size` 1, cnn_dailymail, without expert parallelism (one process), on one NVIDIA GB200 (torch 2.11, transformers 5.13.0, CUDA 13.1). This PR at `27c5f77feac` with `--experts_implementation eager`, against main at `99c3834f62c`, both compared at zero tolerance. Later revisions of this PR leave its own diff unchanged (a rebase) or change tests only. | Scale kind | Tensors | Differing | |---|---|---| | `input_scale` | 18,624 | 0 | | `weight_scale` | 18,624 | 0 | | `weight_scale_2` | 18,624 | 0 | | `k_scale` | 48 | 0 | | `v_scale` | 48 | 0 | - Review (juhi10071998), 2026-10-30: under transformers' own expert parallelism the grouped forward sorts the rows routed to other ranks' experts past `offs[-1]`, which the grouped GEMM neither reads nor writes; before transformers 5.17 nothing zeroes them in the down projection's input, so its input quantizer calibrated on uninitialized rows. While an input quantizer collects statistics it now sees only the rows before `offs[-1]` (a host sync, during calibration only). A test hands the grouped projections a NaN or a large finite tail, as pre-5.17 expert parallelism leaves it, and checks the input amaxes match a run without one; both cases fail without the change. `test_fused_experts.py`, `test_moe_a2a.py`, `test_calib.py` and `test_hf_ptq_args.py` pass on transformers 5.18 (307) and 5.5.4 (301, 6 skipped); pre-commit passes against main. - Rebased on 2026-10-30 onto main `059588b89a5` (after #3131, #3124, #2751, #3104, #3128, #3134 and #3149 merged), with its own diff unchanged and no conflicts. pre-commit passes against main. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ❌ Fused-experts MoE calibration now runs on `grouped_mm`, so some activation scales differ from earlier releases; `experts_implementation="eager"` restores the earlier checkpoint exactly. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ Backward Breaking Changes. - Did you get Claude approval on this PR?: ✅ At `042e4bf1e55`; the error-message change after it is from that review. ### Additional Information Split out of #2745 (part 3 of 21 of expert-parallel PTQ for Hugging Face MoE models), which merged on 2026-10-28. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Behavior Changes** * Hugging Face MoE calibration now uses Transformers’ default `grouped_mm` backend instead of forcing eager execution. Calibrated activation scales may differ; set `experts_implementation="eager"` to use the earlier calibration behavior. * Eager-only expert wrappers continue to use eager execution. Using them with `torch_a2a` raises an error. * **Bug Fixes** * Improved calibration support for grouped and expert-parallel execution, including local expert weights and experts receiving no routed inputs. * Calibration now reports when fused experts run without interception or required inputs are not captured. * Improved FSDP2 parameter handling during weight access and writeback. * **New Features** * Added calibration support for newer Transformers fused-expert implementations, including Qwen3-MoE and native FP8 experts. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Note
Slice [1/3] of #2746's split (2026-10-30), so the model-spec work is reviewed and merged in pieces under the 500-line budget. Merge order: #2741 [1/3] (#3131, merged) → this PR → #2746 [2/3] (#3135) → #2746 [3/3] (#2746).
Base:
main, since #2741 [1/3] (#3131) merged on 2026-10-30. The diff below is this PR's alone.Series layout (the part around this PR)
main([OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741 [1/3], [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131, merged)What does this PR do?
Type of change: New feature
Before it can shard a model's experts, the expert-parallel path needs three facts about the model:
All three are facts about the architecture, so they belong with the model's other facts in
modelopt.torch.models.EPSpec, a newModelSpecsection (ModelSpec.ep_spec). Its fields:base_model_ep_plan, the plan to attach;plan_config_attr, the nested config it goes on;requires_transformers, the release this model's expert parallelism needs;prepare, for a checkpoint that ships its own modeling code;prepared_ep_plan, the planprepareattaches;unsupported, why a model cannot be expert-parallel yet.deepseek_v3carries the plan stock transformers ships only from 5.13;qwen3_moeandminimax_m3_vlcarry only the transformers release their own plan needs (a newminimax_m3_vlpackage);nemotron_hcarries its plan. It needs noprepare: its gate emits logits and the block runs its own top-k, so the experts receive global expert ids, which the all-to-all experts forward routes on directly.EPSpec.prepare, the escape hatch for checkpoints that ship their own modeling code (the Kimi adapters). A spec whose plan onlyprepareattaches states that plan too (prepared_ep_plan). It can then be read without runningprepare, which patches the model's code.llm_configand transformers' 5.18 porttext_config. The spec lists both,text_configfirst, and the first one a config has takes the plan.is_moeby the root type is False for itsNemotronHMoEblocks.EPSpec.unsupportedsays why a model whose own plan would route it into the expert-parallel path cannot take it yet:llama4_textpackage shares Llama-4's reason.Nothing reads the section yet. #2746 [2/3] (#3135) applies it to a config, and #2746 [3/3] (#2746) has the expert-parallel decision and load use it.
Usage
Testing
tests/unit/torch/models/test_ep_specs.py, the registered specs:test_model_specs.pyandtest_specs_vs_transformers.py, which check every registered spec, pass with the new packages.moe_specis left for the PR that next touchestest_specs_vs_transformers.py, as the review suggests; the gap also covers main's own such specs. The three claims added here hold on transformers 5.4, 5.5.4, 5.12, 5.18 and 5.19.tests/unit/torch/modelsandtests/unit/torch/utils, torch 2.14):tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (247) and 5.5.4 (241, 6 skipped).tests/unit/torch/modelsandtest_moe_a2a.pypass on transformers 5.18 (286) and 5.5.4 (268, 7 skipped).Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models. Split out of #2746 (OMNIML-6076), whose review history stays on #2746.
Summary by CodeRabbit