[shengliangxu]

What does this PR do?

Type of change: New feature, with a backward-incompatible change to MoE calibration (see the changelog entry)

The fused-experts wrapper handled one shape of MoE: gated experts (gate_up_proj/down_proj) in the (num_experts, out, in) layout, calibrated through the eager backend, where Hugging Face calls the expert matmul once per expert. Several real architectures are not that shape. Each failed differently, usually by silently collecting no amax rather than by raising.

• grouped_mm. transformers' default experts implementation is grouped_mm, not eager, and it passes one expert-sorted batch per projection instead of one call per expert, so the eager interception had nothing to hook. The wrapper now intercepts the grouped path and keeps the cumulative token offsets, and the local-Hessian capture splits the batch back into per-expert chunks with them. Backends that cannot be intercepted (batched_mm, the opaque sonicmoe kernel) are moved to grouped_mm before calibration rather than quietly producing an uncalibrated model.
• Ungated experts. Nemotron-H has a single up_proj with no gate/up split. The first projection's name is now a class attribute the subclass overrides, the quantizer names derive from it, and _is_gated = False tells the export not to split it.
• Transposed experts ((num_experts, in, out)) without a wrapper of their own are left unquantized, with a warning. The fused-experts export would split gate and up along the wrong axis for them and write wrong weights; main did so silently. GPT-OSS and Llama-4 keep their own wrappers.
• 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.
  • Both now index the rank's own shard, and FSDP2's window leaves DTensors that are not on FSDP's mesh alone.
• Fail loudly. calibrate() raises when a fused-experts module's forward ran but neither hook fired. Every failure above otherwise looks the same: an export that seems complete and is quietly unquantized. An experts module that no calibration token reached is not an error and falls back to weight-derived scales as on main: Ernie-4.5-VL's vision experts, under text-only calibration.

Behavior change. main forced MoE calibration onto eager; transformers now resolves an unrequested backend to grouped_mm, so fused-experts models calibrate there. A model whose experts class has its own quantized wrapper keeps main's forced eager. GPT-OSS's wrapper, for instance, quantizes down_proj's input by intercepting the eager forward's matmuls. A single-GPU Qwen3-30B-A3B export matches main in every weight and weight scale, but 4,647 activation scales move, by up to 10.2%. Loading with experts_implementation="eager" reproduces main's checkpoint byte for byte (all 74,835 tensors for NVFP4 and 56,211 for FP8). The changelog says so. On the grouped path each forward fake-quantizes a projection's experts together, so calibration holds a transient copy of that projection's weights.

Usage

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

Testing

• tests/unit/torch/quantization/plugins/test_fused_experts.py (+611 lines), against synthetic and real transformers experts modules:

  • grouped_mm interception and its offsets;
  • backend normalization;
  • ungated (Nemotron-H) and split-gated layouts;
  • local quantizer indexing;
  • the transposed guard, which also leaves a class with its own wrapper alone;
  • the interception check: it raises when the forward ran past both hooks, and not when no token reached the module;
  • GPT-OSS, a real tiny model, calibrating down_proj's input on the eager forward.

  The last two, and the transposed warning, fail on the code before those fixes.

• Expert parallelism, on CPU under gloo:

  • experts sharded the way transformers shards them calibrate and fold on their own weights, including inside FSDP2's weight-access window; both cases fail on the code before the fix;
  • 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 (mutation-checked twice).

• The GPU tests of QAT under expert parallelism, which this PR carried before the series was reordered, are in [10/21], the first PR with an FSDP2-sharded dense part they also run.

• Against main, on tiny real models:

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

• Local-Hessian calibration captures each expert's own rows under grouped_mm and under the torch_a2a expert-parallel forward ([1/21]); both feed the quantizers one expert-sorted batch.

• Registering those input hooks under expert parallelism reads each rank's own experts: indexing the expert-sharded weight all-gathered every expert of the projection. A two-rank test checks that registration issues no collective.

• Near-duplicate tests are merged into parametrized ones, tests a remaining test already covers are dropped, and the two expert-parallel jobs become cases of the tests they repeated. Each drop is checked by mutation.

• Locally on this revision, the reordered series rebased on main 264fd26: the full unit suite on transformers 5.18, and on 5.5 apart from the tests whose ONNX, TensorRT or diffusers packages that environment lacks; CI's transformers-free unit lane; pre-commit over the whole PR against main; and the Sphinx docs build. The same at every PR of the stack.

• The GPU results here were measured on earlier revisions of this branch; this revision's GPU tests run in CI only.

• Byte-for-byte checkpoint comparison against main on one GPU: Qwen3-30B-A3B, eager and grouped_mm, NVFP4 and FP8. Qwen2.5-0.5B (dense) is byte-identical.

Before your PR is "Ready for review"

• 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?: ❌ Not yet.

Additional Information

Part 3 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after [2/21].

Summary by CodeRabbit

• Compatibility

  • Hugging Face MoE calibration now uses the backend selected by Transformers when supported, rather than always using eager execution. Activation scales may differ from earlier releases; set experts_implementation="eager" to restore the previous calibration setup.
  • Unsupported expert implementations are normalized to a supported backend during calibration.

