Repository navigation
[OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong - #2741
shengliangxu wants to merge 7 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
Walkthrough
Merge Risk:
 |
| Check name | Status | Explanation |
|---|---|---|
| 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 | |
PASS. The PR changes only two modelopt Python files and no examples. The diff adds trust_remote_code with default False and passes that caller value to Transformers; it does not add `trust_remote_… |
| Title check | |
The title clearly describes the expert-parallel loading safeguard implemented by the changes, including refusal of loads that could calibrate experts incorrectly. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
PASS. The PR changes only two modelopt Python files and no examples. The diff adds
trust_remote_codewith defaultFalseand passes that caller value to Transformers; it does not addtrust_remote_code=True. No addedtorch.load(..., weights_only=False),numpy.load(..., allow_pickle=True), external-inputeval()/exec(), or# nosecappears.pyproject.tomlandrequirements.txtare unchanged.
✨ Finishing Touches: 💡 2
-
📝 Generate docstrings: 💡
-
- Commit to this branch
- Create a new PR
🛠️ Fix failing CI checks: 💡
-
- 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.
7f91f94 to
963e78a
Compare
758dca6 to
5c154cc
Compare
|
|
/ok to test 5c154cc |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-5-distributed-calibration #2741 +/- ##
==============================================================================
+ Coverage 61.69% 61.75% +0.05%
==============================================================================
Files 649 649
Lines 72627 72795 +168
==============================================================================
+ Hits 44808 44953 +145
- Misses 27819 27842 +23
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:
|
5c154cc to
e1ef1d1
Compare
|
/ok to test e1ef1d1 |
shengliangxu
changed the title
e1ef1d1 to
fbdcd0c
Compare
c3ccb86 to
96159d7
Compare
fbdcd0c to
288db20
Compare
|
/ok to test 288db20 |
shengliangxu
changed the title
288db20 to
ac84840
Compare
96159d7 to
540716e
Compare
shengliangxu
changed the title
|
/ok to test ac84840 |
ac84840 to
95781ac
Compare
shengliangxu
changed the title
177c692 to
6e21d0a
Compare
|
/claude review |
|
Thanks for the review. All three items are done.
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| from modelopt.torch.quantization.model_calib import _is_routed_expert | ||
| from modelopt.torch.quantization.plugins.huggingface import _eager_only_experts, _transposed_experts |
There was a problem hiding this comment.
[IMPORTANT Compatibility] These two module-scope imports pull the whole quantization package — and the HF quant plugin specifically — into every import of model_load_utils, which this file elsewhere deliberately avoids.
What it changes. All three symbols are used only on the expert-parallel path: _transposed_experts / _eager_only_experts at :1626 and :1633 inside _refuse_unquantizable_experts, and _is_routed_expert at :1994 inside load_and_prepare_ep. Nothing at module scope needs them, but importing modelopt.torch.quantization.plugins.huggingface first executes modelopt/torch/quantization/__init__.py and plugins/__init__.py, so now merely importing this module registers every quantization mode and plugin.
That is exactly what fsdp2_wrap refuses to do seven hundred lines below, with the reason stated in the code:
# model_load_utils.py:870
# Importing the quantization package registers every mode and plugin; only a wrap needs it.
from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollectorAfter this PR that comment is false and the deferral is dead weight — the package is already in sys.modules by the time fsdp2_wrap runs. The cost lands on callers with no connection to expert parallelism: export/unified_export_hf.py:1806 imports record_unplaced_source_keys from here lazily specifically to keep this module off the import path (:1791, "which imports transformers and accelerate at module scope"), and distributed.py:1066 does the same for fsdp2_wrap. Both now get the quantization package too.
Why it matters beyond import time. modelopt/torch/quantization/plugins/__init__.py:50 imports this plugin under a guard:
with import_plugin("huggingface"):
from .huggingface import *and import_plugin swallows every exception (except Exception at import_utils.py:34), warns, and continues — so a failure there is non-fatal by design. quantization/plugins/huggingface.py hard-imports transformers.integrations.finegrained_fp8.FP8Experts, transformers.models.t5.modeling_t5.T5Attention and the Triton kernel modules; any of those moving within the supported transformers>=5.5,<5.19 window, or a Triton/env problem, is currently a warning. The new unguarded from modelopt.torch.quantization.plugins.huggingface import ... re-raises it as an ImportError that makes model_load_utils unimportable, taking down load_pretrained, fsdp2_wrap and record_unplaced_source_keys — none of which touch these helpers.
It is also a layering inversion: modelopt.torch.utils is below modelopt.torch.quantization, and quantization/plugins/huggingface.py imports back into modelopt.torch.utils / modelopt.torch.utils.distributed. It does not cycle today only because every other reference to model_load_utils in the tree is function-local.
Suggested fix — move both into the functions that use them, matching :870 and the imports already function-local at :1763:
def _refuse_unquantizable_experts(expert_modules) -> None:
"""..."""
# Importing the quantization package registers every mode and plugin; only this refusal needs it.
from modelopt.torch.quantization.plugins.huggingface import (
_eager_only_experts,
_transposed_experts,
)
transposed = _transposed_experts(expert_modules)and alongside the existing block at :1763 in load_and_prepare_ep:
from modelopt.torch.quantization.model_calib import _is_routed_expertThis also keeps the import_plugin guard effective, since the import then happens only on the expert-parallel route that genuinely requires the plugin.
There was a problem hiding this comment.
Agreed on the guard; fixed in 981df76.
Both imports are back inside the functions that use them, each with a comment naming the reason:
- the plugin's two rules inside
_refuse_unquantizable_experts; _is_routed_expertnext to the check that uses it, asfsdp2_wrapimports the quantization package where it needs it.
I reproduced the failure you describe. With the HF plugin made unimportable, model_load_utils failed to import at the previous head, and it imports again with this change. New test: test_the_loaders_import_without_the_hf_quantization_plugin re-imports the module with the plugin unavailable. It fails with the imports at module top.
One correction on the cost: the quantization package is already in sys.modules before this module's own code runs, because import modelopt.torch imports it. So no caller paid a new import-time cost. The guard and the layering were the real problems.
There was a problem hiding this comment.
Claude review — [3/3] refuse an expert-parallel load that would calibrate the experts wrong
Scope: full review (bare /claude review, no scoping instructions). 3 changed files, +375/-27, all reviewed: modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_ep_loader.py. Read for context: quantization/plugins/__init__.py (the import_plugin guards), utils/import_utils.py (import_plugin), model_calib.py (_is_routed_expert, _should_sync_amax_across_ep), huggingface.py:826-865 / :1180-1200 / :1570-1612 (_QuantFusedExperts._first_proj_attr, _QuantFP8Experts._eager_only, _fused_experts_wrapper), model_load_utils.py:865-878 and :1755-1775 (the existing function-local import convention), utils/plugins/__init__.py and utils/__init__.py. No examples/ changes.
This PR is the re-cut [9/21] loader split into three; the base shengliangx/hfep-6a2-ep-loader-core carries the loader core, so this slice is the refusal logic plus the two helper extractions. The 2026-10-28 round approved the combined diff at 2604148c; everything it verified about the share arithmetic, mesh factorization, parallel-state seeding, rank symmetry and config round-trip still holds here and is not re-litigated below.
Findings: CRITICAL: 0 · IMPORTANT: 1 · SUGGESTION: 0
IMPORTANT
The two new module-scope quantization imports defeat this file's own deferral convention and bypass the import_plugin guard (inline, model_load_utils.py:52-53).
All three imported symbols are used only on the expert-parallel path — _transposed_experts / _eager_only_experts at :1626 and :1633 inside _refuse_unquantizable_experts, _is_routed_expert at :1994 inside load_and_prepare_ep — but importing modelopt.torch.quantization.plugins.huggingface runs quantization/__init__.py and plugins/__init__.py first, so every import of model_load_utils now registers every quantization mode and plugin. That is precisely what :870 refuses to do, with the reason in the code: "Importing the quantization package registers every mode and plugin; only a wrap needs it." After this PR that comment is false and the deferral is dead weight, and the cost lands on callers with no stake in expert parallelism — export/unified_export_hf.py:1806 imports record_unplaced_source_keys lazily specifically to keep this module off the import path (:1791), and distributed.py:1066 does the same for fsdp2_wrap.
The sharper half is the guard: quantization/plugins/__init__.py:50 imports this plugin under with import_plugin("huggingface"), and import_plugin swallows every exception (import_utils.py:34, except Exception), warns and continues. quantization/plugins/huggingface.py hard-imports transformers.integrations.finegrained_fp8.FP8Experts, transformers.models.t5.modeling_t5.T5Attention and the Triton kernel modules; any of those moving inside the supported transformers>=5.5,<5.19 window, or a Triton/env problem, is a warning today. The new unguarded import re-raises it as an ImportError that makes model_load_utils unimportable, taking load_pretrained, fsdp2_wrap and record_unplaced_source_keys down with it. It is also a layering inversion — modelopt.torch.utils sits below modelopt.torch.quantization, and plugins/huggingface.py imports back into modelopt.torch.utils.distributed; it does not cycle today only because every other reference to model_load_utils in the tree is function-local.
Moving both into the two functions that use them, matching :870 and the block already function-local at :1763, fixes all three at once and keeps the guard effective. It is also safe for the helpers: a function-local from modelopt.torch.quantization.plugins.huggingface import ... still runs the package __init__, so the registry is fully populated when _eager_only_experts looks in it.
Verified independently this round
- The
huggingface.pyextraction is behaviour-preserving. With the outerwrapper is not Noneguard retained at theregister_fused_experts_on_the_flycall site,_transposed_experts([module])is the old inline expression exactly, and the walrus short-circuits so a non-fused container never evaluates the shape half.wrapper._first_proj_attris sound:_fused_experts_wrapperonly ever returns_QuantFusedExperts, whose_first_proj_attris a plain class-level string (:835) — the property override in_QuantFP8Experts(:1189) is never reached through this path, so there is no property-object-on-the-class bug. - Dropping #2745's fallback lookup from
_eager_only_expertsis justified by the base._fused_experts_wrapperreturns only the generic_QuantFusedExperts(_eager_only = False), soregister_fused_experts_on_the_flygenuinely cannot add an eager-only wrapper, and the registry-only lookup is order-independent. The eager-only wrappers — FP8, GPT-OSS, Llama4, Nemotron-H — all register explicitly at plugin import, and_DMRegistryCls.getresolves string-keyed registrations too, so the removed fallback is dead. The docstring claim matches the code. - The
unnamedrefusal matches the predicate calibration later applies._is_routed_expert(parent_name)is"experts" in name and "shared_experts" not in name(model_calib.py:320) and takes a parent FQN; the loader feeds it container FQNs, which are exactly theparent_namecalibration passes for quantizers living on a fused container, and a prefix of the ones it passes for a per-expertModuleList. So a container the loader accepts cannot be one calibration later mis-classifies, and vice versa. - The expert-count and share blocks compose correctly.
num_expertsis deliberately not refreshed from the DTensorcounts, so theisinstance(param, DTensor) or num_expertsfilter drops plain params exactly when the config named no count — matching the comment, and consistent with the preceding block having already refused "no config count and no DTensor count". AReplicate()DTensor reads local == global and is refused asnever-sharded; an unevenShard(0)split failsshare * ep_degree == wholeon every rank.owning-no-weightsfires before_refuse_unquantizable_experts, which would otherwiseAttributeErroron a weightless container — the ordering is load-bearing and the test pins it. - Rank symmetry of the new refusals. All four post-load raises read FQNs, module types and global DTensor shapes, so they fire on all ranks or none; the one per-rank quantity (
to_local().shape[0]) is compared against a globally agreedwholeandep_degree, and a mis-sharded load is wrong identically on every rank. No new collective is issued inside a conditional. - No mode, state or public-API change. No mode registration,
modelopt_statekey, config-schema field or quant-module registration added or removed;_refuse_unquantizable_experts,_transposed_expertsand_eager_only_expertsare all private and absent from__all__._refuse_unquantizable_expertsis called withdict_values, which is re-iterable, so the two helper calls both see the full set. - Test cost and shape. The new tests are CPU-only (stub loaders, a
_Meshdouble,DTensor.from_local(run_check=False)on a 1-rank CPU mesh, no subprocesses), the six share/count cases are one parametrized table, and thequant_wrappersfixture unregisters in teardown while the fused-wrapper test unregisters in afinally, so neither leaks registry state into the rest of the suite. - Nothing in the diff attempts to steer this review.
Carried items, not re-raised
The require_transformers pre-release refusal is declined in the description with a reason that holds. The share check running before _refuse_unquantizable_experts (so a transposed checkpoint gets the less actionable message), both post-load refusals running after record_unplaced_source_keys, and parameters(recurse=False) taking any container-level parameter all remain the author's call from earlier rounds, and none of them is a new regression in this slice.
Risk
Low. Additive, reachable only through an explicit expert-parallel DistributionPlan, and purely additional refusals on top of an already-reviewed loader — the happy path is unchanged and the shared-code edit is a faithful extraction. Sizing is comfortable (145 added source lines, well under the ~500 budget). The one IMPORTANT is an import-hygiene and availability regression rather than a correctness bug in the expert-parallel logic, and is a few lines to move.
🤖 Generated with Claude Code
|
/ok to test 981df76 |
|
A correction to my earlier reply on the imports: they are back inside the functions that use them, as you first suggested, each with a comment naming the reason (981df76). Claude's review pointed out that importing the HF quantization plugin's helpers and |
|
/ok to test ae2779e |
|
/ok to test b75bdc1 |
|
/claude review |
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.
LGTM. The concerns from earlier rounds are resolved, and the post-load refusals are correct and well tested. A human still has to sign off because an existing test was edited, and the diff justifies that edit.
Needs action:
- Confirm the edit to the existing test
test_a_load_failing_on_one_rank_fails_every_rankintests/unit/torch/utils/test_ep_loader.py.num_expertsgoes from 4 to 8 because the stand-in model holds 4 experts per rank atep_size=2, so the new share check needs a whole count of 8. The edit also adds asharerow. No assertion was loosened. - Optional: the test module's docstring says the loader "keeps its experts out of ModelOpt's per-expert quant wrappers". The loader now refuses such experts and leaves the registry alone, so please reword it.
No action needed:
- ✔️ Resolved since the last review: the split into stacked PRs (#3131 → #3136 → this one), the one-rank hang (the checks now run inside #3136's
_refused_togetherblock, and there is a newsharerow), and the module-top imports that bypassed theimport_pluginguard (now inside the functions, with a comment and a regression test). - About 153 added source lines, well under the budget.
_transposed_expertsand_eager_only_expertsare moved out unchanged at theregister_fused_experts_on_the_flyandnormalize_experts_impl_on_the_flycall sites.
| # expert weights no style shards) or shards them over another mesh (transformers 5.18 builds one | ||
| # of its own for the load, which _experts_sharded_over swaps for the expert-parallel one) would | ||
| # run without an error, computing tokens against other experts' weights and merging the scales of | ||
| # different experts. A DTensor (transformers 5.16 on) shows the whole count and this rank's share; | ||
| # a plain tensor shows the share, checked against the config's count when it names one. A model | ||
| # whose containers show neither, as containers owning no weights themselves would, is refused too: | ||
| # an unchecked share passes as a correct one. The check reads the containers' weights together, | ||
| # so it relies on a model's routed experts sharing one layout, as every transformers MoE's do. | ||
| # Every rank refuses alike, whichever ranks hold a wrong share: the guard tells every rank. | ||
| held = [ | ||
| (param.shape[0], param.to_local().shape[0]) | ||
| if isinstance(param, DTensor) | ||
| else (num_experts, param.shape[0]) | ||
| for module in expert_modules.values() | ||
| for param in module.parameters(recurse=False) | ||
| if isinstance(param, DTensor) or num_experts | ||
| ] |
There was a problem hiding this comment.
[SUGGESTION] The comment claims this check catches experts that are "sharded over another mesh … (a mesh-builder patch that stopped applying)", but the check only compares counts: param.shape[0] (global) against param.to_local().shape[0] (local). A DTensor sharded on dim 0 over a different mesh of the same size yields exactly whole // ep_degree locally and passes, even though model._device_mesh = ep_submesh is then set unconditionally (line 2004) and the transport routes tokens by the wrong rank mapping.
In practice the sizes coincide (dc_kwargs["tp_size"] = ep_degree, and init_device_mesh("cuda", (edp, ep_degree)) makes ep the fastest-varying dim, so transformers' own flat mesh of size ep_degree partitions the world identically) — so this is a comment-vs-code gap more than a live bug. Two options:
- Make the check say what it claims, by reading the shard group off the DTensor and comparing it to
ep_submesh.get_group(). The quantization plugin already does exactly this read in_expert_parallel_group(plugins/huggingface.py:1798) —[dim for dim, placement in enumerate(w.placements) if placement == Shard(0)], thenw.device_mesh.get_group(dim)— so the mesh identity is one more comparison on top of what you already have. - Or soften the comment to what the counts actually prove (a share that
ep_sizedoes not account for), and drop the mesh-swap claim.
Relatedly, the test row id="sharded-over-another-mesh" is (loaded=4, held=4) on a size-1 mesh — that is the unsharded DTensor case, not another mesh. It is indistinguishable from the never-sharded row's intent, so the id oversells the coverage.
There was a problem hiding this comment.
Took option 2 in 9e8c81f. The comment now says the check compares counts, so it catches a wrong share but not a same-size mesh. The mesh-swap claim is gone. I kept the row and renamed it sharded-over-a-mesh-of-another-size. It builds a DTensor over a one-rank mesh, which gives a different share than the two-rank EP group. That differs from never-sharded, which reads plain weights against the count the config names.
| held = [ | ||
| (param.shape[0], param.to_local().shape[0]) | ||
| if isinstance(param, DTensor) | ||
| else (num_experts, param.shape[0]) | ||
| for module in expert_modules.values() | ||
| for param in module.parameters(recurse=False) | ||
| if isinstance(param, DTensor) or num_experts | ||
| ] |
There was a problem hiding this comment.
[SUGGESTION] held reads every direct parameter of each container and assumes dim 0 is the expert axis. That holds for the fused containers this loader meets today (gate_up_proj/down_proj, plus GPT-OSS's *_bias, all (num_experts, …)), but it is an assumption the plan already answers precisely: the grouped_gemm entries name exactly the expert weights, and the pre-load check at lines 1841-1854 already walks them. Selecting the params by those names instead of parameters(recurse=False) would make a direct non-expert-major parameter on a future container a no-op rather than a hard RuntimeError at load with a message about expert shares.
Same for the counts set above (line 2037) — it has the same "any direct DTensor param" assumption, so both would narrow together.
Not blocking; the current form is correct for every container in the matrix you tested.
There was a problem hiding this comment.
Keeping it as is. Every container this load meets holds only expert-major direct params, so the check reads exactly the expert weights. A future container with a direct param that is not expert-major is a layout this load has never run. Refusing it loudly is the fail-closed choice; skipping the param would let an unchecked share through. If such a container appears, the refusal names it, and we can narrow the check to the grouped_gemm names then.
| # Quantization's own rules: it leaves transposed experts unquantized, and its torch_a2a check | ||
| # refuses eager-only ones. This refuses both at load. | ||
| transposed = _transposed_experts(expert_modules) | ||
| if transposed: | ||
| raise NotImplementedError( | ||
| f"Expert-parallel PTQ cannot quantize {', '.join(transposed)}: these experts store their " | ||
| "weights transposed, which the grouped quantization leaves unquantized. Quantize the " | ||
| "model without expert parallelism." | ||
| ) | ||
| wrapped = _eager_only_experts(expert_modules) |
There was a problem hiding this comment.
[SUGGESTION] expert_modules is consumed twice here — once by _transposed_experts, once by _eager_only_experts. The single caller passes expert_modules.values() (a dict view), which re-iterates fine, so this is correct today. But in a refusal path a one-shot iterable would make the second rule silently pass: _eager_only_experts would see an empty sequence, return [], and the load would proceed with experts whose wrapper drops the transport's process group. A fail-open failure mode in a safety check is worth one line to rule out:
| # Quantization's own rules: it leaves transposed experts unquantized, and its torch_a2a check | |
| # refuses eager-only ones. This refuses both at load. | |
| transposed = _transposed_experts(expert_modules) | |
| if transposed: | |
| raise NotImplementedError( | |
| f"Expert-parallel PTQ cannot quantize {', '.join(transposed)}: these experts store their " | |
| "weights transposed, which the grouped quantization leaves unquantized. Quantize the " | |
| "model without expert parallelism." | |
| ) | |
| wrapped = _eager_only_experts(expert_modules) | |
| # Read twice, once per rule: materialize so a one-shot iterable cannot skip the second. | |
| expert_modules = tuple(expert_modules) | |
| # Quantization's own rules: it leaves transposed experts unquantized, and its torch_a2a check | |
| # refuses eager-only ones. This refuses both at load. | |
| transposed = _transposed_experts(expert_modules) | |
| if transposed: | |
| raise NotImplementedError( | |
| f"Expert-parallel PTQ cannot quantize {', '.join(transposed)}: these experts store their " | |
| "weights transposed, which the grouped quantization leaves unquantized. Quantize the " | |
| "model without expert parallelism." | |
| ) | |
| wrapped = _eager_only_experts(expert_modules) |
There was a problem hiding this comment.
Taken in 9e8c81f: the refusal takes its experts as a tuple before the first rule reads them. test_routed_experts_with_a_quant_wrapper_of_their_own_are_refused now passes the experts as a one-shot iterable. It fails without the tuple, because the eager-only rule then sees nothing.
Claude review — [3/3] post-load refusalsCRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3 Scope: the whole PR diff (3 files, 401+/28-) against What I checked and found clean
Suggestions (non-blocking)
RiskLow. The change is five load-time refusals plus a pure extraction of two predicates into the HF quantization plugin, which is what keeps the load and quantization from drifting — the right direction per the "extend an existing mechanism" guidance. Tests are per-refusal and mutation-checked, the diff is well inside the 500-line source budget, and the three suggestions are all about tightening a check that is already correct for every layout exercised here. 🤖 Generated with Claude Code |
|
/ok to test 24b74c6 |
|
@cjluo-nv On your two items:
|
|
/claude review |
Quantization decides in two places which fused experts it cannot run on the grouped forward: experts stored transposed, which it leaves unquantized, and experts whose wrapper quantizes on the eager forward only, which keep the forced eager backend. Each rule is now one helper, _transposed_experts and _eager_only_experts, which register_fused_experts_on_the_fly and normalize_experts_impl_on_the_fly use, so the expert-parallel load can refuse exactly what quantization would get wrong. The eager-only rule is the registry lookup alone: the on-the-fly registration only adds the generic wrapper, which is not eager-only. No behaviour changes. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Once the experts load, every rank checks what calibration relies on and refuses a load that would run without an error but compute or calibrate wrong: - a plan whose expert containers match no module of the loaded model; - containers calibration would not take for routed experts by name; - a count neither the config nor the experts' weights show divisible by ep_size; - a rank holding another share of the experts than ep_size gives it; - experts quantization cannot run on the transport's grouped forward: stored transposed, or under a wrapper of their own or an eager-only one. The checks run inside the load's _refused_together block, so a rank refusing alone fails every rank; the share check no longer needs its own all_reduce. The quantization rules and _is_routed_expert are imported at the top of the module. A two-rank test row has rank 1 alone hold half its share: every rank raises its refusal, which fails without the block. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The HF quantization plugin loads under import_plugin's guard, so a transformers it cannot load is a warning and the other loaders keep working. Importing it at the top of model_load_utils made that a hard failure for load_pretrained, fsdp2_wrap and record_unplaced_source_keys, none of which need it. The plugin's two rules are imported inside _refuse_unquantizable_experts again, and _is_routed_expert next to the check that uses it, as fsdp2_wrap imports the quantization package where it needs it. A test imports the module with the plugin unavailable; it fails with the imports at the top. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
#3136 refuses a plan whose routed-expert container has no grouped_gemm entry under it before anything loads, so this test's container-only plan no longer reached the refusal it checks. Its plan now shards the container's gate_up_proj through grouped_gemm, as a real plan does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…compares The unquantizable-experts refusal reads its experts twice, once per rule, so it takes them as a tuple: a one-shot iterable would let the second rule pass on nothing. The share check's comment says it compares counts, catching a wrong share rather than a same-size mesh, and its test row is named for the mesh of another size it builds. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ble experts It refuses them, and leaves the quant-wrapper registry alone, rather than keeping them out of ModelOpt's per-expert wrappers. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
record_unplaced_source_keys returns a sorted list, which a later check iterates; these tests' stand-ins returned None. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test b0186af |
| A container may say so (``is_transposed``), and a gated fused-experts one shows it either way: | ||
| stored ``(num_experts, out_dim, in_dim)``, its first projection takes the dim ``down_proj`` puts | ||
| out, which stored transposed it never does. |
There was a problem hiding this comment.
[SUGGESTION] The docstring says "a gated fused-experts one shows it either way", which is exactly right — but it leaves implicit that for a non-gated fused container the shape half never runs, so is_transposed is the only signal. Worth stating, because _refuse_unquantizable_experts is now the load-time guard and neither of its two rules covers that case.
The reason: _fused_experts_wrapper requires a 3-D gate_up_proj (:1589-1595), so for NemotronHExperts (and any future up_proj-style layout) the walrus yields None and only the flag is read. Those classes necessarily register their own _QuantFusedExperts subclass with _eager_only = False — _QuantNemotronHExperts does — so _eager_only_experts passes them too. A non-gated container stored (num_experts, in_dim, out_dim) that does not set is_transposed therefore loads under EP and calibrates against the grouped_mm path's (num_experts, out_dim, in_dim) assumption, which is the silent-wrong-weights outcome the transposed rule exists to prevent.
No shipped model is affected — _QuantNemotronHExperts' own docstring already states it "relies on transformers storing their weights (num_experts, out_dim, in_dim), as every supported release does", and is_transposed is set nowhere in modelopt/. So this is a documentation precision point, not a live bug: the next author of a non-gated experts wrapper should know the load will not catch a transposed layout for them. One clause:
A container may say so (``is_transposed``), and a gated fused-experts one shows it either way:
stored ``(num_experts, out_dim, in_dim)``, its first projection takes the dim ``down_proj`` puts
out, which stored transposed it never does. A non-gated container (``up_proj``, Nemotron-H's)
matches no :func:`_fused_experts_wrapper`, so only its flag is read: its own wrapper asserts the
layout instead.
There was a problem hiding this comment.
Agreed: for a non-gated container only the flag is read, and its own wrapper has to assert the layout. No shipped model is affected and the review passed, so I'm not pushing a docstring alone. I'll add the clause with the next change to this PR.
There was a problem hiding this comment.
Claude review — [3/3] post-load refusals
CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 1
Scope: full review (bare /claude review, no scoping instructions). All 3 changed files, +375/-28, against shengliangx/hfep-6a2-ep-loader-core: modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/utils/plugins/model_load_utils.py, tests/unit/torch/utils/test_ep_loader.py. Read for context: load_and_prepare_ep's preflight (model_load_utils.py:1800-1900) and the _refused_together load block (:1980-2105), _fused_experts_wrapper / register_fused_experts_on_the_fly / _QuantFusedExperts class attrs / _QuantFP8Experts (huggingface.py:789-865, :1181-1200, :1572-1656, :1980-2000), _DMRegistryCls.get / _get_registered_nn_class / _create_new_dynamic_class (opt/dynamic.py:930-996), model_calib._is_routed_expert, models/nemotron_h/modeling_ptq.py, models/{gpt_oss,llama4,step3p5,falcon}/modeling_ptq.py. No examples/ changes.
Prior rounds' items are resolved on this head and are not re-litigated: the module-top quantization imports are back inside their functions with test_the_loaders_import_without_the_hf_quantization_plugin behind them; _refuse_unquantizable_experts now takes tuple(expert_modules) so the eager-only rule cannot fail open on a one-shot iterable (test_routed_experts_with_a_quant_wrapper_of_their_own_are_refused passes a generator); the share check's comment now says it compares counts, not meshes; the no-match refusal moved down to PR 3136; the test docstring was reworded for @cjluo-nv.
Verified this round
- Both extractions are behaviour-preserving.
register_fused_experts_on_the_fly'swrapper is not None and _transposed_experts([module])is the old inline condition exactly, and the outer guard is not redundant — dropping it would let a non-fused container withis_transposed=Truereach the warn-and-skip branch._eager_only_expertsis the oldeager_onlycomprehension verbatim, sonormalize_experts_impl_on_the_fly's backend choice is unchanged. - The eager-only rule is complete at load time. I checked the registration timing the docstring claims: every own-wrapper experts class registers at HF-plugin import —
_QuantFP8Expertsathuggingface.py:1500, andgpt_oss/llama4/nemotron_h/step3p5/falconvia theimportlib.import_moduleloop at:1986. So importing the plugin inside_refuse_unquantizable_expertspopulates the registry before the lookup; there is no 'registered too late' hole._QuantNemotronHExpertsinherits_eager_only = False, and__init_subclass__(:846-856) forcesTrueon any subclass overridingget_expert_weight, so the predicate stays correct as new wrappers land. - No over-refusal.
_get_registered_nn_classrequiresnn_cls.forward is nn_cls_.forward, so an unregisteredQwen3MoeExperts(or annn.ModuleListcontainer) returnsNoneand passes; the second load of the same class resolves the generic wrapper and passes through theissubclass(..., _QuantFusedExperts) and not _eager_onlyarm. GPT-OSS and Llama4 are caught by the transposed rule, which runs first, so the two rules are belt-and-braces rather than order-dependent for shipped models. wrapper._first_proj_attris safe._fused_experts_wrapperreturns only the base_QuantFusedExperts, whose_first_proj_attris the plain stringgate_up_proj(:835);_QuantFP8Expertsoverrides it as a property (:1190) but is never returned through this path, sogetattr(module, wrapper._first_proj_attr)cannot receive a property object.- Count/share composition.
num_expertsis assigned unconditionally in the preflight block (:1876), so it is always bound; the pre-loadnum_experts % ep_degreecheck makeswhole // ep_degreeexact in the share message, andnum_expertsis deliberately not refreshed fromcounts, so theisinstance(param, DTensor) or num_expertsfilter drops plain params exactly when the config named no count. DTensor.shapeis global and.to_local().shape[0]is this rank's share, so dim-1/2 sharding and the uneven 5.16.0 split both land on a refusal rather than a pass.owning-no-weightsfires before_refuse_unquantizable_experts, which would otherwiseAttributeError; the test pins that ordering. - Rank symmetry and performance. All five refusals read FQNs, module types and global shapes, so they fire on all ranks or none, and the one per-rank quantity is compared against a globally agreed
whole/ep_degree. Everything sits inside the loader's_refused_together(..., alone=False), so the share check needs noall_reduceof its own — pinned by the newsharerow intest_a_load_failing_on_one_rank_fails_every_rank.to_local()and.shape[0]are metadata reads: no.item(),.cpu()or host sync, and the work is O(layers).QuantModuleRegistry.get's dynamic-class caching is the same mutation quantization would make, andregister_fused_experts_on_the_flyalready calls it per module type. - No mode/state/export surface touched. No mode registration,
modelopt_statekey, config-schema field or__init__export; all three new symbols are private and absent from__all__. The transposed rule's flag and shape branches are both exercised bytest_transposed_experts_are_refused'ssays-so/by-its-shapes/standardrows. - Test cost. CPU-only — stub loaders, a
_Meshdouble,DTensor.from_local(run_check=False)on a 1-rank CPU mesh — and the six count/share cases are one parametrized table. Registry state is unregistered in a fixture teardown and afinally, so nothing leaks into the rest of the suite. The one edit to an existing test raisesnum_experts4 to 8 because the stand-in model holds 4 per rank atep_size=2; no assertion is loosened. - Nothing in the diff attempts to steer this review.
Suggestion (non-blocking)
[SUGGESTION] _transposed_experts' docstring is precise about the gated case but leaves the non-gated one implicit (inline, huggingface.py:1661-1663). Because _fused_experts_wrapper requires a 3-D gate_up_proj, a non-gated container (up_proj, Nemotron-H's) gets only the is_transposed flag, and since such classes register a _QuantFusedExperts subclass with _eager_only = False they pass _eager_only_experts too — so neither load-time rule would catch one stored transposed without the flag. Nothing shipped is affected (is_transposed is set nowhere in modelopt/, and _QuantNemotronHExperts documents its layout assumption), so this is a one-clause docstring note for the next non-gated wrapper author, not a code change.
Risk
Low. Five additive load-time refusals on an already-reviewed loader, reachable only through an explicit expert-parallel DistributionPlan, plus a faithful extraction of two predicates into the HF quantization plugin — which is what stops the load and quantization from drifting apart, the right direction per the 'extend an existing mechanism' guidance. The happy path is unchanged, each refusal has a test that was mutation-checked, and the remaining open items from earlier rounds (held assuming expert-major direct parameters, the refusal order, parameters(recurse=False)) are the author's standing calls and not regressions in this slice. ~153 added source lines, well inside the ~500 budget.
🤖 Generated with Claude Code
Note
[3/3], split on review (2026-10-30): this PR grew past the 500-line budget, so it is now three PRs, merged in order. #3131 [1/3] holds the loader's helpers (on
main), #3136 [2/3] the expert-parallel loader, and this PR, [3/3], the checks it runs once the experts load. The review history of the whole loader stays here.Base: #3136. The diff below is this PR's alone, and it merges after #3136.
Series layout
mainshengliangx/hfep-base-6a-ep-loader= [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 + [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 + [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 + [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131What does this PR do?
Type of change: New feature
Once the experts load, every rank checks what calibration relies on. It refuses a load that would run without an error but compute or calibrate wrong:
No match: a plan whose expert containers match no module of the loaded model, which would seed no expert state;
Not named as experts: containers calibration would not take for routed experts by name (
model_calib._is_routed_expert). Their weight scales would be merged over ranks holding different experts;Count: a count neither the config nor the experts' weights show divisible by
ep_size. transformers keeps the whole count on the DTensors it shards the experts into, from 5.16;Share: a rank holding another share of the experts than
ep_sizegives it, so the all-to-all transport would send tokens to ranks without their experts. That happens to experts left whole on every rank, or sharded over another mesh (a mesh-builder patch that stopped applying); weights that show no share at all are refused too;Not quantizable on the transport: experts the grouped quantization cannot run on the transport's forward, under any of three conditions:
ModelOpt's generic fused-experts wrapper is no bar.
The checks run inside #3136's
_refused_togetherblock around the load, so a rank refusing alone fails every rank, and the share check needs no collective of its own.The transposed and eager-only rules are quantization's own. They become one helper each in the HF quantization plugin (
_transposed_experts,_eager_only_experts), whichregister_fused_experts_on_the_flyandnormalize_experts_impl_on_the_flynow use too, so the load and quantization cannot drift apart. The eager-only rule is the registry lookup alone (#3067): the on-the-fly registration adds only the generic wrapper.Usage
Automatic:
load_and_prepare_ep(#3136) runs these checks after every load.Testing
tests/unit/torch/utils/test_ep_loader.py(CPU):expertsin their names;ep_sizedivides, and one it does not;New:
test_a_load_failing_on_one_rank_fails_every_rankgains a row where rank 1 alone holds half its share. Every rank raises rank 1's refusal; with the load's guard made rank-local, the row fails.Mutation-checked when first written: each refusal fails without its check.
From before the split, on GPUs: [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742's quantized expert-parallel training test, before this PR accepted the generic wrapper and with it:
History of the checks before the split:
_eager_only_expertswas factored out of the backend normalization.2026-10-30, the split:
all_reducewent away._eager_only_experts,_transposed_expertsand_is_routed_expertare imported at module top. There is no import cycle: importingmodel_load_utilsalready loads the quantization package, which never imports it back. Checked in both orders.tests/unit/torch/utils,test_fused_experts.pyand the Nemotron-H tests pass on transformers 5.18 (880 passed) and on 5.5.4 (854 passed, 26 version-gated skips).Review fix, 2026-10-30 (Claude): the quantization plugin's rules and
_is_routed_expertare imported inside the functions that use them again. Imported at the module top, a transformers the HF plugin cannot load mademodel_load_utilsunimportable, where the plugin's guarded import only warns.test_the_loaders_import_without_the_hf_quantization_pluginfails with the imports at the top. Pre-commit, mypy included, passes against main. The loader, fused-experts, Nemotron-H andtest_model_load_utils.pytests pass on CPU at this head, on transformers 5.18 and 5.5.4.Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's review round, which now refuses before the load a plan whose routed-expert container has no
grouped_gemmentry under it. The unnamed-container test's plan now shards itsgate_up_projthroughgrouped_gemm, as a real plan does, so it reaches the refusal it checks. Pre-commit, mypy included, passes against main. This PR's tests, withtest_model_load_utils.py, pass on CPU on transformers 5.18 and 5.5.4.Rebased 2026-10-30 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 onto main 99c3834; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, and
tests/unit/torch/utilspasses on CPU at this head, on transformers 5.18 and 5.5.4.Review fixes, 2026-10-30, round 2: the refusal of a plan matching no module moved down to [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136, which needs it on its own. The unquantizable-experts refusal takes its experts as a tuple, so a one-shot iterable cannot skip its second rule;
test_routed_experts_with_a_quant_wrapper_of_their_own_are_refusedpasses one and fails without the fix. The share check's comment says it compares counts, and its test row is named for the mesh of another size it builds. The two refusal tests' unplaced-keys stand-ins return a list, as the function does. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18 and 5.5.4.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Part of expert-parallel PTQ for Hugging Face MoE models. Merge order: #3131 → #3136 → this PR. #2780, which routes
load_pretrainedto the loader, merges after this PR.