[shengliangxu]

Note

[3/3] of the OMNIML-6073 split: #3088 (FSDP2's weight-access window), #2745 (grouped_mm quantization) and this PR (calibration on grouped_mm and torch_a2a). It was split out of #2745 on 2026-10-27, night, and the FSDP2 window moved out of it into #3088 on 2026-10-28, at the bot review's request. #3067 (Nemotron-H) sits on this PR.

Base: main. #3088, #2745 and #2736 merged on 2026-10-28, so this PR sits on main, rebased with its own diff unchanged.

In the tree: #2741's helper base merges #3067, which carries this PR and #2745.

Series layout

• main
  • [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 FSDP2 weight-access window (merged)
  • [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 fused experts under grouped_mm (merged)
    • [OMNIML-6073] [3/3] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 calibration on grouped_mm (this PR)
      • [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 Nemotron-H's experts under grouped_mm
  • [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 meta-DCP checkpoint load (merged)
  • [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 distribution plan, run group, collectives (merged)
    • [OMNIML-6071] Calibrate a distributed model on per-rank data #2740 calibration on per-rank data
    • [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750 keep-fused fold
    • [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751 distributed-writer pieces
      • [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it #2752 no-gather export (on [OMNIML-6082] [3/3] Write a sharded model's HF checkpoint shards from each tensor's owner rank #2751, [OMNIML-6081] [2/2] Fold fused experts' quantizers into their weights shard-locally #2750, [OMNIML-6054] Route load_pretrained by the distribution plan #2779 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
        • [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753 expert-parallel export
    • [OMNIML-6054] Decide how a model is distributed, from the model built on the meta device #2737 plan decision
      • [OMNIML-6054] Route load_pretrained by the distribution plan #2779 load routing
        • [OMNIML-6060] Load an MoE expert-parallel: shard the experts at load time #2741 expert-parallel loader (on [OMNIML-6054] Route load_pretrained by the distribution plan #2779, [OMNIML-6073] Quantize Nemotron-H's experts under grouped_mm, with their wrapper in its model folder #3067 and [OMNIML-6071] Calibrate a distributed model on per-rank data #2740)
          • [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742 dense sharding
            • [OMNIML-6062] Dequantize block-FP8 checkpoints as they load expert-parallel #2743 block-FP8 dequantization
              • [OMNIML-6060] Load an expert-parallel plan through load_pretrained #2780 expert-parallel load routing
                • [OMNIML-6076] Describe expert parallelism in the model specs #2746 expert parallelism in the model specs
                  • [OMNIML-6063] Decompress compressed-tensors MoE experts on the fly as they load #2747 on-the-fly decompression
                  • [OMNIML-6078] Add the Kimi-K2.5 and Kimi-VL expert-parallel adapters #2748 Kimi-K2.5 and Kimi-VL adapters
                  • [OMNIML-6079] Add the Kimi-K3 expert-parallel adapter #2749 Kimi-K3 adapter
                  • [OMNIML-6085] Run expert-parallel PTQ from hf_ptq.py #2754 expert-parallel PTQ from hf_ptq.py (on [OMNIML-6076] Describe expert parallelism in the model specs #2746 and [OMNIML-6084] Export an expert-parallel model through the no-gather writer #2753)

What 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_mm experts forward, transformers' default, while calibration still runs every model on eager. This PR moves calibration onto grouped_mm and the torch_a2a expert-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.

• Backend normalization. normalize_experts_impl_on_the_fly replaces force_eager_experts_impl_on_the_fly. eager, grouped_mm and torch_a2a are kept; a backend nothing intercepts (batched_mm, the opaque sonicmoe kernel, hub kernels) is moved to grouped_mm before calibration, or to eager on 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 quantizes down_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 on torch_a2a, before any config changes, instead of calibrating on raw FP8 bytes or not at all.
• Per-expert captures. grouped_mm and torch_a2a feed 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, and a capture fed by the other path than the experts' config selects is refused rather than attributed to the wrong experts.
• Expert parallelism. Each rank holds only its slice of experts, so expert weight quantizers are indexed by local position.
  • transformers keeps those expert weights as DTensors sharded over the experts. Outside the forward, max calibration's weight pass and fold_weight indexed 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.
  • Either way, every rank but rank 0 calibrated each expert as the max with rank 0's expert at the same local index, so an expert's weight scale came out too large wherever rank 0's expert had the larger one, and folding would have written the wrong experts.
  • The quantizers, captures and folds now index the rank's own shard, and [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 makes FSDP2's window gather only the parameters FSDP2 manages, in any shard state.
• Fail loudly. calibrate() raises, naming the modules, when a fused-experts forward that had to reach a quantizer ran but neither hook fired: a hand-written bmm forward on a class without @use_experts_implementation, which no config change redirects. Each calibrate() judges its own forwards and raises with the model back in its training mode. 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 with experts_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.

Moved activation scales, Qwen3-30B-A3B NVFP4, 1 GPU	Tensors
Experts' down_proj input (33 layers, 128 experts each)	4,224
Experts' shared gate_proj and up_proj input (1 layer)	256
Attention input and KV-cache	167
Total, in 45 of 48 layers (none in layers 0, 1 and 5)	4,647

The largest moves are 10.2% on the experts (layer 25's down_proj input) 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, and it zeroes them), the routing weights apply after down_proj on both, and neither casts. They round the block's output differently. Eager adds a token's 8 expert outputs into a bf16 buffer with index_add_, rounding after each add, while grouped_mm_experts_forward sums 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 feeds down_proj another expert's activations, and down_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' shared gate_proj/up_proj input, whose rows are every token however it is routed, moved in one layer only, against 33 for down_proj.

Why grouped_mm off expert parallelism too, rather than eager there and grouped_mm on torch_a2a only: 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_a2a runs 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

# To calibrate exactly as earlier releases did:
model = AutoModelForCausalLM.from_pretrained(path, experts_implementation="eager")

# The same from hf_ptq.py:
python hf_ptq.py --pyt_ckpt_path <model> --qformat nvfp4 --experts_implementation eager

Testing

• tests/unit/torch/quantization/plugins/test_fused_experts.py, against synthetic and real transformers experts modules:
  • backend normalization, including the fallback to eager on a transformers without _grouped_linear;
  • GPTQ's Hessians and local Hessian's captures under grouped_mm and torch_a2a match eager's per expert, and a capture follows whichever path ran;
  • the interception check raises when the forward ran past both hooks, also after an earlier calibration reached them, and not for unreached experts, an eager forward routed only to other ranks' experts, or weight-only experts beside a calibrated dense input quantizer under max (each expert's weight amax equals its weight's per-channel abs max). Under local Hessian and GPTQ, which read the experts' inputs, those weight-only experts still raise;
  • GPT-OSS, a real tiny model, calibrating down_proj's input on the eager forward.
• Expert parallelism, on CPU under gloo:
  • experts sharded the way transformers shards them calibrate and fold on their own weights, including inside FSDP2's weight-access window, under [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088's ownership filter, and with unnamed meshes; these fail on the code before the fixes;
  • registering the weight-input hooks reads each rank's own experts without a collective, and GPTQ builds its expert helpers from them;
  • with fake quantizers inserted, experts split over two ranks train through the torch_a2a forward like the whole quantized module: the fake-quantized case of test_moe_a2a.py's whole-module test;
  • the transport's hop slicing under a graph: 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_linear override does. Its hop-slicing comments say what the slicing guarantees, every hop within all_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 with experts_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_mm is kept, and redirected to, only where _grouped_linear can be intercepted; one capability constant gates the interception, the supported set and the redirect target;
  • an eager experts forward counts as work for the interception check only when it routes a row to this rank's experts, so a rank whose experts calibration routed nothing to no longer raises alone;
  • an enabled expert weight quantizer needs the interception only when the calibration captures the experts' inputs to calibrate their weights (GPTQ, local Hessian, imatrix), so weight-only experts under max or MSE no longer raise; the message names what falls back.
• 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_a2a test 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_offs when 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 main 7092db6ac5b and 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 the grouped_mm fallback does not call F.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 names normalize_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 from sys.modules instead of importing it, so a plugin that failed to load (which import_plugin only warns about) no longer fails every calibration with its error. test_a_plugin_that_failed_to_load_leaves_calibration_alone breaks 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.py and test_utils.py pass on CPU (200 passed, 6 skipped), and pre-commit, mypy included, passes against main.
• Review suggestions, 2026-10-28, evening: _local_experts refuses 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's unpack_weight refuses 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 distributed calibrate()). 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_mm and a2a_moe build 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_offs comment 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_experts checks for a DTensor through the plugin's guarded torch.distributed.tensor import, 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.py gains --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 with eager and fails without the plumbing. The tests/examples/hf_ptq unit 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, as build_meta_causal_lm sets attn_implementation, and a test captures the config the FSDP2 loader receives (it fails without the assignment); get_model passes 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 main bc0b6e03237 after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 merged; focused and tests/examples/hf_ptq tests 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 None backend 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 main e9f432affbb (Bump minimum torch to 2.10 and drop older-torch code paths #3091): DTensor now comes from the plugin's plain torch.distributed.tensor import, since the guarded import is gone with the torch 2.10 floor. Focused and tests/examples/hf_ptq tests 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 names normalize_experts_impl_on_the_fly, which keeps GPT-OSS on eager here. test_fused_experts.py and test_utils.py pass (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_implementation setter (5.5.4 and 5.18; a tiny Qwen3.5-MoE built on meta dispatches its experts on eager with the flag and on grouped_mm without 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. The tests/examples/hf_ptq unit 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 1ca8bd1fc4d after [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 now main, 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.py in the trtllm hf_ptq lane (2 failed, 267 passed): its stand-in for hf_ptq's parsed arguments lacked the new experts_implementation, which load_model now 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's test_distributed.py and the hf_ptq unit 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.

Before your PR is "Ready for review"

• Is this change backward compatible?: ❌ Fused-experts MoE calibration now runs on grouped_mm, so some activation scales differ from earlier releases; experts_implementation="eager" restores the earlier checkpoint exactly.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ Backward Breaking Changes.
• Did you get Claude approval on this PR?: ✅ At 042e4bf1e55; 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

• Behavior Changes
  • Hugging Face MoE calibration now uses Transformers’ default grouped_mm backend instead of forcing eager execution. Calibrated activation scales may differ; set experts_implementation="eager" to use the earlier calibration behavior.
  • Eager-only expert wrappers continue to use eager execution. Using them with torch_a2a raises an error.
• Bug Fixes
  • Improved calibration support for grouped and expert-parallel execution, including local expert weights and experts receiving no routed inputs.
  • Calibration now reports when fused experts run without interception or required inputs are not captured.
  • Improved FSDP2 parameter handling during weight access and writeback.
• New Features
  • Added calibration support for newer Transformers fused-expert implementations, including Qwen3-MoE and native FP8 experts.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: f1a1bc3
• Latest commit, not reviewed: ec9e09c

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: shengliangx/hfep-8a-grouped-mm-calibration
• File: .coderabbit.yaml
• Commit: f1a1bc3
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: fc3acd1f-f7a2-4d1c-95a5-7e1fd625cde4

Commits

Reviewing files that changed from the base of the PR and between 3390594 and f1a1bc3.

Files selected for processing: 1

• modelopt/torch/quantization/plugins/huggingface.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

Hugging Face MoE calibration supports grouped execution when available. Calibration captures grouped inputs and GPTQ Hessians, checks fused-expert interception, and uses rank-local expert weights. FSDP2 writeback skips DTensor parameters that are not owned by its FSDP units.

Changes: MoE Calibration and Expert-Parallel Support

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.

Priority: Normal

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

Loading

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

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
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 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=…
Title check		The title clearly and specifically describes the main change: fused-expert calibration now supports grouped_mm and torch_a2a instead of forcing eager execution.

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.}

[shengliangxu]

/ok to test e5a7530

[shengliangxu]

/ok to test ec9e09c

[juhi10071998]

Review: I'm blocking on evidence and scope, not code correctness. The code looked careful.

1. The default backend change has no accuracy data, only scale deltas. The PR shows 4,647 activation scales moving by up to 10.2% on Qwen3-30B-A3B and attributes it to combine rounding that flips near-tied top-k routes. That explains the change, but doesn't show whether it matters. Before the default flips I'd ask for:

• downstream accuracy for eager vs grouped_mm on at least one model, NVFP4 (e.g. GSM8K/MMLU on Qwen3-30B-A3B);
• a noise floor: how much down_proj input amax moves under eager when only the calibration subset or seed changes. If that is also around 10%, the backend move is within noise. If it's around 1%, the claim that grouped_mm is the "more precise" path needs data.

2. No maintainer other than the author has signed off. Every approval so far is from a bot, and the bot's own latest review still lists the maintainer sign-off on the grouped_mm default as open.

3. The evidence comes from older revisions. The byte-for-byte checkpoint comparisons, including "experts_implementation="eager" reproduces main", were measured on earlier revisions of #2745. Since eager is the advertised way back to the old behaviour, please re-verify it at this head.

4. No determinism check for grouped_mm. The PR shows main's eager repeat is byte-identical, but not a grouped_mm repeat, and not a cross-architecture comparison. torch._grouped_mm takes different kernel paths on sm90, sm100 and sm120, so calibrated scales could become hardware-dependent.

5. Scope: the EP correctness fix should land on its own. Indexing EP DTensors by global expert, which inflates expert weight scales and makes fold_weight write the wrong experts on non-zero ranks, is a silent wrong-result bug. It's bundled here with two backward-breaking changes: the new hard error in calibrate() and the default flip. Splitting it out, with a Bug Fixes changelog entry, would let it land now, and would let the default flip be reverted on its own if item 1 shows a regression.

Smaller points, not blocking:

• assert_fused_experts_intercepted is rank-local, so a rank that doesn't raise waits in its next collective until timeout. Reducing over the experts' own EP group, which every member already joins for the a2a, avoids this without a world-wide collective.
• forward() reads top_k_index as args[1] by position. A class with a different argument order would compare routing weights against num_experts and could falsely raise. A test that checks this ordering across transformers' fused-experts classes would guard it.
• --experts_implementation has no choices, so a typo is only caught late, inside transformers.