Repository navigation
[OMNIML-6073] [3/21] Quantize fused experts under grouped_mm, ungated and transposed layouts - #2745
Open
shengliangxu wants to merge 27 commits into
Open
shengliangxu wants to merge 27 commits into
shengliangxu wants to merge 27 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. |
Contributor
|
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 | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 58.06% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 93 functions across 8 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
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 identifies the main change: fused-expert quantization for grouped_mm, ungated, and transposed layouts. It is concise and specific. 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.
shengliangxu
added this pull request to stack #2755
October 12, 2026 00:57
shengliangxu
force-pushed
the
shengliangx/hfep-8-fused-experts
branch
from
October 12, 2026 02:35
6e6ac9c to
bad2de9
Compare
Contributor
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## shengliangx/hfep-7-amax-sync #2745 +/- ##
================================================================
- Coverage 79.13% 78.82% -0.31%
================================================================
Files 661 661
Lines 74977 75536 +559
================================================================
+ Hits 59332 59543 +211
- Misses 15645 15993 +348
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:
|
Collaborator
Author
|
/ok to test bad2de9 |
shengliangxu
force-pushed
the
shengliangx/hfep-8-fused-experts
branch
from
October 12, 2026 08:39
bad2de9 to
16e5c17
Compare
Collaborator
Author
|
/ok to test 16e5c17 |
shengliangxu
changed the title
shengliangxu
force-pushed
the
shengliangx/hfep-8-fused-experts
branch
2 times, most recently
from
October 13, 2026 01:32
5c0d5fa to
698f459
Compare
Collaborator
Author
|
/ok to test 698f459 |
shengliangxu
removed this pull request from stack #2755
October 13, 2026 17:42
shengliangxu
added this pull request to stack #2775
October 13, 2026 17:43
shengliangxu
changed the title
shengliangxu
force-pushed
the
shengliangx/hfep-8-fused-experts
branch
from
October 13, 2026 19:38
698f459 to
eaa41c2
Compare
shengliangxu
removed this pull request from stack #2775
October 13, 2026 19:38
shengliangxu
added this pull request to stack #2781
October 13, 2026 19:40
shengliangxu
changed the title
Collaborator
Author
|
/ok to test eaa41c2 |
shengliangxu
removed this pull request from stack #2781
October 14, 2026 00:13
shengliangxu
force-pushed
the
shengliangx/hfep-8-fused-experts
branch
from
October 14, 2026 00:14
eaa41c2 to
3ff1d97
Compare
shengliangxu
added this pull request to stack #2786
October 14, 2026 00:14
shengliangxu
changed the title
… experts Three ways the fused-experts changes broke models that worked on main: - GPT-OSS: its own wrapper quantizes down_proj's input by intercepting the eager forward's bmm/matmul. Main forced every fused-experts model onto eager; the series kept grouped_mm, where that interception never fires, so down_proj's input quantizer was never calibrated and its exported scale fell back to gate_up_proj's. A model with any fused-experts class that has its own (non-_QuantFusedExperts) wrapper keeps main's forced eager again. - The interception check raised for an experts module no calibration token reached, e.g. Ernie-4.5-VL's vision experts under text-only calibration, which main exported with weight-derived scales. It now counts the module's forwards and fails only when the forward ran past both hooks. - A transposed fused-experts class without its own wrapper made quantization raise. It is now left unquantized with a warning: still never exported wrong, without aborting the rest of the model. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
transformers 5.18's grouped_mm_experts_forward reads _is_expert_parallel, which its @use_experts_implementation sets on real experts classes. The test's synthetic experts call the forward directly, so they failed with an AttributeError on CI's latest transformers -- and the failures left state behind that broke later tests in the same session. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Under expert parallelism transformers keeps fused expert weights as DTensors sharded over the experts, with num_experts set to the local count. Outside the forward, max calibration's weight pass and fold_weight sliced them by local expert index, which on a DTensor reads the global expert; with the dense part FSDP2-sharded, FSDP2's weight-access window had first gathered every expert onto every rank. Either way every rank but rank 0 calibrated each expert as the max with rank 0's expert of the same local index, and folded into the wrong experts: expert weight scales came out too large. Index the local shard of an expert-sharded DTensor, and keep DTensors that do not live on FSDP's mesh out of FSDP2's weight-access window. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
With fake quantizers inserted, fused experts split over two ranks must train through the torch_a2a forward like the whole quantized module. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
main opened 0.49.0 for what merges after 0.48.0's cut, which this PR does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Each of them spawned two ranks, about 14 s of the linux unit job apiece. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The torch_a2a expert-parallel forward runs its expert matmuls through _grouped_linear as the grouped_mm backend does, but _running_grouped_mm recognized only grouped_mm. local_hessian therefore took the eager per-call hook, which never moved past expert 0: expert 0's Hessian got every expert's rows and the others fell back to plain MSE. The grouped hook also reads each weight the way the quantizers index it, rank-local: indexing an expert-parallel DTensor names a global expert and redistributes it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…r EP Hook registration for local Hessian and imatrix handed on_weight weight[0] of each fused-experts projection. Under expert parallelism the fused weight is a DTensor sharded over the experts, so indexing it named a global expert and all-gathered every expert of the projection to report one expert's shape, then paired that global expert with this rank's first quantizer. It now reads this rank's experts, as the grouped capture hook does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…eir imports A second copy of the Qwen3-VL section header, constants and helpers sat before the NemotronH tests, shadowed by the original block the Qwen3-VL tests use. The imports the expert-parallel and grouped tests made inside functions move to module scope: none is optional at the transformers 5.5 and torch 2.8 floors, and a missing name now fails at collection instead of inside a pooled worker. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
- The ungated synthetic classes and TestNemotronHUngatedExperts become a gated/ungated parameter of the grouped_mm passthrough and calibration tests; TestNonGatedFusedExperts keeps the detection and quantizer checks the Nemotron-H class repeated. - The three "leaves a supported value unchanged" tests become one, parametrized over eager, grouped_mm and None. - test_intercept_fires_twice_and_uses_full_3d_weight asserted on a spy's arguments; a missed or wrong interception fails the calibration and local-Hessian tests instead. - test_a_transposed_class_with_its_own_wrapper_is_left_to_it is covered by the real GPT-OSS calibration test. - The QAT-through-torch_a2a job becomes the fake-quantized case of the a2a file's whole-module backward test, and the FSDP-wrapped EP job the fsdp2-dense case of the EP-sharded calibration test. Each regression a dropped test caught fails a remaining test. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Main's native FP8 experts (#2820) route the forward's and calibration's per-expert weights through get_expert_weight, which indexed the fused weight directly. Under expert parallelism that weight is a DTensor sharded over the experts, so indexing it read the global expert: from rank 1 on, another rank's. Index this module's own experts, as the rest of the wrapper does, for the weight and for the FP8 experts' scales. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Main's GPTQ for HF fused-MoE experts feeds each expert's Hessian from the shared input quantizer, keyed by _current_expert_idx: one call per expert, as the eager loop runs. This stack calibrates fused experts on grouped_mm, where that quantizer runs once per projection on the whole expert-sorted batch, so the first expert took every token and the others none (a tiny Qwen3-MoE: expert 0 got all 64 rows, experts 1-3 none, so they fell back to the plain search). Split the batch by _grouped_offs, as the weight-input hooks already do, and note both dispatches in _fused_expert_projections' docstring. On that Qwen3-MoE every expert now gets the token counts eager calibration gives it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Under expert parallelism a fused experts' weight is a DTensor sharded over the experts. GPTQ built each expert's helper from weight[idx], where idx is the local index its quantizer has: that names a global expert and gathers every expert's weight, once per local expert. The helper only reads its expert's shape and device, so the values were right, but the gathers cost the memory and time of the whole expert weight each time. Index this rank's experts instead (_local_experts), as the weight-input hooks beside it already do. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…rride runs them The a2a transport left quantization out of its docs, since ModelOpt's quantized _grouped_linear override and the torch_a2a exception to eager experts arrive here. With them, the transport's local experts run through ModelOpt's per-expert quantizers over the tokens dispatched from every rank, and its docs say so again. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Its torch_a2a case gave the model a one-rank EP group and left it registered on the session's persistent cpu_dist_workers until the pool shut down; the job's cleanup now destroys it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The chunked-hops test ran only forward, under no_grad, so the slicing branch that runs while a graph is built, where the slices are concatenated rather than written in place, never ran in any test. The test now also runs forward and backward through hops sliced below the limit, and checks the output and the input gradients against the unsliced run, bit for bit. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The transport's comments said a hop past 2**31 elements overflows all_to_all_single's int32 element offsets and hangs. With the slicing disabled, such hops have since completed exactly on NCCL 2.29, 2.30 and 2.31, up to 32 ranks over InfiniBand, and no hang has been reproduced. The slicing stays, since it is exact and costs two extra hops only past 2**31; the comments now say what it guarantees, that every hop stays within the int32 element offsets. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The eager forward's weight-input hook indexed the fused weight by the active expert directly. Under expert parallelism that weight is a DTensor whose index names a global expert, so every rank past the first paired its quantizers with another rank's expert, and the read all-gathered every expert. The hook now indexes this rank's experts, as the grouped hook beside it does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
_QuantFP8Experts applies the checkpoint's block scales in get_expert_weight, which only the eager per-expert forward reaches; transformers' FP8 grouped_mm forward runs its own kernels, which nothing here intercepts. main forced every MoE onto eager, but now only wrappers outside _QuantFusedExperts are forced there, and FP8's is a subclass, so a native-FP8 MoE checkpoint stayed on grouped_mm and calibration found no expert it could quantize. A wrapper now says when it quantizes only on the eager forward, and the FP8 one does. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The local-Hessian, imatrix and GPTQ capture hooks pick the eager or the grouped split from config._experts_implementation, while the interceptors run whichever path the forward takes. A module whose forward ran the other path was captured silently wrong: the grouped hook returned on a missing offset list, and the eager one split nothing. The interceptors now record the path that ran (the grouped one sets _grouped_offs, the eager one clears it) and the hooks raise on a mismatch. The forward that raises leaves the quantized F.linear and _grouped_linear installed, because replace_function restored them only on a normal exit; it now restores them in a finally. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…thing reaches The on-the-fly registration leaves transposed (num_experts, in_dim, out_dim) experts unquantized, and the classes that store them so (GPT-OSS, Llama-4) have wrappers of their own, so _QuantFusedExperts never wraps a transposed module, as its docstring says. Its transposed branches in the grouped weight quantization, the calibration iterator and fold_weight could not run; they are gone, and fold_weight's docstring no longer claims to reuse iter_weights_for_calibration. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Without autograd, _quantize_grouped_weight writes each expert's fake-quantized slice into a preallocated output instead of stacking a list of them, so a projection's slices are no longer all alive beside the stacked copy. Under autograd (QAT) it still stacks: slices written into one buffer chain one CopySlices node per expert, each copying the whole gradient in the backward. It takes this rank's experts before the passthrough return, so both branches hand _grouped_linear the same kind of tensor. normalize_experts_impl_on_the_fly refuses, before it changes any config, a model whose experts quantize only on the eager forward (native FP8 experts, a class with its own wrapper such as GPT-OSS's) when it runs them on torch_a2a, which cannot fall back to eager: they would calibrate on raw FP8 bytes, or not at all. register_fused_experts_on_the_fly also leaves a gated container unquantized when its shapes say it is stored transposed, whether or not it sets is_transposed: stored (num_experts, out, in), the first projection takes the dim down_proj puts out. calibrate() zeroes the fused-experts interception counters before the algorithm runs the forward loop, so each calibration judges its own forwards, and raises the interception error only once the model is back in its training mode. fold_weight goes back to folding the pairs iter_weights_for_calibration yields through _fold_weight_quantizer, now that get_expert_weight reads this rank's experts. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The recipe test's stand-in experts stored gate_up_proj as (experts, hidden, 2 * intermediate), the transposed layout, which registration now detects from the shapes and leaves unquantized, so the test failed looking for the experts' weight quantizers. Qwen3-MoE's experts store it as (experts, 2 * intermediate, hidden), as the other recipe tests' stand-ins already do. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Collaborator
Author
|
/ok to test 9913069 |
Collaborator
Author
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Comment on lines
+1530
to
+1535
| if wrapper is not None and ( | ||
| getattr(module, "is_transposed", False) | ||
| # Stored `(num_experts, out_dim, in_dim)`, the first projection takes the dim down_proj | ||
| # puts out; a gated container stored transposed never does, whether it says so or not. | ||
| or getattr(module, wrapper._first_proj_attr).shape[-1] != module.down_proj.shape[1] | ||
| ): |
Contributor
There was a problem hiding this comment.
[SUGGESTION] The shape half of this transposed guard can only ever fire for the gated layout; the ungated one slips through.
For ungated experts the two layouts are shape-symmetric:
| layout | up_proj |
down_proj |
up.shape[-1] |
down.shape[1] |
|---|---|---|---|---|
normal (E, out, in) |
(E, I, H) |
(E, H, I) |
H |
H → equal |
transposed (E, in, out) |
(E, H, I) |
(E, I, H) |
I |
I → also equal |
So a transposed ungated container is not detected by shape, only by an explicit is_transposed attribute. If such a class ever appears without that flag, _QuantNemotronHExperts claims it and every per-expert weight amax/scale is reduced over the input dim instead of the output dim — silently wrong scales in the export, which is exactly the failure this guard was added to prevent (the gate/up split doesn't apply to ungated, but the scale axis does).
The code comment is correctly scoped ("a gated container stored transposed never does"), but the condition itself is applied to both wrappers, so the asymmetry isn't visible at the call site. Worth either restricting the shape test to wrapper._is_gated and saying in the comment that the ungated layout relies on is_transposed alone, or finding a shape invariant that works for both (e.g. for gated, first.shape[1] == 2 * down.shape[2]).
Low severity — all transposed classes in transformers today (GptOssExperts, Llama4TextExperts) set is_transposed and are short-circuited by the QuantModuleRegistry.get(mod_type) is not None check above anyway.
Comment on lines
+572
to
+580
| dim_names = original_device_mesh.mesh_dim_names | ||
| if ( | ||
| dim_names is not None | ||
| and fsdp_device_mesh.mesh_dim_names is not None | ||
| and dim_names[:fsdp_dim] != fsdp_device_mesh.mesh_dim_names | ||
| ): | ||
| # Not FSDP's shard: a DTensor on another mesh, such as an expert-parallel weight that | ||
| # fully_shard was told to ignore. Gathering it would hand every rank every expert, which | ||
| # its module indexes by local expert. |
Contributor
There was a problem hiding this comment.
[SUGGESTION] This guard is a no-op whenever either mesh has unnamed dims, which lets the exact bug it fixes back through.
DeviceMesh.mesh_dim_names is None for a mesh built without mesh_dim_names= (plain DeviceMesh(...) / init_device_mesh without names). In that case both is not None conditions short-circuit and the EP DTensor is not skipped:
- if
fsdp_dim == original_device_mesh.ndim, the assert below is also skipped and the EP weight getsredistributed toReplicate()— every rank ends up with every expert, which_local_expertsthen hands to a module that indexes by local expert. Silently wrong scales, the failure mode this PR is fixing. - if
fsdp_dim != original_device_mesh.ndim, the next assert doesoriginal_device_mesh.mesh_dim_names[:fsdp_dim]onNone→TypeErrorinstead of a useful message.
Transformers' EP path does name its mesh dims, so this is latent rather than live. Still, a positive identity test would be robust regardless of naming, e.g. skip unless the param's mesh is (or starts with) the FSDP mesh:
original_device_mesh = param.device_mesh
if original_device_mesh[: fsdp_device_mesh.ndim] != fsdp_device_mesh:
# Not FSDP's shard: a DTensor on another mesh, such as an expert-parallel weight that
# fully_shard was told to ignore. Gathering it would hand every rank every expert,
# which its module indexes by local expert.
continue(DeviceMesh.__eq__ compares the underlying mesh tensor and dim names, and slicing an unnamed mesh works.) If you'd rather keep the name comparison, at least turn an unnamed mesh into an explicit error rather than letting it fall through.
Comment on lines
+923
to
+927
| # Each slice goes into the output as it is made: a stack would hold all of them beside it. | ||
| out = torch.empty_like(weight) | ||
| for idx in range(self.num_experts): | ||
| out[idx] = quantizers[idx](weight[idx]) | ||
| return out |
Contributor
There was a problem hiding this comment.
[SUGGESTION] The changelog entry for this backend switch mentions only the activation-scale drift, not the peak-memory step it also brings.
On the eager path only one expert slice (out, in) is fake-quantized at a time. On the grouped path — now the default for every fused-MoE model — each _grouped_linear call materializes a full extra copy of the projection's fused 3-D weight, alive across _orig_grouped_linear, so it coexists with the grouped GEMM's own allocation. For Qwen3-30B-A3B that is ~0.8 GiB per gate_up_proj (128 × 1536 × 2048 × 2 B), more on larger MoEs; the autograd branch briefly holds two. That is intrinsic to grouped_mm (the GEMM needs the whole weight), so there's nothing to fix in the code here — the comments already explain the two branches well.
What's worth adding is the user-facing half: someone calibrating near the memory ceiling on the previous release can OOM after upgrading with no code change, and the changelog currently points them only at scale differences. Suggest extending the CHANGELOG.rst entry with a clause along the lines of "and raises calibration peak memory by about one fused expert-weight copy per projection", so experts_implementation="eager" reads as the remedy for both symptoms.
Contributor
Claude review — [3/21] Quantize fused experts under grouped_mm, ungated and transposed layoutsScope: full review (the trigger comment carried no scoping instructions). All 7 changed Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 3
All three are latent or documentation-level. None blocks merge. What I traced and found correct
One thing outside this diff worth a follow-up
RiskLow for the code, moderate for users by design. The implementation is careful and the loud-failure machinery is a real improvement over the silent miscalibration on main. The declared risk is the deliberate default change (eager -> |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: New feature, with a backward-incompatible change to MoE calibration (see the changelog entry)
The fused-experts wrapper handled one shape of MoE: gated experts (
gate_up_proj/down_proj) in the(num_experts, out, in)layout, calibrated through the eager backend, where Hugging Face calls the expert matmul once per expert. Several real architectures are not that shape. Each failed differently, usually by silently collecting no amax rather than by raising.grouped_mm, not eager, and it passes one expert-sorted batch per projection instead of one call per expert, so the eager interception had nothing to hook. The wrapper now intercepts the grouped path and keeps the cumulative token offsets, and the local-Hessian capture splits the batch back into per-expert chunks with them. Backends that cannot be intercepted (batched_mm, the opaquesonicmoekernel) are moved togrouped_mmbefore calibration rather than quietly producing an uncalibrated model.up_projwith no gate/up split. The first projection's name is now a class attribute the subclass overrides, the quantizer names derive from it, and_is_gated = Falsetells the export not to split it.(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.fold_weightindexed them by local position too, which on a DTensor reads the global expert; with the dense part FSDP2-sharded, FSDP2's weight-access window had first gathered every expert onto every rank.calibrate()raises when a fused-experts module's forward ran but neither hook fired. Eachcalibrate()judges its own forwards (it resets the interception counters first) and raises with the model back in its training mode. Every failure above otherwise looks the same: an export that seems complete and is quietly unquantized. An experts module that no calibration token reached is not an error and falls back to weight-derived scales as on main: Ernie-4.5-VL's vision experts, under text-only calibration. A model whose experts quantize only on the eager forward (native FP8 experts, GPT-OSS) is refused on thetorch_a2aexpert-parallel forward, before any config changes, instead of calibrating on raw FP8 bytes or not at all.Behavior change. main forced MoE calibration onto eager; transformers now resolves an unrequested backend to
grouped_mm, so fused-experts models calibrate there. A model whose experts class has its own quantized wrapper keeps main's forced eager. GPT-OSS's wrapper, for instance, quantizesdown_proj's input by intercepting the eager forward's matmuls. A single-GPU Qwen3-30B-A3B export matches main in every weight and weight scale, but 4,647 activation scales move, by up to 10.2%. Loading withexperts_implementation="eager"reproduces main's checkpoint byte for byte (all 74,835 tensors for NVFP4 and 56,211 for FP8). The changelog says so. On the grouped path each forward fake-quantizes a projection's experts together: without autograd (calibration, evaluation) it writes them into one copy, expert by expert; under autograd (QAT) they are stacked, which briefly holds two.Weight-input captures check the path the experts ran on: one fed by the other path than the experts' config selects is refused rather than attributed to the wrong experts, and native FP8 experts stay on the eager forward, where their dequantization runs. Under expert parallelism the eager capture reads this rank's own experts.
Usage
Testing
tests/unit/torch/quantization/plugins/test_fused_experts.py(+611 lines), against synthetic and real transformers experts modules:down_proj's input on the eager forward.The last two, and the transposed warning, fail on the code before those fixes.
Expert parallelism, on CPU under gloo:
torch_a2aforward like the whole quantized module: the fake-quantized case oftest_moe_a2a.py's whole-module test (mutation-checked twice);test_moe_a2a.py's chunked-hops test now also runs forward and backward through hops sliced below the limit, and checks the output and input gradients bit for bit against the unsliced run. Only the forward slicing ran in a test before; zeroing a slice in the graph path fails the new case.The GPU tests of QAT under expert parallelism, which this PR carried before the series was reordered, are in [10/21], the first PR with an FSDP2-sharded dense part they also run.
Against main, on tiny real models:
down_proj's input amax equal to main's;Local-Hessian calibration captures each expert's own rows under
grouped_mmand under thetorch_a2aexpert-parallel forward ([1/21]); both feed the quantizers one expert-sorted batch.Registering those input hooks under expert parallelism reads each rank's own experts: indexing the expert-sharded weight all-gathered every expert of the projection. A two-rank test checks that registration issues no collective.
The transport's ([1/21], merged) docs say again that ModelOpt quantizes its local experts, now that this PR's
_grouped_linearoverride does. Its hop-slicing comments now say what the slicing guarantees, every hop withinall_to_all_single's int32 element offsets, rather than that it prevents a hang: with the slicing disabled, hops past 2**31 elements completed exactly on NCCL 2.29, 2.30 and 2.31, including at 32 ranks over the network on 2.29 and 2.30.Near-duplicate tests are merged into parametrized ones, tests a remaining test already covers are dropped, and the two expert-parallel jobs become cases of the tests they repeated. Each drop is checked by mutation.
tests/unit/torch/quantization/plugins/test_fused_experts.pyalso pins the eager capture of an expert-parallel layer reading the local expert's weight, FP8 experts forced onto the eager forward, and a capture fed by the other path refused, with the patched functions restored. It also pins one fake-quantized expert slice alive at a time without autograd, thetorch_a2arefusal of eager-only experts with the configs left unchanged, transposed experts detected by their shapes, and a later calibration that bypasses both hooks raising with the training mode restored; each fails with its fix reverted.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.Locally on this revision, the stack rebased on main 0638a27: pre-commit over the whole PR against main, at every PR of the stack; each fix's own tests at its PR; and at the stack's tip the unit tests of the changed areas (
tests/unit/torch/utils,export,models, the fused-experts quantization tests and the hf_ptq example tests) on transformers 5.18 and 5.5 and on torch 2.8. They pass apart from the diffusers-export and GGML tests, which fail the same way on main in that environment. CI runs the full unit suite.The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.
Byte-for-byte checkpoint comparison against main on one GPU: Qwen3-30B-A3B, eager and grouped_mm, NVFP4 and FP8. Qwen2.5-0.5B (dense) is byte-identical.
Before your PR is "Ready for review"
grouped_mm, so some activation scales differ from earlier releases;experts_implementation="eager"restores the earlier checkpoint exactly.CONTRIBUTING.md: N/AAdditional Information
Part 3 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [2/21].
Summary by CodeRabbit
Compatibility
experts_implementation="eager"to restore the previous calibration setup.Bug Fixes