fix(export): unblock AWQ export for fused MoE experts (Nemotron-H)#2017
Draft
cjluo-nv wants to merge 1 commit into
Draft
fix(export): unblock AWQ export for fused MoE experts (Nemotron-H)#2017cjluo-nv wants to merge 1 commit into
cjluo-nv wants to merge 1 commit into
Conversation
int4_awq PTQ on NVIDIA-Nemotron-3-Super-120B-A12B-BF16 could not produce a
checkpoint. Calibration succeeded; export failed. Two independent defects, both
only reachable on AWQ formats, and both specific to the fused-experts layout that
transformers >= 5.5 uses for NemotronH (NemotronHExperts -> 3-D packed params,
wrapped as _QuantNonGatedFusedExperts). FP8/NVFP4 were unaffected.
1) requantize_resmooth_fused_llm_layers() called get_experts_list() for every MoE
block, gated on "awq" in quantization_format. That helper addresses experts
positionally (module.experts[i].<linear>), which only exists for the
sequential nn.ModuleList layout, so a fused wrapper raised
"TypeError: object of type 'QuantNemotronHExperts' has no len()".
Fused blocks are now skipped: resmoothing exists to give every expert in a
group one common input pre_quant_scale, and a fused wrapper already shares a
single input quantizer across all its experts, so there is nothing to
reconcile.
2) With that unblocked, _export_fused_experts()'s uncalibrated-expert fallback
assigned a scalar amax. That is only valid for per-tensor formats; for
block-wise formats (INT4 AWQ, block_sizes={-1: 128}) the exporter then
evaluated weight.shape[-1] // weights_scaling_factor.shape[-1] on a 0-d
tensor and raised "IndexError: tuple index out of range". The fallback now
derives a per-block amax shaped (out_features, n_blocks), matching what
MaxCalibrator produces for a calibrated expert. This path is easy to hit on
this model: 512 routed experts at top-22 leaves many experts unrouted for any
modest calibration set.
Verified end-to-end on 4xB200: hf_ptq.py --qformat int4_awq on the full 120B now
writes a complete checkpoint (8 shards, 68G, W4A16_AWQ group_size=128), with
expert weight_scale shapes (2688, 8) and (1024, 21) matching 1024/128 and
2688/128 blocks. Adds a regression test that fails on the pre-fix fallback.
Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Contributor
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2017 +/- ##
==========================================
- Coverage 78.36% 78.35% -0.01%
==========================================
Files 518 518
Lines 58658 58669 +11
==========================================
+ Hits 45965 45973 +8
- Misses 12693 12696 +3
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
int4_awqPTQ on nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 could not produce a checkpoint. Calibration succeeded; export failed. Two independent defects, both reachable only on AWQ formats, and both specific to the fused-experts layouttransformers>=5.5uses for Nemotron-H (NemotronHExperts→ 3-D packed params, wrapped as_QuantNonGatedFusedExperts). FP8 / NVFP4 were unaffected, which is why only AWQ was reported broken.1.
TypeError: object of type 'QuantNemotronHExperts' has no len()requantize_resmooth_fused_llm_layers()calledget_experts_list()for every MoE block, gated on"awq" in quantization_format. That helper addresses experts positionally (module.experts[i].<linear>), which only exists for the sequentialnn.ModuleListlayout — a fused wrapper has no__len__and no per-expert submodules.Fused blocks are now skipped. This is correct rather than merely crash-avoiding: resmoothing exists to give every expert in a group one common input
pre_quant_scale, and a fused wrapper already shares a single input quantizer across all its experts, so there is nothing to reconcile.2.
IndexError: tuple index out of rangeWith #1 unblocked,
_export_fused_experts()'s uncalibrated-expert fallback assigned a scalaramax. That is only valid for per-tensor formats; for block-wise formats (INT4 AWQ,block_sizes={-1: 128}) the exporter then evaluatedweight.shape[-1] // weights_scaling_factor.shape[-1]on a 0-d tensor.The fallback now derives a per-block
amaxshaped(out_features, n_blocks), matching whatMaxCalibratorproduces for a calibrated expert. This path is easy to hit on wide MoEs: 512 routed experts at top-22 leaves many experts unrouted for any modest calibration set.The two are sequential — fixing #1 alone just moves the crash to #2.
Usage
Testing
quant_algo: W4A16_AWQ,group_size: 128, zero tracebacks. Expertweight_scaleshapes verified at real dims:up_proj (2688, 8)= 1024/128 blocks,down_proj (1024, 21)= 2688/128 blocks; packeduint8weights at half the output rows; all scales finite.test_uncalibrated_expert_block_quant_amax_is_per_block— asserts an uncalibrated expert under a block-wise format gets a 2-D per-blockamax. Verified it fails on the pre-fix fallback and passes after.tests/unit/torch/export+tests/unit/torch/quantization: 962 passed, 8 skipped (one unrelated file skipped —test_pytorch_geometric_plugin.py, missing optional dep).Before your PR is "Ready for review"
nn.ModuleListpath is untouched. The block-awareamaxonly changes the uncalibrated fallback, and only for block-wise formats; per-tensor formats still get the scalaramaxas before.CONTRIBUTING.md: N/A — no copied code, no new dependencies.Additional Information
Scope is not Nemotron-H-specific: any model using
transformers>=5.5fused experts (_QuantFusedExperts/_QuantNonGatedFusedExperts) with an AWQ format hits the same two crashes.Two related issues found while debugging, deliberately left out of this PR as separate concerns:
--moe_calib_experts_ratiois a silent no-op on the fused path —mode.pyapplies it viahasattr(module, "_moe_calib_experts_ratio"), and that attribute is defined on_QuantSparseSequentialMoeonly. Export now succeeds, but unrouted experts get a weight-derivedamaxrather than an activation-calibrated one, so the accuracy tail is still open.hf_quant_config.jsonexclude_modulesentries contain a literal NUL byte followed by apt_name_sentinelsuffix (i.e.lm_head.+\x00+backbone.pt_name_sentinel). Pre-existing and format-independent (identical under FP8), but it looks like a hazard for downstream vLLM / TRT-LLM config parsing.🤖 Generated with Claude Code