Skip to content

[OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager - #3074

Merged
shengliangxu merged 42 commits into
mainfrom
shengliangx/hfep-8a-grouped-mm-calibration
Oct 30, 2026
Merged

shengliangxu merged 42 commits into
mainfrom
shengliangx/hfep-8a-grouped-mm-calibration

Conversation

@shengliangxu

@shengliangxu shengliangxu commented Oct 27, 2026 •

Copy link
Copy Markdown
Collaborator

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_mm by default and the new calibrate() 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_mm quantization), 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

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, 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.
  • Expert parallelism. [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 makes everything that reads the expert weights outside the forward index each rank's own experts: max calibration's weight pass, 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.
  • 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. 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 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, where calibration skips them; transformers zeroes them only from 5.17), 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  # or grouped_mm

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:

    • the tests that experts sharded the way transformers shards them calibrate, fold, capture and build GPTQ helpers on each rank's own experts moved to [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 with that fix;
    • every rank of the experts' expert-parallel group raises when only one rank's experts ran past both hooks;
    • 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.

  • 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_implementation takes the routing indices second (transformers 5.5.4 and 5.18); and --experts_implementation offers eager and grouped_mm as argparse choices. pre-commit, mypy included, passes against main; this PR's tests, test_moe_a2a.py and the hf_ptq unit 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 new experts_implementation, which load_model's FSDP2 branch reads. It now passes None. 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 GPU max_memory_allocated comparison is queued. The changelog entry now says the experts stay on grouped_mm for every forward after mtq.quantize or mto.restore (QAT/QAD, AutoQuantize scoring, fake-quant evaluation), with experts_implementation="eager" as the way back. test_fused_experts.py and test_moe_a2a.py pass 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_experts call in _quantize_grouped_weight is 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.py and test_hf_ptq_args.py pass 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_allocated around the mtq.quantize(...) call only, after reset_peak_memory_stats. Qwen3-30B-A3B on the default grouped_mm backend, recipe general/ptq/fp8_default-kv_fp8, calib_size 32, batch_size 1, one NVIDIA GB200 for both revisions (torch 2.11, transformers 5.13.0, CUDA 13.1).

    Revision Calibration peak (bytes) MiB
    Before the pass-through change (3f3544435d8) 864,901,632 824.8
    After (27c5f77feac) 174,759,936 166.7
    Drop 690,141,696 658.2
  • Rebased on 2026-10-30 onto main 55898ee4e3b after [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. The linux job'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.

  • eager against main, byte for byte, 2026-10-30 (re-verifying at this PR's head the claim that experts_implementation="eager" reproduces main): PTQ of Qwen/Qwen3-30B-A3B with recipe general/ptq/nvfp4_default-kv_fp8, calib_size 512, calibration sequence length 2048, batch_size 1, cnn_dailymail, without expert parallelism (one process), on one NVIDIA GB200 (torch 2.11, transformers 5.13.0, CUDA 13.1). This PR at 27c5f77feac with --experts_implementation eager, against main at 99c3834f62c, both compared at zero tolerance. Later revisions of this PR leave its own diff unchanged (a rebase) or change tests only.

    Scale kind Tensors Differing
    input_scale 18,624 0
    weight_scale 18,624 0
    weight_scale_2 18,624 0
    k_scale 48 0
    v_scale 48 0
  • Review (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 before offs[-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.py and test_hf_ptq_args.py pass 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"

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

@shengliangxu
shengliangxu requested review from a team as code owners October 27, 2026 23:40
@shengliangxu
shengliangxu requested review from kevalmorabia97 and removed request for a team October 27, 2026 23:40
@coderabbitai

coderabbitai Bot commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

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.

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

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: Passed 5
Check name Status Explanation
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns Passed 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 Passed The title clearly and concisely 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
shengliangxu requested review from kaix-nv and removed request for a team October 27, 2026 23:40
This was referenced Oct 27, 2026
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/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>
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test a6758d7

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants