Repository navigation
[OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager - #3074
Conversation
|
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.
|
| Layer / File(s) | Summary |
|---|---|
Backend normalization and interception checks modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/quantization/model_quant.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, CHANGELOG.rst |
Backend normalization preserves supported implementations and redirects unsupported ones according to grouped-linear availability. Calibration tracks whether its forward loop ran and checks interception for applicable fused experts. Tests cover normalization and missed interception. |
Grouped execution capture and GPTQ modelopt/torch/quantization/model_calib.py, modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/utils/moe.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, tests/unit/torch/utils/test_moe_a2a.py |
Capture hooks split grouped batches by expert offsets for input calibration and GPTQ. Tests cover grouped capture, Hessian behavior, fake quantization, and transport gradients. MoE documentation describes quantization of the grouped linear operation. |
Rank-local weights and distributed validation modelopt/torch/quantization/plugins/huggingface.py, modelopt/torch/quantization/model_calib.py, modelopt/torch/quantization/utils/core_utils.py, tests/unit/torch/quantization/plugins/test_fused_experts.py |
Expert-weight lookup and FP8 scale access use local shards. FSDP2 writeback skips DTensor parameters not owned by FSDP units. Tests cover expert-parallel calibration, folding, hooks, and GPTQ helpers. |
Estimated code review effort: 4 (Complex) | ~45 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant Calibration
participant FusedExperts
participant GroupedLinearHook
participant CaptureHooks
Calibration->>FusedExperts: run calibration forward
FusedExperts->>GroupedLinearHook: pass grouped inputs and expert offsets
GroupedLinearHook->>CaptureHooks: split inputs by expert offsets
CaptureHooks->>Calibration: collect per-expert inputs or Hessians
Suggested reviewers: cjluo-nv, chadvoegele
Merge Risk: 
Minimal · previous review f1a1b
-
No actionable issue was identified in the reviewed calibration and expert-parallel paths. The PR is mergeable after normal checks.
Pre-merge checks: 
5
-
Full details: Security Anti-Patterns
-
Explanation
No security anti-patterns were introduced. The reviewed diff changes no example Python files and no pyproject.toml or requirements.txt files. Added-line searches found no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval/exec on input, or # nosec comments. The only eval-related addition is the safe model.eval() call in a test.
-
Finishing Touches: 💡 1
-
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
📝 Generate docstrings
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Commit to this branch
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/ok to test e5a7530 |
calibrate() imported the HF plugin for its interception check and caught ImportError only. A plugin whose body fails otherwise (a transformers version skew raising AttributeError, say) is one import_plugin only warns about, and the import reran it and failed every calibration with that error. The check now reads the plugin from sys.modules: one that did not load installed no fused-experts wrappers to check. A test breaks the plugin's import and calibrates; it fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The weight-input captures split an expert-sorted batch by the grouped offsets, read as exclusive cumulative ends, one per expert, as transformers' grouped_mm forward and a2a_moe build them. Offsets of another form (a leading zero, too few entries, ends past the rows) would credit experts with other experts' rows without a word; the split now raises. A test feeds each form and fails without the check. The _grouped_offs comment names every capture that reads it, and its other job: telling a capture which forward ran. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The offsets check covered the count of ends and the last one, but ends out of order passed it and credited an expert with another expert's rows. It now also requires the ends never to decrease, from zero, and the test gains that case, which fails without it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A first end of zero is fine (expert 0 got no rows); what the check refuses is the inclusive [0, c0, c0 + c1, ...] form, one end too many. The docstring and the test case's id now say so. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Calibration now runs fused experts on the backend transformers picks, grouped_mm by default, and the way back to earlier releases' scales is loading with experts_implementation="eager", which hf_ptq.py had no way to pass. The flag reaches every load path: get_model (the compressed- tensors branch included), the low-memory load, and, through the config the meta model is built from, the FSDP2 one. A test loads a tiny Qwen3-MoE with eager and fails without the plumbing; the changelog entry names the flag. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
calibrate() now raises when a fused-experts forward ran but reached neither quantizer hook, where main exported weight-derived input scales with a warning: a script that used to produce a checkpoint now stops. The Backward Breaking Changes section says so. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…ly when given The FSDP2 load builds its meta model from a pre-fetched config, so the flag is set on that config's _experts_implementation, as build_meta_causal_lm sets a pre-fetched config's attn_implementation, rather than through an AutoConfig kwarg. A test captures the config the FSDP2 loader receives and fails without it. get_model passes the flag on its compressed-tensors branch only when given, as on every other load, instead of an explicit None. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
normalize_experts_impl_on_the_fly rewrote the experts backend without a word. On main it always forced eager, but a caller can now ask for one (a load's experts_implementation, hf_ptq's --experts_implementation), and an unintercepted backend redirected to grouped_mm, or experts that quantize on the eager forward only, override that choice. Each change now warns once, naming both backends and why; a supported value, the common case, stays quiet. The redirect test expects the warning and fails without it. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
None already dispatches to the eager F.linear forward, so moving it to eager for experts that quantize on the eager forward only changes nothing; it warned anyway, for every such model, through the experts' own config too. The FP8 test now covers a grouped_mm and a None backend, and counts the warnings: one, then none. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Every rank redirects its experts alike, so an expert-parallel run printed one copy of the notice per rank. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Transformers' _experts_implementation setter carries the value to the sub-configs, so the one assignment on the pre-fetched config also reaches the text config Qwen3.5-MoE's experts dispatch on. The FSDP2 test now runs a VLM config too and checks its text config; it fails if the value reaches the root config only. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A class no backend change redirects needs its own quantized wrapper; meanwhile the rest of the model quantizes with these modules' quantizers disabled, which the check honours. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…tation load_model now reads args.experts_implementation; the test's stand-in for hf_ptq's parsed arguments lacked it, so both of its cases failed with an AttributeError in the trtllm hf_ptq lane. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
A typo in the backend name surfaced only inside transformers' model load. Quantization redirects every backend but eager and grouped_mm to one of them, so those two are the choices. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The fused-experts forward reads the routing indices by position to tell the forwards that route rows to this rank's experts. A test now reads every experts class transformers decorates with use_experts_implementation and checks that its forward takes them right after the hidden states. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The check was rank-local, so under expert parallelism a rank whose experts saw no rows did not raise and waited in its next collective while the others raised. Each experts module's verdict is now reduced over the group its experts are sharded over, which every member already joins in the experts' forward: one all-reduce per group, issued whatever the rank found, and only when experts are sharded that way. A rank that raises for another rank's experts says so. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The expert-parallel group is read from the experts' weight, a DTensor sharded over the experts under transformers 5.16 and later. Earlier versions' expert parallelism keeps each rank's experts as a plain tensor, which names no group, so the check stays rank-local there, as it was. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Under expert parallelism the check all-reduces over the experts' group, so whether calibrate() runs it (the forward loop ran, the HF plugin is loaded) must come out the same on every rank, or the ranks that run it wait for the ones that did not. A comment at the gate says so. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
load_model's FSDP2 branch now reads args.experts_implementation; this 2-GPU test's stand-in for hf_ptq's parsed arguments lacked it, so it would fail with an AttributeError once run (the nightly runs it; the 1-GPU PR runners skip it). Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…atistics During calibration the weight quantizers collect statistics without quantizing, so each hands its slice back unchanged. The grouped forward still built its result slice by slice into a new weight-sized buffer: a copy of every expert's weight per projection on each calibration forward, now that calibration runs on grouped_mm by default. Such quantizers now collect per slice and the weight itself is returned, as for disabled ones. A test checks the result is the weight and each expert's statistics are collected; it fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…r restore The backend normalization runs at convert and at restore, so after mtq.quantize or mto.restore the fused experts keep grouped_mm for every forward, not only calibration: QAT/QAD training, AutoQuantize scoring and fake-quant evaluation too. The changelog entry now says so, and names experts_implementation="eager" as the way back for all of them. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
…needs _quantize_grouped_weight is reached only from the grouped interceptor, inside the experts forward. Under expert parallelism transformers swaps each DTensor expert parameter for its local shard around that forward (the torch_a2a experts style asks for local tensors), and before 5.16 the shards were plain tensors, so the weight it receives is already this rank's experts and _local_experts returned it unchanged on every path. The docstring now says why no indexing is needed. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The pass-through called every enabled weight quantizer on its slice, but only one that calibrates does anything with it: enable_stats_collection leaves a quantizer with a constant amax neither quantizing nor calibrating, so its calls were no-ops. The loop now calls only the calibrating ones. The calibration test gains a constant-amax case that calls no quantizer and fails without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The interception check counts an eager forward only when a row routes to this module's experts, but a grouped one (grouped_mm, torch_a2a) whatever its routing indices hold: those forwards call _grouped_linear whatever rows they are handed, and torch_a2a's indices are global expert ids. A test runs a grouped-configured forward that reaches neither hook with every index past the module's experts, and checks calibration still raises; without the grouped term the forward goes uncounted. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The experts forward clears the grouped offsets when it returns, so a weight-input capture reached outside a forward does not split its input by a finished forward's offsets. The capture test now calls an input quantizer directly after a grouped forward and checks that nothing is captured; without the clear, it splits by the stale offsets and raises. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
The test read each transformers modeling file with the platform's default encoding, which is cp1252 on Windows; 16 of them, such as modeling_altclip.py, hold characters cp1252 cannot decode, so the test failed in the windows unit job. It now reads them as UTF-8. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Under transformers' own expert parallelism the grouped forward sorts the rows routed to other ranks' experts past offs[-1], where the grouped GEMM neither reads nor writes. Before transformers 5.17 nothing zeroes the up projection's output there, so the down projection's input holds those rows uninitialized, and the shared input quantizer calibrated on them: an input scale from garbage, silently when the garbage is finite. While an input quantizer collects statistics, it now sees only the rows before offs[-1]; the tail passes to the GEMM unread. The host sync this needs is paid during calibration only. A test hands the grouped projections a NaN or a large finite tail, as pre-5.17 expert parallelism leaves it, and checks the input amaxes match a run without one; both cases fail without the change. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
|
/ok to test a6758d7 |
Note
[2/2]: #3122 ([1/2]) is this PR's expert-parallel correctness fix, split out on 2026-10-29 at a reviewer's request so that it lands on its own, with a Bug Fixes changelog entry. This PR keeps the two backward-breaking changes, calibration on
grouped_mmby default and the newcalibrate()error, and sat on #3122 until #3122 merged on 2026-10-30. Before that, it was [3/3] of the OMNIML-6073 split, after #3088 (FSDP2's weight-access window) and #2745 (grouped_mmquantization), both merged. #3067 (Nemotron-H) sits on this PR.Base:
main. #3122 merged on 2026-10-30, so this PR sits on main, its own diff unchanged.In the tree: #2741's helper base merges #3067, which carries this PR and #2745.
Series layout
mainWhat does this PR do?
Type of change: New feature, with a backward-incompatible change to MoE calibration (see the changelog entry)
#2745 quantizes the
grouped_mmexperts forward, transformers' default, while calibration still runs every model on eager. This PR moves calibration ontogrouped_mmand thetorch_a2aexpert-parallel forward, and makes everything that reads the experts' inputs or weights during calibration work there. Each gap below otherwise failed silently, by collecting no amax or the wrong one rather than by raising.normalize_experts_impl_on_the_flyreplacesforce_eager_experts_impl_on_the_fly.eager,grouped_mmandtorch_a2aare kept; a backend nothing intercepts (batched_mm, the opaquesonicmoekernel, hub kernels) is moved togrouped_mmbefore calibration, or toeageron a transformers without_grouped_linear. A model whose experts quantize only on the eager forward keeps main's forced eager: native FP8 experts, a class with its own wrapper (GPT-OSS's quantizesdown_proj's input by intercepting the eager forward's matmuls), and Nemotron-H's non-gated experts until [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067. Such a model is refused ontorch_a2a, before any config changes, instead of calibrating on raw FP8 bytes or not at all.grouped_mmandtorch_a2afeed a projection's input quantizer one expert-sorted batch. The weight-input captures (local Hessian, imatrix, GPTQ) split it back per expert by the cumulative token offsets the grouped interception records, following whichever path ran. The input quantizers calibrate on the rows the grouped GEMM computes only: under transformers' own expert parallelism, the rows routed to other ranks' experts are sorted past the offsets, and before transformers 5.17 the down projection's input holds them uninitialized.fold_weight, the weight-input captures and GPTQ. Here the grouped captures split the expert-sorted batch per local expert. The grouped weight quantization needs no indexing: transformers swaps expert-parallel DTensors for their local shards around the experts forward.calibrate()raises, naming the modules, when a fused-experts forward that had to reach a quantizer ran but neither hook fired: a hand-writtenbmmforward on a class without@use_experts_implementation, which no config change redirects. Eachcalibrate()judges its own forwards and raises with the model back in its training mode. Under the expert parallelism of transformers 5.16 and later a module's verdict is reduced over its experts' expert-parallel group, so every rank of the group raises, not just the ranks whose experts saw rows. Not reported: experts no calibration token reached (Ernie-4.5-VL's vision experts under text-only calibration, which fall back to weight-derived scales as on main), an eager forward whose rows were all routed to other ranks' experts, and weight-only experts under max or MSE, whose weight quantizers calibrate on the weights alone.Behavior change. main forced MoE calibration onto eager; transformers resolves an unrequested backend to
grouped_mm, so fused-experts models now calibrate there. A single-GPU Qwen3-30B-A3B export, NVFP4 or FP8, matches main in every weight and weight scale, but 4,647 activation scales move, by up to 10.2%. main's own repeat is byte-identical, and loading withexperts_implementation="eager"reproduces main's checkpoint byte for byte (all 74,835 tensors for NVFP4 and 56,211 for FP8), so the backend alone moves them. The changelog says so.down_projinput (33 layers, 128 experts each)gate_projandup_projinput (1 layer)The largest moves are 10.2% on the experts (layer 25's
down_projinput) and 7.8% on attention.Why they move. Both forwards feed an input quantizer the same rows: a single GPU has no padding rows (transformers' sentinel rows exist only under its own expert parallelism, where calibration skips them; transformers zeroes them only from 5.17), the routing weights apply after
down_projon both, and neither casts. They round the block's output differently. Eager adds a token's 8 expert outputs into a bf16 buffer withindex_add_, rounding after each add, whilegrouped_mm_experts_forwardsums them in fp32 and rounds once, and its GEMMs run other kernels. From the first MoE block's output on, the residual stream differs by rounding, and where a router's 8th and 9th scores nearly tie, the routing flips. A flipped token feedsdown_projanother expert's activations, anddown_proj's input scale is a max over routed (token, expert) rows, so it moves by more than rounding would. The moved scales fit that: the experts' sharedgate_proj/up_projinput, whose rows are every token however it is routed, moved in one layer only, against 33 fordown_proj.Why
grouped_mmoff expert parallelism too, rather than eager there andgrouped_mmontorch_a2aonly: it is the forward transformers runs the model on when nothing else is asked for, its fp32 sum is the more precise of the two combines, and expert-parallel calibration runs the grouped path anyway (torch_a2aruns the local experts through_grouped_linear), so a model calibrates on one forward whether or not it is sharded.experts_implementation="eager"keeps earlier releases' checkpoints.Usage
Testing
tests/unit/torch/quantization/plugins/test_fused_experts.py, against synthetic and real transformers experts modules:_grouped_linear;grouped_mmandtorch_a2amatch eager's per expert, and a capture follows whichever path ran;down_proj's input on the eager forward.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;test_moe_a2a.py's chunked-hops test also runs forward and backward through hops sliced below the limit, and checks the output and input gradients bit for bit against the unsliced run.The transport's ([OMNIML-6059] [1/21] Add the DP=EP all-to-all expert transport #2738, merged) docs say again that ModelOpt quantizes its local experts, now that the
_grouped_linearoverride does. Its hop-slicing comments 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.Against main, on tiny real models: GPT-OSS calibrates on eager with
down_proj's input amax equal to main's; Ernie-4.5-VL-MoE quantizes in FP8 and NVFP4 under text-only calibration, as on main.Byte-for-byte checkpoint comparison against main on one GPU, measured on earlier revisions of [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745: Qwen3-30B-A3B, eager and
grouped_mm, NVFP4 and FP8; Qwen2.5-0.5B (dense) is byte-identical. On GPU, Nemotron-3-Nano-30B-A3B's NVFP4 calibration with its experts pinned to eager and [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067's withexperts_implementation="eager"give byte-identical exports, all 18,012 tensors.The GPU tests of QAT under expert parallelism are in [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742, the first PR with an FSDP2-sharded dense part they also run.
Review fixes, 2026-10-27, carried from [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745: the FSDP2 window's ownership test (unnamed meshes, every shard state); an empty experts forward is no work the interception check expects; the non-gated experts are pinned to eager whichever plugin runs first.
Review fixes, 2026-10-27, last round, each with a test that fails without it:
grouped_mmis kept, and redirected to, only where_grouped_linearcan be intercepted; one capability constant gates the interception, the supported set and the redirect target;Split from [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 on 2026-10-27, night: [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's and this PR's commits together give [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's previous tree exactly, plus this round's fixes. Pre-commit, mypy included, and this PR's unit tests pass at this head on transformers 5.18 and 5.5.4 and on torch 2.8, and the series' PRs merged together add no unit-test failure over main f03651e (tests/unit/torch utils, export, quantization and models). This revision's GPU tests run in CI only.
Review suggestions, 2026-10-28: the weight-input captures (local Hessian, imatrix, GPTQ) read the path that ran, not the one the config named when they were registered, so they follow whichever forward runs and the refusal is gone, pinned by a test that fails on the code before; the interception check's forward count, and the host sync its routing test costs, stop after the first hook call of each calibration; the check's docstring says it is rank-local. Rebased over [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's fixes, taking back the
torch_a2atest fixtures [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 no longer carries.Rebased over [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's second-round fixes: the grouped interceptor records the offsets and clears the eager expert index, and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's test that captures skip the grouped batch gives way to this PR's per-expert capture tests.
Review follow-up, 2026-10-28, third round: drops the note [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 now adds, that the grouped captures skip their batch, since here they split it per expert. A new test pins that a calibration forward stops counting, and skips its host-syncing routing test, once a hook has fired; it fails with that gate removed.
Review suggestions, 2026-10-28, fourth round: the experts forward clears
_grouped_offswhen it returns, so a capture reached outside a forward skips instead of splitting its input by a finished forward's offsets; the interception check's docstring says that a rank whose experts saw no rows waits in its next collective until the launcher stops it or the collective times out, so the error is in another rank's log. Rebased over [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's cleanup.Rebased on 2026-10-28, evening, onto [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's head
e7c449fcf87, after [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 moved onto main7092db6ac5band took its last review rounds. Conflicts only in the changelog (both sides' entries kept),_QuantFusedExperts's docstring (it keeps [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's note that thegrouped_mmfallback does not callF.linear) and the fused-experts tests ([OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's transposed-container tests kept). The FP8 experts wrapper's comment from [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 now namesnormalize_experts_impl_on_the_fly, this PR's name for the function that keeps it on eager. pre-commit, mypy included, passes at this head against main.Review fix, 2026-10-28, evening:
calibrate()'s interception check reads the HF plugin fromsys.modulesinstead of importing it, so a plugin that failed to load (whichimport_pluginonly warns about) no longer fails every calibration with its error.test_a_plugin_that_failed_to_load_leaves_calibration_alonebreaks the plugin's import and fails without the change. pre-commit, mypy included, passes at this head.Split, 2026-10-28, evening, at the bot review's request: the FSDP2 window's ownership filter and its ownership test moved to [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088, with a test that the window leaves an ignored expert-parallel DTensor sharded. This PR sits on a helper base merging [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 and [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088, and its source diff is 334 lines. At this head
test_fused_experts.py,test_moe_a2a.pyandtest_utils.pypass on CPU (200 passed, 6 skipped), and pre-commit, mypy included, passes against main.Review suggestions, 2026-10-28, evening:
_local_expertsrefuses a DTensor whose local shard does not hold the module's experts (one replicated over the ranks, say) instead of indexing the global experts, and the native FP8 wrapper'sunpack_weightrefuses expert weights sharded over the experts instead of filling a buffer for every expert. One two-rank test covers both and fails without either. The suggested all-reduce in the interception check is declined (a collective in every distributedcalibrate()). Focused tests pass (201 passed, 6 skipped), and pre-commit passes against main.Review suggestions, 2026-10-28, late evening: the per-expert split of a grouped batch requires exactly one exclusive end per expert, within the rows, as transformers'
grouped_mmanda2a_moebuild the offsets, and raises on any other form (a leading zero would credit each expert with the previous one's rows). A test feeds three such forms and fails without the check. The_grouped_offscomment names all three captures and the attribute's role as the which-forward-ran signal. Rebased onto the helper base re-merged with [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088's update; focused tests pass (204 passed, 6 skipped), and pre-commit passes against main.Review suggestion, 2026-10-28, late evening:
_local_expertschecks for a DTensor through the plugin's guardedtorch.distributed.tensorimport, so a torch without distributed does not fail every fused-MoE call. Rebased onto the helper base re-merged with [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088's update; focused tests pass (204 passed, 6 skipped), and pre-commit passes against main.2026-10-28, night:
hf_ptq.pygains--experts_implementation, passed on every load path:get_model(the compressed-tensors branch included), the low-memory load, and the FSDP2 load, through the config its meta model is built from. The changelog's way back to earlier releases' scales is now reachable from the script. A test loads a tiny Qwen3-MoE witheagerand fails without the plumbing. Thetests/examples/hf_ptqunit tests pass (110), and pre-commit passes against main.Review fixes, 2026-10-28, night: on the FSDP2 path the flag is set on the pre-fetched config's
_experts_implementation, asbuild_meta_causal_lmsetsattn_implementation, and a test captures the config the FSDP2 loader receives (it fails without the assignment);get_modelpasses the flag only when given, on every branch.calibrate()'s new error for a fused-experts forward that reaches neither hook has its own changelog entry. Rebased onto [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745, itself rebased onto mainbc0b6e03237after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 merged; focused andtests/examples/hf_ptqtests pass (316 passed, 6 skipped), and pre-commit passes against main.Review suggestions and a rebase, 2026-10-28, late night: the backend normalization warns once for each change it makes to a caller's choice, but not for a
Nonebackend pinned to eager, which already dispatches to the eager forward. Rebased onto [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 on maine9f432affbb(Bump minimum torch to 2.10 and drop older-torch code paths #3091):DTensornow comes from the plugin's plaintorch.distributed.tensorimport, since the guarded import is gone with the torch 2.10 floor. Focused andtests/examples/hf_ptqtests pass (325 passed, 6 skipped), and pre-commit passes against main.Review suggestion and a rebase, 2026-10-28, late night: the backend-redirect notice goes through
warn_rank_0, so an expert-parallel run prints one copy. Rebased onto [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745's comment and warning-text update;_is_fused_experts_module's new docstring namesnormalize_experts_impl_on_the_fly, which keeps GPT-OSS on eager here.test_fused_experts.pyandtest_utils.pypass (197), and pre-commit passes against main.Bot review of
2e253a509d0, 2026-10-28, late night: the FSDP2 test runs a Qwen3.5-MoE config too and checks the text config its experts dispatch on, which the one assignment reaches through transformers' recursive_experts_implementationsetter (5.5.4 and 5.18; a tiny Qwen3.5-MoE built on meta dispatches its experts oneagerwith the flag and ongrouped_mmwithout it). The test fails when the value reaches the root config only, and the comment at the assignment says the setter carries it. Declined: resolving the routing indices' parameter name from the forward's signature, since every transformers MoE block passes them positionally. Thetests/examples/hf_ptqunit tests pass on 5.18 and 5.5.4 (120), and pre-commit passes against main.Bot review of
042e4bf1e55(approved), 2026-10-28, late night: the interception error names the ways forward: a class no backend change redirects needs its own quantized wrapper, and meanwhile the rest of the model quantizes with these modules' quantizers disabled, which the check honours. Declined: caching the grouped offsets' host copy across the two projections' captures (one small sync per experts module per batch under GPTQ, local Hessian and imatrix, whose per-expert matmuls dominate). pre-commit passes against main.Rebased on 2026-10-28, late night, onto main
1ca8bd1fc4dafter [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged (squashed with the tree of its last head, so this PR's tree is unchanged); the base is nowmain, so the GPU tests run. pre-commit passes against main.CI fix, 2026-10-28, late night: the first run on main failed two cases of
tests/examples/hf_ptq/test_hf_ptq_vision_quantization.pyin the trtllmhf_ptqlane (2 failed, 267 passed): its stand-in for hf_ptq's parsed arguments lacked the newexperts_implementation, whichload_modelnow reads. The stand-in gains it; the file passes on transformers 5.18 and 5.5.4.Rebased on 2026-10-28, late night, onto main
bdc45368d10([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged), with its own diff unchanged. Its tests with [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067's, [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736'stest_distributed.pyand thehf_ptqunit tests pass (376), and pre-commit passes against main.CI, 2026-10-29: the first full run on main at
df2e3fe4a76, GPU lanes included, passes.Split, 2026-10-29, at a reviewer's request: the expert-parallel indexing fix, its tests and a Bug Fixes changelog entry moved to [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 ([1/2], on main), and this PR now sits on it, its own diff otherwise unchanged. The review's three smaller items are taken: the interception check reduces its verdict over the experts' expert-parallel group (a two-rank test fails with the rank-local check); a test checks that every experts class transformers decorates with
use_experts_implementationtakes the routing indices second (transformers 5.5.4 and 5.18); and--experts_implementationofferseagerandgrouped_mmas argparse choices. pre-commit, mypy included, passes against main; this PR's tests,test_moe_a2a.pyand thehf_ptqunit tests pass (315).Review suggestion (Claude), 2026-10-29: the interception check reads the expert-parallel group from the experts' weight, a DTensor sharded over the experts from transformers 5.16 on; from 5.6 to 5.15 the shards are plain tensors that name no group, so the check stays rank-local there, as before. The docstrings say so. Restacked onto [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122's two follow-ups; pre-commit passes against main.
Comment, 2026-10-29: at
calibrate()'s gate for the interception check, a sentence says that both conditions, the forward loop having run and the HF plugin being loaded, must come out the same on every rank, since under expert parallelism the check all-reduces over the experts' group. pre-commit passes against main.Test fix, 2026-10-29:
tests/gpu/torch/quantization/test_fsdp2_autoquant.py's 2-GPU KV AutoQuantize test builds hf_ptq's arguments by hand and lacked the newexperts_implementation, whichload_model's FSDP2 branch reads. It now passesNone. The 1-GPU PR runners skip that test, and main's nightly runs it. No other test or example builds hf_ptq's arguments by hand on a path that reads the flag. pre-commit passes against main.Rebased on 2026-10-30 onto [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 on main
99c3834f62c, with its own diff unchanged and no conflicts. pre-commit passes against main.Review (sychen52), 2026-10-30: with every weight quantizer collecting statistics without quantizing (calibration), the grouped forward returns the stored weight itself instead of building a weight-sized copy slice by slice on each forward. A CPU test checks the identity and the collected per-expert amax, and fails on the previous code. On CPU, one call on a Qwen3-30B-A3B-sized
gate_up_proj(128 × 1536 × 2048, bf16, 768 MiB) raised peak RSS by 755 MiB before and 0 MiB after; a GPUmax_memory_allocatedcomparison is queued. The changelog entry now says the experts stay ongrouped_mmfor every forward aftermtq.quantizeormto.restore(QAT/QAD, AutoQuantize scoring, fake-quant evaluation), withexperts_implementation="eager"as the way back.test_fused_experts.pyandtest_moe_a2a.pypass on transformers 5.18 (220) and 5.5.4 (214, 6 skipped); pre-commit passes against main.Review (kevalmorabia97), 2026-10-30: the three changes no test covered now have one each, or are gone. The
_local_expertscall in_quantize_grouped_weightis dropped as unreachable: the weight arrives as this rank's local shard inside the experts forward. New tests pin that a grouped forward reaching neither hook is counted whatever its routing indices hold, and that a capture after a grouped forward does not split by its offsets. The calibration pass-through calls only quantizers that calibrate (Claude's suggestion), with a constant-amax test case. Each test fails with its change reverted.test_fused_experts.py,test_moe_a2a.pyandtest_hf_ptq_args.pypass on transformers 5.18 (283) and 5.5.4 (277, 6 skipped); pre-commit passes against main.GPU calibration peak memory, 2026-10-30 (sychen52's request):
torch.cuda.max_memory_allocatedaround themtq.quantize(...)call only, afterreset_peak_memory_stats. Qwen3-30B-A3B on the defaultgrouped_mmbackend, recipegeneral/ptq/fp8_default-kv_fp8,calib_size32,batch_size1, one NVIDIA GB200 for both revisions (torch 2.11, transformers 5.13.0, CUDA 13.1).3f3544435d8)27c5f77feac)Rebased on 2026-10-30 onto main
55898ee4e3bafter [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 merged, dropping [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122's pre-squash commits; its own diff is unchanged, with no conflicts. pre-commit passes against main.CI fix, 2026-10-30: the routing-order test read transformers' modeling sources with the platform's default encoding, cp1252 on Windows, which 16 of them (
modeling_altclip.py, say) cannot be decoded with; it now reads them as UTF-8, and the windows unit job's only failure goes with it. Thelinuxjob's one failure on the previous head, a NAS sorting test this PR does not touch, passed on rerun at the same commit. pre-commit passes against main.eageragainst main, byte for byte, 2026-10-30 (re-verifying at this PR's head the claim thatexperts_implementation="eager"reproduces main): PTQ of Qwen/Qwen3-30B-A3B with recipegeneral/ptq/nvfp4_default-kv_fp8,calib_size512, calibration sequence length 2048,batch_size1, cnn_dailymail, without expert parallelism (one process), on one NVIDIA GB200 (torch 2.11, transformers 5.13.0, CUDA 13.1). This PR at27c5f77feacwith--experts_implementation eager, against main at99c3834f62c, both compared at zero tolerance. Later revisions of this PR leave its own diff unchanged (a rebase) or change tests only.input_scaleweight_scaleweight_scale_2k_scalev_scaleReview (juhi10071998), 2026-10-30: under transformers' own expert parallelism the grouped forward sorts the rows routed to other ranks' experts past
offs[-1], which the grouped GEMM neither reads nor writes; before transformers 5.17 nothing zeroes them in the down projection's input, so its input quantizer calibrated on uninitialized rows. While an input quantizer collects statistics it now sees only the rows beforeoffs[-1](a host sync, during calibration only). A test hands the grouped projections a NaN or a large finite tail, as pre-5.17 expert parallelism leaves it, and checks the input amaxes match a run without one; both cases fail without the change.test_fused_experts.py,test_moe_a2a.py,test_calib.pyandtest_hf_ptq_args.pypass on transformers 5.18 (307) and 5.5.4 (301, 6 skipped); pre-commit passes against main.Rebased on 2026-10-30 onto main
059588b89a5(after [OMNIML-6060] [1/3] Add require_transformers, and build a VLM without its LM head #3131, Keep GatedDeltaNet beta / alpha gates unquantized in Megatron fake quant #3124, [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, Skip long-lived objects in the per-step gc.collect() of QAT and QAD #3104, [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128, [OMNIML-6076] [1/3] Describe expert parallelism in the model specs #3134 and ci: establish a shared review baseline and adopt it in Claude reviews #3149 merged), with its own diff unchanged and no conflicts. pre-commit passes against main.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/A042e4bf1e55; the error-message change after it is from that review.Additional Information
Split out of #2745 (part 3 of 21 of expert-parallel PTQ for Hugging Face MoE models), which merged on 2026-10-28.
Summary by CodeRabbit
grouped_mmbackend instead of forcing eager execution. Calibrated activation scales may differ; setexperts_implementation="eager"to use the earlier calibration behavior.torch_a2araises an error.