Repository navigation
[OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized - #2745
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:
No actionable comments were generated in the recent review. ℹ️ Recent review info
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 | |
No stated security anti-pattern was introduced. The PR changes only modelopt Python files and tests; it adds no examples or dependency-file changes. Added modelopt lines contain no unsafe torch.load, … |
| Title check | |
The title clearly summarizes the primary change: quantizing fused experts under grouped_mm while leaving transposed layouts unquantized. |
| Description Check | |
Check skipped - CodeRabbit’s high-level summary is enabled. |
Full details: Security Anti-Patterns
-
Explanation
No stated security anti-pattern was introduced. The PR changes only modelopt Python files and tests; it adds no examples or dependency-file changes. Added modelopt lines contain no unsafe torch.load, pickle-enabled numpy.load, hardcoded trust_remote_code=True, eval/exec on external input, or # nosec comments. The existing trust_remote_code=True call in huggingface.py is unchanged from the base revision.
✨ 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.
06a5f2d to
5f115f2
Compare
6e6ac9c to
bad2de9
Compare
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2745 +/- ##
==========================================
+ Coverage 73.82% 79.74% +5.92%
==========================================
Files 662 662
Lines 75949 75979 +30
==========================================
+ Hits 56069 60590 +4521
+ Misses 19880 15389 -4491
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:
|
|
/ok to test bad2de9 |
5f115f2 to
d6a4b30
Compare
bad2de9 to
16e5c17
Compare
|
/ok to test 16e5c17 |
shengliangxu
changed the title
d6a4b30 to
e28bd09
Compare
5c0d5fa to
698f459
Compare
|
/ok to test 698f459 |
shengliangxu
changed the title
698f459 to
eaa41c2
Compare
235ecbd to
9572dc5
Compare
shengliangxu
changed the title
|
/ok to test eaa41c2 |
9572dc5 to
6e7ca4f
Compare
eaa41c2 to
3ff1d97
Compare
shengliangxu
changed the title
…efusal The torch_a2a expert-parallel forward runs the grouped linear too, so loading with experts_implementation="eager" alone does not avoid it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The grouped interceptor clears the expert index, so on grouped_mm local Hessian, imatrix and GPTQ calibrate the experts' weights without their inputs, while max calibration is unaffected. The class docstring said both backends calibrate the same quantizers without that exception. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…d forward The passthrough shortcut skips only quantizers that would leave their slice unchanged; the rotation case fails with its rotate_is_enabled term dropped. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_fused_experts_wrapper's body comment said again, at greater length, what its docstring says about act_fn, and _is_quant_fused_experts_module had moved above _get_fused_experts_quantizer_attr_names with nothing calling for it. Both now read as on main. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…onger restores Restore converts the model as quantize does, which no longer registers the transposed experts type, so a saved state's quantizer keys for those experts match no module (ApplyModeError). The entry now tells users so. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test imported transformers.integrations.moe and read its private _grouped_linear unconditionally, while the plugin and the unit tests treat both as optional. It now skips only the grouped_mm case without them, and spies on _grouped_linear only where it exists. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ntizers The per-slice loops ran over num_experts while the stacked kernel goes by the weight's rows: a weight with more rows than quantizers would have sent unwritten rows into the grouped matmul. A strict zip makes any mismatch fail. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Without autograd the per-slice path wrote each fake-quantized slice into a buffer of the weight's dtype, so a pre-quant scale wider than the weights (FP32 scales on a BF16 model) was rounded back to BF16, while the autograd path's stack kept FP32: the same forward computed from different weights in the two grad modes. The buffer now takes the quantizers' output dtype. The new test runs the grouped forward both ways on BF16 weights with FP32 input and weight scales and fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The changelog says the Hugging Face export refuses an MoE block whose transposed experts were left unquantized, and that a state saved when an earlier release quantized them no longer restores. One test runs the export's MoE preparation on a block is_moe recognizes and expects NotImplementedError; the other quantizes such a container the way an earlier release did, then restores its state into a fresh model and expects ApplyModeError. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A modelopt_state saved for a model with transposed generic fused experts no longer restores; that is a break for whoever holds one, not a bug fix, so it gets its own entry there. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…sagree Without autograd the per-slice buffer took the first slice's dtype, so an expert whose pre-quant scale widens it past the first one's had its slice narrowed, while the autograd path's torch.stack promotes every slice to the widest dtype. The buffer now moves to the promoted dtype when a later slice is wider. The widening test gains a case where only a later expert is widened; it fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Installing the F.linear and _grouped_linear patches together is safe only because the grouped_mm fallback does not route through F.linear; the class docstring now says so, and that the CPU tests running the fallback pin it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
… intercepted force_eager_experts_impl_on_the_fly's docstring said every non-eager backend bypasses the quantizer hooks; after this change only batched_mm and sonicmoe do. grouped_mm is intercepted, and calibration stays on eager until the per-expert input captures follow it there. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its _eager_only refusal sits in the grouped interceptor, which transformers' FP8 grouped_mm forward never calls, so for this wrapper the flag cannot fire; force_eager_experts_impl_on_the_fly, on convert and restore, is what keeps it on eager. The comment now says so. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The shape test separates the layouts of a gated container, whose first projection is 2 * intermediate; a non-gated container's two layouts have the same shapes, so only its own is_transposed gives it away. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
juhi10071998
left a comment
juhi10071998
left a comment
There was a problem hiding this comment.
Thanks for this, the grouped_mm interception is a nice piece of work and the comments explaining the layout assumptions are really helpful. I left three inline notes. The first two are the ones I'd most like your thoughts on, since both can end in a quiet wrong result rather than an error.
| # puts out. That separates the layouts for a gated container, whose first projection is | ||
| # `2 * intermediate`, whether it says so or not; a non-gated one's two layouts have the | ||
| # same shapes, so only its own `is_transposed` gives it away. | ||
| or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1] |
There was a problem hiding this comment.
Nice catch on the gated case, and I appreciate that the comment above spells out the limit for non-gated containers.
That limit is what worries me a bit. For a non-gated container the two layouts have identical shape relationships (transposed: up (E,H,I) / down (E,I,H), so up.shape[-1] == down.shape[1]; non-transposed: up (E,I,H) / down (E,H,I), which also satisfies it). So the shape check can't tell them apart, and the guard is only as good as the container exposing is_transposed. A transposed non-gated class without that attribute would still be registered as _QuantNonGatedFusedExperts, and then hit the wrong-axis export slice this PR is trying to avoid.
A couple of thoughts, happy to be told I'm missing something:
- Would it make sense to fail loudly rather than
warnings.warnhere? As written, PTQ runs through calibration with those experts unquantized and the problem only shows up later (or not at all, if the warning scrolls by). If skipping is intentional, maybe say in the warning that the model will export with unquantized experts. - Since the filter lives here and not in
_fused_experts_wrapper,_is_fused_experts_module(andforce_eager_experts_impl_on_the_fly) still see these modules as fused. Is that intended? Moving the check into_fused_experts_wrapperwould keep all three consistent.
There was a problem hiding this comment.
Thanks. I took part of this in 75265cb and kept the rest, for reasons I checked:
-
Warning text. It now says the experts stay unquantized and that
export_hf_checkpointwill refuse their MoE block. That is what happens: the export finds no handler for unconverted experts and raisesNotImplementedError, whichtest_a_transposed_containers_moe_block_is_refused_at_exportpins. Nothing gets exported with wrong weights, and nothing gets exported with silently unquantized experts. -
Warn rather than raise. Registration runs inside
mtq.quantizebefore the config's patterns apply, so it cannot tell whether the recipe quantizes these experts at all. Raising there would block two things that work with the experts unquantized: a recipe that leaves the experts alone, and fake-quant evaluation without an export. The export, which is where wrong weights would land, does fail loudly. -
The filter stays out of
_fused_experts_wrapperon purpose._is_fused_experts_moduleonly feedsforce_eager_experts_impl_on_the_fly, and transposed containers need that forcing.GptOssExpertssetsis_transposed=True, and GPT-OSS's own wrapper quantizesdown_proj's input by intercepting the eager forward'storch.bmm/Tensor.__matmul__, whichgrouped_mmnever calls. I checked this on a tiny GPT-OSS (transformers 5.4, CPU,INT8_DEFAULT_CFG, default experts backend):- As the code is, the model is forced to eager and
down_proj_input_quantizercalibrates (amax 0.377). - With the filter moved into
_fused_experts_wrapper, the model stays ongrouped_mmand that amax staysNone.
_is_fused_experts_module's docstring now says this. - As the code is, the model is forced to eager and
-
The non-gated gap. Agreed: the shape test cannot tell those two layouts apart, and the comment says so. Every class that can dispatch to
grouped_mmgetsis_transposedfrom@use_experts_implementation. In transformers 5.4, 5.18 and 5.19 the transposed classes areGptOssExperts,AriaExpertsandOpenAIPrivacyFilterExperts, all gated. The only non-gated class,NemotronHExperts, is not transposed.
| # The batch holds every expert's rows: no single expert for the eager loop's index to | ||
| # name, so a weight-input capture keyed on it skips the batch rather than credit | ||
| # one expert with all of it. | ||
| self._current_expert_idx = None |
There was a problem hiding this comment.
I like that the comment explains why this is None (crediting one expert with the whole batch would be worse). My concern is what happens downstream.
I checked the hooks in model_calib.py: the weight-input capture does if idx in enabled: and the GPTQ hook does per_expert.get(experts._current_expert_idx). With None, both just skip, silently. So someone who enables local_hessian / imatrix / GPTQ on a model running grouped_mm gets experts calibrated with no input statistics and no message at all. Also, the _fused_expert_projections docstring still says the index is "already set to that expert", which is no longer true on this path.
Could we either (a) raise or warn once when one of those algorithms is active and the grouped path is hit, or (b) force the eager path during calibration for those algorithms? At minimum I'd update that docstring. Plain max calibration looks unaffected, so this may only matter for the heavier algorithms, but that's exactly where the silence is hardest to notice.
There was a problem hiding this comment.
I updated the docstring in 75265cb. Under eager, the input quantizer runs once per expert with the index set. Under grouped_mm it runs once per projection, on every expert's rows, with the index cleared, so a capture keyed on the index skips that batch.
On (a) and (b): this PR already does (b), for every algorithm. force_eager_experts_impl_on_the_fly still pins the model to eager on convert and on restore, so mtq.quantize calibrates on eager. The grouped interceptor runs during calibration only if the caller switches the backend back between quantize and a later mtq.calibrate. The class docstring names that case: there, local Hessian, imatrix and GPTQ calibrate the experts' weights without their inputs, and max is unaffected.
#3074, stacked on this PR, is where calibration moves onto grouped_mm. It splits the grouped batch per expert by its offsets, so those captures follow it there. A warning added here would come out again in that PR.
| if _transformers_moe is not None and hasattr(_transformers_moe, "_grouped_linear"): | ||
| _orig_grouped_linear = _transformers_moe._grouped_linear | ||
|
|
||
| def _quantized_grouped_linear(input, weight, offs, bias=None, is_transposed=False): |
There was a problem hiding this comment.
Smaller one, and more of a defensive suggestion than a bug I can point to: is_transposed is accepted and forwarded, but the weight quantization below assumes the (E, out, in) layout. The registration filter should keep transposed containers out, so I don't think this is reachable today. Still, a two-line if is_transposed: raise RuntimeError(...) here would turn any future path that slips past the filter into a clear error instead of wrong scales applied to a transposed weight. Totally fine to skip if you'd rather not add the branch.
There was a problem hiding this comment.
Leaving it out, as in the earlier round that asked for the same check. Through registration the flag cannot be true here:
- Transformers' grouped forward passes
is_transposed=self.is_transposed, the container's own attribute. - Registration leaves every container with
is_transposedset unquantized, so no_QuantFusedExpertsit registers sees the flag. - The transposed experts that have wrappers of their own (GPT-OSS, Llama-4) are not
_QuantFusedExperts.
The only way past this is registering _QuantFusedExperts by hand for a transposed class. That skips the filter on purpose, and the eager forward would be just as wrong for such a class, so a check here alone would not cover it.
|
/ok to test 1274887 |
The two ownership tests (#3088) counted fully_shard.state calls to show the weight-access window does not rescan a unit. #3091 made _get_fsdp2_mesh call fully_shard.state on every window entry, so the count no longer measured scans and both tests failed on main. They now count _fsdp2_param_groups, the scan's own walk; removing the cache, or rescanning on every miss, still fails them. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 3cfe6e0 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| # puts out. That separates the layouts for a gated container, whose first projection is | ||
| # `2 * intermediate`, whether it says so or not; a non-gated one's two layouts have the | ||
| # same shapes, so only its own `is_transposed` gives it away. | ||
| or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1] |
There was a problem hiding this comment.
[SUGGESTION] This reads _first_proj_attr off the class, which happens to work only because the two wrappers _fused_experts_wrapper can return (_QuantFusedExperts, _QuantNonGatedFusedExperts) define it as a plain class-level string. A sibling wrapper in the same hierarchy — _QuantFP8Experts — already defines it as a @property, and the moment _fused_experts_wrapper can return something like it, wrapper._first_proj_attr yields a property object and getattr(module, <property>) raises TypeError: attribute name must be string from inside registration, i.e. at mtq.quantize/restore time rather than at definition time.
Since the module is already known to carry the attribute (that is what _fused_experts_wrapper just matched on), you can decouple from the class and read it from the module instead:
first_proj = getattr(module, "gate_up_proj", None)
if first_proj is None:
first_proj = module.up_proj
...
or first_proj.shape[-1] != module.down_proj.shape[1]or, if you prefer keeping the class as the source of truth, make the contract explicit where it is consumed (assert isinstance(wrapper._first_proj_attr, str)). Either way the failure mode moves from a confusing TypeError deep in registration to something a future wrapper author sees immediately.
Not blocking — today's two return values are both strings.
There was a problem hiding this comment.
Leaving it as is:
wrapperhere is only ever what_fused_experts_wrapperreturned two lines up,_QuantFusedExpertsor_QuantNonGatedFusedExperts. Both define_first_proj_attras a class string._QuantFP8Experts, whose property you mention, is registered explicitly and never reaches this line.- If
_fused_experts_wrapperever returned a class with a property, this line would raiseTypeError: attribute name must be stringon that class's first registration, in every test of it, so it would not go unnoticed.
Reading the attribute off the module instead would repeat _fused_experts_wrapper's layout logic.
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review (the trigger comment carried no scoping instructions). 7 files changed (+747/−41); I read all of modelopt/ (plugins/huggingface.py, models/step3p5/modeling_ptq.py, CHANGELOG.rst) plus the three test files, and traced _grouped_linear / _current_expert_idx / _first_proj_attr / _is_fused_experts_module across model_calib.py, tensor_quantizer.py, utils/moe.py and plugins/custom.py.
Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 2
What I verified
- Grouped/eager equivalence. The interceptor applies the same shared input quantizer to the expert-sorted batch and the same per-expert weight quantizer to each dim-0 slice, so for max calibration the amax over the concatenated batch equals the max over eager's per-expert subsets (same token multiset), and the
(num_experts, out, in)layout makes sliceidxidentical in both paths. gate/upchunk(2, dim=-1)lands on the same halves, so anaxis=0weight quantizer scales the same rows. - Toggle safety.
_down_proj_linearis shared by both interceptors and reset inforward(); a desync would be loud (quantizers[None]->TypeError) rather than silent, and the CPU tests that run transformers' realgrouped_mmfallback plustest_fake_quantized_grouped_forward_and_backward_match_eagerpin that_grouped_linearfires exactly twice and thatF.linearis not re-entered inside it. _current_expert_idx = None. Both consumers tolerate it without aKeyError/TypeError—idx in enabledin_register_weight_input_hooksandper_expert.get(...)ingptq— so local-Hessian / imatrix / GPTQ genuinely skip the grouped batch instead of crediting it to expert 0.test_weight_input_captures_skip_the_grouped_forwardpins it, and the class docstring now states the consequence.- Passthrough shortcut. The disabled-quantizer early return is sound: the only other things
TensorQuantizer.forwarddoes when_disabledarepre_quant_scale, rotation, and_input_dtype.rotate_back_is_enabledimpliesrotate_is_enabledin both theRotateConfigand legacy-dict branches, sonot q.rotate_is_enabledalready covers it;_input_dtypeis assigned but read nowhere in the tree, so skipping it is inert.SequentialQuantizeris not aTensorQuantizer, so W4A8-style weight stages correctly fall through to the slice-by-slice path. - Dtype promotion without autograd. The incremental
promote_typesfold reaches the same final dtypetorch.stackwould (promotion is associative),out.to(same_dtype)is a no-op returning self so narrower later slices cost nothing, and_reset_to_original_shapekeeps static-block quantizers shape-preserving soout[idx] = quantizednever mismatches. torch_a2ainteraction.modelopt/torch/utils/moe.pydoes itsfrom transformers.integrations.moe import _grouped_linearinsidetorch_a2a_experts_forward, so it resolves the patched attribute at call time and expert-parallel forwards are intercepted too — which is exactly why the_eager_onlymessage naming expert parallelism is the right wording.- Transposed detection. The shape test is correct for every layout combination I enumerated, including
hidden == intermediateandhidden == 2 * intermediate; the non-gated transposed case is genuinely indistinguishable by shape and correctly falls back to the module's ownis_transposed, as the comment says. Qwen3-VL-MoE (hidden=32,moe_intermediate=12) is non-transposed, so the existingTestQwen3VLMoeTextExpertscoverage still exercises real quantization rather than silently degrading to a no-op. - Backward compatibility. The restore break is real, deliberate, documented under Backward Breaking Changes, and pinned by
test_a_state_saved_with_transposed_experts_quantized_no_longer_restores; the checkpoints that state produced had gate/up split on the input axis, so there is indeed nothing to migrate._fused_experts_wrapper_class->_fused_experts_wrapperis module-private with no call sites outside this file andmodeling_ptq.py's docstring, both updated. - Memory. The extra full-size 3-D buffer is inherent to handing
_grouped_mma fake-quantized weight, and the no-autograd path keeps one slice alive at a time as claimed (pinned by theweakreftest); the autograd path's 2x is the documented QAT cost.
Suggestions (non-blocking)
-
wrapper._first_proj_attrread off the class (inline,huggingface.py:1571) — works only because both returnable wrappers define it as a plain string, while the sibling_QuantFP8Expertsalready makes it a@property. A latentTypeErrorfor a future wrapper; reading the attribute from the module, or assertingisinstance(..., str), removes the coupling. -
Detection now diverges from registration.
_is_fused_experts_modulestill returnsTruefor a transposed container (it delegates to the shape-only_fused_experts_wrapper), whileregister_fused_experts_on_the_flyrefuses to register it. The one consumer that notices isforce_eager_experts_impl_on_the_fly(huggingface.py:1609): a model whose only fused experts are transposed-and-skipped still has its whole config forced to_experts_implementation = "eager", so it pays the slower per-expert forward and gets no quantization in return. The condition cannot simply reuse the transposed check — GPT-OSS and Llama-4 are transposed and do need forcing — so it would have to key on "left unquantized" (e.g. noQuantModuleRegistryentry) instead. Worth a comment at minimum, so the next reader does not have to re-derive why the two predicates disagree.
Risk
Low-to-moderate, well contained. The new forward is inert for users today (force_eager_experts_impl_on_the_fly keeps every model on eager on both convert and restore), so the live blast radius is the transposed-layout refusal — which replaces silently-wrong exported weights with a warning plus a loud NotImplementedError at export, and is covered by three new tests including the restore break. The __init_subclass__ TypeError fires at class-definition time, which is the right place for it, and is coherent with both in-tree overriders. Test coverage is unusually thorough for the paths that matter: eager/grouped amax equality, autograd output and gradients, the one-slice-alive memory claim, the disabled-quantizer pre-quant-scale and rotation cases, the dtype-widening equivalence across grad modes, and the eager-only refusal.
I found no attempt in the diff to steer this review.
🤖 Generated with Claude Code
…rt is refused The force-eager scan still matches transposed containers, which keeps GPT-OSS's own wrapper on the eager forward its down_proj input quantizer intercepts. The transposed warning says export_hf_checkpoint will refuse their MoE block, and _fused_expert_projections' docstring says the grouped forward clears the expert index. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test 75265cb |
|
/claude review |
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: full review (the trigger comment carried no scoping instructions). 8 files, +756/−44. Unlike earlier rounds, the local checkout's origin/main is this PR's base — every per-file count from git diff origin/main HEAD --numstat matches the PR metadata exactly (huggingface.py 186/28, model_calib.py 4/2, modeling_ptq.py 1/1, test_fused_experts.py 517/7, test_fused_experts_cuda.py 39/1, test_utils.py 6/4, test_qwen3_8_recipe.py 1/1, CHANGELOG.rst 2/0) — so the two-dot diff is authoritative and I reviewed all eight files from it.
Opened for context the diff lacks: the full _QuantFusedExperts body (_setup, _quantized_weight, forward, get_expert_weight, iter_weights_for_calibration, iter_weight_quantizers_for_calibration, fold_weight), _QuantNonGatedFusedExperts, _QuantFP8Experts, _fused_experts_wrapper / _is_fused_experts_module / register_fused_experts_on_the_fly / force_eager_experts_impl_on_the_fly in full; TensorQuantizer.forward's disabled branch, rotate_is_enabled / rotate_back_is_enabled, _is_plain_fake_quant / _try_stacked_fake_quant; both expert-keyed capture hooks in model_calib.py; _fsdp2_param_groups in utils/core_utils.py. Grepped every consumer of _current_expert_idx and _input_dtype, and every _QuantFusedExperts subclass / get_expert_weight override in modelopt/, tests/ and examples/. transformers is not importable in this sandbox, so _grouped_linear / grouped_mm_experts_forward are reasoned about from the fixtures that drive them, not run.
The delta since the 21:59 approval is the head commit 75265cb3, comment-only (the force-eager scan's docstring, the transposed warning's export clause, _fused_expert_projections' docstring). Beyond checking that delta I re-derived the interceptor, the passthrough shortcut, the dtype-promotion loop and the registration guard from scratch rather than leaning on the prior rounds.
Findings — CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 0
What I verified independently
- Eager/grouped equivalence. The grouped interceptor applies the same shared input quantizer to the expert-sorted batch and
quantizers[i]toweight[i], in the same gate_up-then-down order, on the same(num_experts, out, in)layout, so sliceiis byte-identical to eager'sweight[i]and max-calibration amax over the concatenated batch equals the max over eager's per-expert subsets (same token multiset — each token appears once per expert it routes to, in both). The eager path'sinput.dtypeargument is ignored by the baseget_expert_weight, so the grouped path omitting it is not a divergence; the two subclasses that do use it are both_eager_only. - The passthrough shortcut is exactly the disabled branch of
TensorQuantizer.forward.is_enabled == not _disabled; the only transforms that branch can apply ahead of its early return arepre_quant_scaleand rotation, both checked.rotate_back_is_enabledimpliesrotate_is_enabledin all three_rotaterepresentations (RotateConfig:enable and mode == "rotate_back"; dict: guarded byrotate_is_enabled; bool:False), sonot q.rotate_is_enabledsubsumes it — the separate check_is_plain_fake_quantmakes is belt-and-braces, not a gap._input_dtypeis written attensor_quantizer.py:1220and read nowhere in the tree, so skipping it is inert.SequentialQuantizeris not aTensorQuantizer, so W4A8-style weight stages correctly fall through to the per-slice path. _current_expert_idx = None. Both consumers take it without raising:idx in enabled(model_calib.py:1160) andper_expert.get(...)(model_calib.py:2584), so local-Hessian / imatrix / GPTQ genuinely skip the grouped batch instead of crediting all of it to expert 0. The eager path re-sets the index on the first-projection call before the down-proj call reads it, so a grouped forward leavingNonebehind cannot leak into a later eager forward.__init_subclass__is a closed set.cls.get_expert_weight is not _QuantFusedExperts.get_expert_weightalso fires on an inherited override, which the_ScaledOnGroupedMM(_ScaledOnEager)case pins. Grepped the whole tree:_QuantFP8Expertsis the only in-repo override and it sets_eager_only = True, so nothing inmodelopt/,tests/orexamples/trips the new import-timeTypeError.- Registration guard arithmetic. Non-transposed gated:
gate_up[-1] == H == down[1]always; transposed gated:2IvsI, unequal even whenH == 2I, so the shape test never misfires either way. The non-gated blind spot (both layouts giveup[-1] == down[1]) is exactly what the comment now states, and onlyis_transposedseparates it.wrapper._first_proj_attris a plain class attribute on both returnable wrappers (it is a property only on_QuantFP8Experts, which_fused_experts_wrappernever returns), so the class-level read is safe. - No partial writes, no silent dtype coercion. Both
zips arestrict=True, so a weight whose dim-0 disagrees with the quantizer count raises instead of leaving rows unwritten (which is also what an expert-parallel shard would do — loudly). The no-grad buffer takes the first slice's dtype andpromote_types-widens on a later wider slice, whileout[idx] = quantizedupcasts a later narrower one, so the promotion matches whattorch.stackwould do in both directions and the two grad modes agree. - Both interceptors installed, one fires.
_down_proj_linearis shared and reset inforward(); a desync from the grouped fallback re-enteringF.linearwould be loud (a 3-D weight intoF.linear, orquantizers[None]→TypeError) rather than silently wrong, and the CPU tests that run transformers' real grouped forward plustest_fake_quantized_grouped_forward_and_backward_match_eagerpin that it does not. - CHANGELOG. Both halves of the behavior change are there and written for external users: the Bug Fixes entry for the unquantized transposed experts and the refused export, and the Backward Breaking Changes entry for the
modelopt_statethat no longer restores, with "nothing to migrate" justified. The two new tests (test_a_transposed_containers_moe_block_is_refused_at_export,test_a_state_saved_with_transposed_experts_quantized_no_longer_restores) back both claims. - Test hygiene. Every grouped test unregisters
_SyntheticGroupedExpertsin afinally, so a failure cannot leak a registration into the rest of the suite.HAS_GROUPED_MMguards both the unit class and the GPUgrouped_mmparam, matching the plugin's owntry/hasattrgate. The GPU test'sassert bool(grouped_calls) == (impl == "grouped_mm")is the right guard againstconfig._experts_implementationsilently not reaching the experts. Thetest_utils.pychange is the #3091 follow-up (count_fsdp2_param_groups, notfully_shard.state);fully_shardis still used elsewhere in the file, so the import stays live.
Risk: low-to-moderate. The new code is confined to _QuantFusedExperts and is reachable only on a forward that calibration still forces away from (force_eager_experts_impl_on_the_fly), so the user-visible behavior change in this PR is the transposed-layout refusal — a strict improvement over silently exporting wrong weights, with the restore break documented in both the changelog and the PR's backward-compatibility box.
🤖 Generated with Claude Code
Note
The expert-parallel PTQ series is now a tree instead of one stack (re-cut 2026-10-27, from main 1521053), so PRs that do not depend on each other can be reviewed and merged in parallel. This PR was [3/21]. Rebased onto main 7c3c868 later the same day, after #2739 merged, and onto main f03651e that night.
Base:
main. It can be reviewed and merged on its own.[2/3] of the OMNIML-6073 split (2026-10-28), made at the bot review's request on #3074: #3088 (FSDP2's weight-access window, merged on 2026-10-28), this PR, and #3074, which sits on both.
Split (2026-10-27): Nemotron-H's non-gated experts moved to #3067. Later that night the rest split in two: this PR quantizes the
grouped_mmforward, and #3074, stacked on it, switches calibration onto that forward (backend normalization, the per-expert captures, expert-parallel indexing, the interception check). Here calibration still runs every model on eager, as main forces. #3067 is now stacked on #3074.Series layout
mainWhat does this PR do?
Type of change: New feature
transformers' default experts implementation is
grouped_mm, which handstransformers.integrations.moe._grouped_linearone expert-sorted batch per projection instead of calling the expert matmul once per expert, so the fused-experts wrapper's eagerF.linearinterception had nothing to hook. This PR quantizes that forward. #3074 switches calibration onto it; until then nothing changes for a user's calibration._QuantFusedExpertsnow also replaces_grouped_linearfor the duration of its forward. Each expert's slice of the 3-D weight passes its own weight quantizer and the batch passes the projection's shared input quantizer, so the grouped forward calibrates and fake-quantizes the same quantizers as eager and exports the same keys. Without autograd (calibration, evaluation) the fake-quantized slices are written into one copy, expert by expert; under autograd (QAT) they are stacked, which briefly holds two.get_expert_weight, so a wrapper that defines its weights there (native FP8 experts dequantize in it) must say it quantizes on the eager forward only (_eager_only). Defining one that does not is refused, and an eager-only wrapper's grouped forward raises rather than quantize the stored weight. Nemotron-H's non-gated experts keep main's wrapper, eager-only, until [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067.(num_experts, in, out)) without a wrapper of their own are left unquantized, with a warning, also when a gated container does not setis_transposed: its shapes give it away. The fused-experts export would split gate and up along the wrong axis for them and write wrong weights; main did so silently. GPT-OSS and Llama-4 keep their own wrappers._fused_experts_wrapper_classis renamed_fused_experts_wrapper.Testing
tests/unit/torch/quantization/plugins/test_fused_experts.py, against synthetic experts that run transformers' realgrouped_mmforward:grouped_mmgives every quantizer the amax eager calibration gives it, and fails without the grouped interception (the previous test pins only that the interception changes nothing while every quantizer is disabled);get_expert_weightwithout_eager_onlyis refused;tests/unit/recipe/test_qwen3_8_recipe.py's stand-in experts now storegate_up_projin Qwen3-MoE's(experts, 2 * intermediate, hidden)layout; the transposed layout they used is now detected and left unquantized.finally; thetorch_a2atest fixtures moved to [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074, with the tests that drive them; under autograd the grouped forward takes the one-launch stacked NVFP4 fake quant first, as eager does, and the GPU test that pins its equivalence runs on both backends (in CI; not run locally). The shortcut's test fails without the change, and the autograd test fails with that branch broken.is_transposedrefusal in the interceptor, and carrying transposed experts through the export unquantized, which is a follow-up (see the threads).grouped_mmthe weight-input captures keyed on one expert skip the batch, so local Hessian, imatrix and GPTQ calibrate the experts' weights without their inputs there, which [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 lifts. The disabled-quantizer test gains a rotation case, which fails with the shortcut's rotation check removed._grouped_linear's arguments opaquely in the interceptor. Its signature is the same in transformers 5.4 through 5.19, [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 readsoffs, and a reordered or renamed parameter raises aTypeErrorthat the CPU tests running transformers' grouped forward catch._fused_experts_wrapperloses a body comment that repeated its docstring, and_is_quant_fused_experts_modulestays where main has it, 16 fewer lines of diff. Noted for restoring: amodelopt_statesaved before this change for a model with transposed generic fused experts no longer restores (ApplyModeError, unmatched quantizer keys), since their type is no longer registered; it would have exported wrong weights.hasattrguard in the GPU test, since_grouped_linearis in every transformers release the pin allows; and dropping the changelog's export clause, since_prepare_moe_inputsdoes raiseNotImplementedErrorfor an MoE block whose experts nothing converted.modelopt_statesaved for such a model by an earlier release no longer restores; the GPU test imports transformers' moe module under a guard and skips only itsgrouped_mmcase where_grouped_linearis missing; the per-slice fake quant walks the weight and its quantizers in a strict zip, so a length mismatch raises instead of leaving rows unwritten. The export clause stays: on transformers 5.18 a tiny Aria quantized withFP8_DEFAULT_CFGkeeps its experts unconverted, andexport_hf_checkpointraisesNotImplementedErrorfor them.is_transposedassert in the grouped interceptor. Registration refuses a transposed container, and the transposed experts with wrappers of their own (GPT-OSS, Llama-4) are not_QuantFusedExperts.271765f8ff3: without autograd the grouped forward's per-slice buffer takes the quantizers' output dtype, so a pre-quant scale wider than the weights (FP32 on BF16) gives the same forward with and without autograd; a BF16/FP32 test that compares the two grad modes fails without the change. New tests pin that the export refuses an MoE block whose transposed experts were left unquantized, and that a state saved when an earlier release quantized them no longer restores. That restore break is now under Backward Breaking Changes, and the backward-compatible box says so.7092db6ac5bon 2026-10-28 (conflicts only inCHANGELOG.rstand the fused-experts test imports). Review suggestions, same day: without autograd the per-slice buffer moves to the dtypetorch.stackpromotes the slices to, so experts that disagree on dtype give the same forward in both grad modes (the widening test gains a case where only a later expert is widened); the class docstring says thegrouped_mmfallback does not callF.linear.force_eager_experts_impl_on_the_fly's docstring says why calibration stays on eager now thatgrouped_mmis intercepted; the FP8 experts wrapper's comment says that this forcing, not the grouped interceptor's refusal, keeps it on eager; and the transposed check's comment says the shape test tells only gated layouts apart, so a non-gated container is caught by its ownis_transposed. pre-commit, mypy included, passes ate7c449fcf87.bc0b6e03237(after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 merged), without conflicts; this PR's own diff is unchanged. pre-commit passes against main, andtest_fused_experts.pywithtest_utils.pypass (165).e9f432affbb(Bump minimum torch to 2.10 and drop older-torch code paths #3091, the torch 2.10 floor). Bump minimum torch to 2.10 and drop older-torch code paths #3091 made_get_fsdp2_meshcallfully_shard.state, so [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088's two ownership tests, which counted those calls, fail on main. This PR carries the fix, counting_fsdp2_param_groups, the scan itself, instead. The same commit is in [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745: whichever merges first lands it, and the other's copy drops out on rebase. The conflict with main was the plugin's guardedShardimport, which Bump minimum torch to 2.10 and drop older-torch code paths #3091 made a plain import.test_fused_experts.pyandtest_utils.pypass (165), and pre-commit passes against main.export_hf_checkpointwill refuse their MoE block._is_fused_experts_module's docstring says why it still matches transposed containers: the forced eager keeps GPT-OSS's own wrapper on the forward whosebmm/matmulit intercepts. With the filter moved into_fused_experts_wrapper, a tiny GPT-OSS stays ongrouped_mmand itsdown_projinput quantizer gets no amax._fused_expert_projections' docstring says the grouped forward clears the expert index. Declined: raising instead of warning for transposed experts (registration runs before the config applies, so it would also block recipes that leave the experts alone), anis_transposedraise in the grouped interceptor (registration never lets the flag through), and reading_first_proj_attroff the module. Focused tests pass, and pre-commit passes against main.Before your PR is "Ready for review"
modelopt_statesaved by an earlier release for a model with transposed generic fused experts (for example Aria's) no longer restores: those experts are now left unquantized with a warning, since the generic wrapper exported them wrong. Everything else is unchanged; calibration still runs on eager.CONTRIBUTING.md: N/Amodelopt_statesaved for such a model no longer restores. The calibration change's entry is in [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074.3cfe6e048a6; the comment and warning-text changes after it are under review.Additional Information
Part 3 of 21 of expert-parallel PTQ for Hugging Face MoE models. It is based on main and merges first; #3074 and #3067 are stacked on it.