• Bug Fixes

  • Improved calibration and quantization support for grouped and distributed expert execution, including expert-parallel and FSDP2 layouts. Calibration now reports an error if a quantized fused-expert module runs without being intercepted.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[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: 691ece4
• Latest commit, not reviewed: aa80229

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

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 57622c8a-2115-413f-9fc1-86b04e66d1d8

📥 Commits

Reviewing files that changed from the base of the PR and between eb9c0ae and 691ece4.

📒 Files selected for processing: 2

• tests/unit/torch/quantization/plugins/test_fused_experts.py
• tests/unit/torch/utils/test_moe_a2a.py

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

📝 Walkthrough

Fused-expert quantization now supports grouped dispatch and backend normalization, including ungated NemotronH experts. Calibration captures grouped expert inputs, validates interception when its forward loop runs, and handles transposed weights and local expert shards.

Changes: MoE Calibration

Layer / File(s)	Summary
Backend selection and fused-expert interception modelopt/torch/quantization/plugins/huggingface.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, modelopt/torch/export/moe_utils.py, modelopt/torch/models/step3p5/modeling_ptq.py	Wrappers intercept eager and grouped execution. Backend normalization preserves supported implementations and redirects unsupported implementations to grouped_mm, with separate handling for modules that have a registered wrapper. Ungated NemotronH experts use _QuantNemotronHExperts. Tests cover backend normalization, wrapper selection, transposed layouts, and forward equivalence.
Grouped calibration and interception validation modelopt/torch/quantization/model_calib.py, modelopt/torch/quantization/model_quant.py, modelopt/torch/quantization/plugins/huggingface.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, CHANGELOG.rst	Calibration hooks use grouped offsets to capture inputs for enabled experts. calibrate checks interception only when its forward loop ran. Weight calibration and folding handle transposed orientations. Tests compare grouped and eager calibration and cover interception checks. The changelog documents the grouped_mm calibration backend and the eager setting.
Sharded expert calibration and folding modelopt/torch/quantization/utils/core_utils.py, tests/unit/torch/quantization/plugins/test_fused_experts.py, tests/unit/torch/utils/test_moe_a2a.py	FSDP2 weight access skips DTensor parameters whose named mesh dimensions do not match the FSDP mesh. Tests cover expert-parallel gradients, per-rank calibration and folding, and hook registration without collectives.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Calibration
  participant QuantFusedExperts as _QuantFusedExperts
  participant GroupedLinear as _grouped_linear
  Calibration->>QuantFusedExperts: run expert forward
  QuantFusedExperts->>GroupedLinear: intercept grouped linear call
  GroupedLinear-->>QuantFusedExperts: return grouped output
  QuantFusedExperts-->>Calibration: return quantized output

Loading

Merge Risk: ⚪ Minimal · up to 691ec

This change adds tests for fused-expert quantization under grouped_mm and expert parallelism. No remaining test-hygiene or correctness concerns were established, so merge risk is minimal.

Pre-merge checks: 5 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 58.06% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 93 functions across 8 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 5 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		No stated security anti-pattern was introduced. The PR changes only modelopt Python files and tests; it adds no examples or dependency-file changes. Added modelopt lines contain no unsafe torch.load, …
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly identifies the main change: quantizing fused experts for grouped_mm, ungated, and transposed layouts.

Full details: Security Anti-Patterns

Explanation

No stated security anti-pattern was introduced. The PR changes only modelopt Python files and tests; it adds no examples or dependency-file changes. Added modelopt lines contain no unsafe torch.load, pickle-enabled numpy.load, hardcoded trust_remote_code=True, eval/exec on external input, or # nosec comments. The existing trust_remote_code=True call in huggingface.py is unchanged from the base revision.

✨ Finishing Touches: 💡 2

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2745/
Built to branch gh-pages at 2026-10-23 05:39 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 95.03106% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.63%. Comparing base (719a36d) to head (aa80229).

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/plugins/huggingface.py	95.76%	5 Missing ⚠️
modelopt/torch/quantization/model_calib.py	90.00%	3 Missing ⚠️

Additional details and impacted files

@@                       Coverage Diff                        @@
##           shengliangx/hfep-7-amax-sync    #2745      +/-   ##
================================================================
+ Coverage                         79.61%   79.63%   +0.02%     
================================================================
  Files                               660      660              
  Lines                             74774    74906     +132     
================================================================
+ Hits                              59528    59653     +125     
- Misses                            15246    15253       +7     

Flag	Coverage Δ
unit	61.21% <95.03%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test bad2de9

[shengliangxu]

/ok to test 16e5c17

[shengliangxu]

/ok to test 698f459

[shengliangxu]

/ok to test eaa41c2

[shengliangxu]

/ok to test f2be168

[shengliangxu]

/ok to test 233d659

[shengliangxu]

/ok to test 2a18b13

[shengliangxu]

/ok to test 6c1e2d3

[shengliangxu]

/ok to test e322d20

[shengliangxu]

/ok to test b452b7b

[shengliangxu]

/ok to test a38ba3e

[shengliangxu]

/ok to test 99ff96e

[shengliangxu]

/ok to test 220ad2d

[shengliangxu]

/ok to test 0124a4e

[shengliangxu]

/ok to test efce4ee

[shengliangxu]

/ok to test 4038ec7

[shengliangxu]

/ok to test aa80229