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
-
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
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 |
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.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # 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 |
|
/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
|
/ok to test c8b1bf0 |
|
/ok to test c81227c |
|
/ok to test 0c96451 |
|
/ok to test b2105f7 |
|
/ok to test ed2f8a1 |
|
/ok to test 0115662 |
|
/ok to test b42e86c |
|
/ok to test caf9b01 |
|
/ok to test 9f7e896 |
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>
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 5b3b34b |
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
main([OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122, [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 [1/3] and [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 have merged)shengliangx/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 #2740What 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.Rebased 2026-10-30 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 onto the rebuilt
shengliangx/hfep-base-6a-ep-loader, which after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131 merged holds [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 (on [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140, which movesfsdp2_wrapintomodelopt/torch/utils/fsdp2_utils.py), [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 on main 62a5c48. This PR's own diff is unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18, 5.9.0 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 e251f6c (meenchen's review round). The bounded two-rank test's docstring now names [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's new
no-gpurow beside this PR'ssharerow; this PR's own diff is otherwise unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18, 5.9.0 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 1491b66 (its MXFP4 refusal and review fixes); 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, 5.9.0 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 ad9559d (meenchen's review, round 3). [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136's tests now build their model with the shared
model_with_decoder_at, so this PR's quant-wrapper and transposed-experts tests keep a small_Experts/_SharedExpertspair of their own, which the quantization registry and the refusal messages go by, and pass the experts directly rather than through a model. This PR's own diff is otherwise unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.Rebased 2026-10-30 onto main 059588b with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 eb6c382, on
shengliangx/hfep-base-6a-ep-loaderrebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 3b814a6, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 fbff3f3 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 186efae. [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 now also refuses a rank not yet bound to its GPU. This PR's own diff is unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4. The GPU tests have not run on this head.Rebased 2026-10-30 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 56353ef, whose new commits change only tests and one comment; 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, 5.9.0 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 ab06860 onto main d464d30, on
shengliangx/hfep-base-6a-ep-loaderrebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 8f9d877, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 d452768 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 9a87d15; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18, 5.9.0 and 5.5.4.Rebased 2026-10-31 onto [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 7747c5d, which moves
test_ep_loader.py's cases intotest_model_load_utils.pyand theexpert_parallel_meshones intotest_moe_a2a.py. This PR's test edits moved with those cases, unchanged, and its source diff is unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilscollects and passes the same cases as before on CPU, on transformers 5.18, 5.9.0 and 5.5.4.Rebased 2026-10-31 with [OMNIML-6060] [2/3] Load an MoE expert-parallel: shard the experts at load time #3136 930bb0f onto main e352cac, on
shengliangx/hfep-base-6a-ep-loaderrebuilt there from [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 c78d2e5, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 adb9757 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 dea9b79; this PR's own diff is unchanged. Pre-commit, mypy included, passes against main, andtests/unit/torch/utilspasses on CPU at this head, on transformers 5.18, 5.9.0 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.