From 1231e7b56334ace3dee20d830df05de913670779 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 27 Apr 2026 16:26:36 +0000 Subject: [PATCH 1/3] [NVBUG: 6103846] Fix nvfp4_awq export for uncalibrated MoE experts awq_lite postprocess used to leave the input_quantizer without a pre_quant_scale whenever an MoE expert ended up disabled after calibration (NaN in act/weight scales or no search-pass tokens). get_quantization_format then returned 'nvfp4' for that expert while sibling experts kept 'nvfp4_awq', and export_hf_checkpoint failed in preprocess_linear_fusion with: AssertionError: Modules have different quantization formats Unify the disabled-expert paths in the postprocess loop so any expert with awq_lite.is_enabled == False gets max-calibrated weights plus a neutral all-ones pre_quant_scale, matching the existing handling of num_cache_steps == 0. Emit a warning whenever this fallback fires so users notice that calibration coverage is incomplete and accuracy may degrade. Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 4 +++ modelopt/torch/quantization/model_calib.py | 33 ++++++++++++++-------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 539cdeef819..edfbfd301da 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,10 @@ Changelog - Add FP8 MHA quantization support for vision transformers. Adds an attention-aware ONNX post-processing pass (scale Mul / K-transpose move before Q, Q→DQ insertion on softmax output) in :class:`FP8QuantExporter `, per-instance nested-attention-wrapper skipping in the HF plugin, and ``nn.LayerNorm`` registration in ``QuantModuleRegistry`` so BMM input quantizers and LayerNorm output quantizers defined in FP8_DEFAULT_CFG are honored end-to-end. See `examples/torch_onnx/torch_quant_to_onnx.py `_ for the general timm-model quantize→ONNX workflow. +**Bug Fixes** + +- Fix ``nvfp4_awq`` export ``AssertionError: Modules have different quantization formats`` for MoE models (e.g. Qwen3-30B-A3B) when some experts are not exercised by the calibration data. ``awq_lite`` now applies a neutral all-ones ``pre_quant_scale`` to any expert that ends up disabled (no cache-pass tokens, NaN scales, or no search-pass tokens) so its format remains ``nvfp4_awq``, consistent with the rest of the MoE block. A warning is emitted whenever this fallback fires. + 0.44 (2026-05-xx) ^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 04aaa88a519..67c63a99ada 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1270,11 +1270,29 @@ def postprocess(module, name): for name, module in model.named_modules(): if hasattr(module, "awq_lite"): - if module.awq_lite.num_cache_steps == 0: - # Uncalibrated expert: max calibrate weights and apply neutral - # (all-ones) pre_quant_scale for export consistency. - # NOTE: ones_scale must be registered OUTSIDE enable_weight_access_and_writeback + # Flag modules whose search pass missed them despite cache hits, so + # they fall through to the neutral-scale path below. + if module.awq_lite.num_cache_steps > 0 and module.awq_lite.num_search_steps == 0: + module.awq_lite.is_enabled = False + warnings.warn( + "awq_lite: Calling `forward_loop(model)` the second time did not forward" + f" data through the {name}. Please provide a valid `forward_loop` function" + " that can be used to forward data through the model many times." + ) + + if module.awq_lite.num_cache_steps == 0 or not module.awq_lite.is_enabled: + # Expert was either uncalibrated (no cache-pass tokens), had NaN + # in act/weight scales, or saw no search-pass tokens. Max-calibrate + # weights and apply a neutral (all-ones) pre_quant_scale so the + # exporter sees a consistent nvfp4_awq format across all expert + # linears in an MoE group. + # NOTE: ones-scale must be registered OUTSIDE enable_weight_access_and_writeback # because HF accelerate post_forward drops newly-registered submodule buffers. + warnings.warn( + f"awq_lite: Forcing pre_quant_scale=1 for {name} because the expert " + "was not properly exercised during calibration. This may degrade accuracy; " + "consider increasing calibration size or using a more diverse dataset." + ) with enable_weight_access_and_writeback(module, model, name_to_module): max_calibrate(module, lambda module: module.weight_quantizer(module.weight)) w_shape, w_dtype, w_device = ( @@ -1289,13 +1307,6 @@ def postprocess(module, name): device=w_device, ) else: - if module.awq_lite.num_search_steps == 0: - module.awq_lite.is_enabled = False - warnings.warn( - "awq_lite: Calling `forward_loop(model)` the second time did not forward" - f" data through the {name}. Please provide a valid `forward_loop` function" - " that can be used to forward data through the model many times." - ) with enable_weight_access_and_writeback(module, model, name_to_module): postprocess(module, name) From 051e4d32088073c507d2dd2d1613aa478f2b65d9 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 27 Apr 2026 16:30:50 +0000 Subject: [PATCH 2/3] [NVBUG: 6103846] Move changelog entry to 0.44 Signed-off-by: Chenjie Luo --- CHANGELOG.rst | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index edfbfd301da..44d89438323 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,10 +19,6 @@ Changelog - Add FP8 MHA quantization support for vision transformers. Adds an attention-aware ONNX post-processing pass (scale Mul / K-transpose move before Q, Q→DQ insertion on softmax output) in :class:`FP8QuantExporter `, per-instance nested-attention-wrapper skipping in the HF plugin, and ``nn.LayerNorm`` registration in ``QuantModuleRegistry`` so BMM input quantizers and LayerNorm output quantizers defined in FP8_DEFAULT_CFG are honored end-to-end. See `examples/torch_onnx/torch_quant_to_onnx.py `_ for the general timm-model quantize→ONNX workflow. -**Bug Fixes** - -- Fix ``nvfp4_awq`` export ``AssertionError: Modules have different quantization formats`` for MoE models (e.g. Qwen3-30B-A3B) when some experts are not exercised by the calibration data. ``awq_lite`` now applies a neutral all-ones ``pre_quant_scale`` to any expert that ends up disabled (no cache-pass tokens, NaN scales, or no search-pass tokens) so its format remains ``nvfp4_awq``, consistent with the rest of the MoE block. A warning is emitted whenever this fallback fires. - 0.44 (2026-05-xx) ^^^^^^^^^^^^^^^^^ @@ -51,6 +47,7 @@ Changelog - Fix Minitron pruning (``mcore_minitron``) for MoE models. Importance estimation hooks were incorrectly registered for MoE modules and NAS step was hanging before this. - Fix TRT support for remote autotuning in ONNX Autotune from 10.16+ to 10.15+ and fix TRT versioning check to the ``trtexec`` version instead of the TRT Python API when using ``trtexec`` backend. - Exclude MatMul/Gemm nodes with K or N < 16 from ONNX INT8 and FP8 quantization. Such small-dimension GEMMs cannot efficiently use INT8/FP8 Tensor Cores and the added Q/DQ layers cause perf regressions in TensorRT. Honors Gemm ``transB`` when deriving K. +- Fix ``nvfp4_awq`` export ``AssertionError: Modules have different quantization formats`` for MoE models (e.g. Qwen3-30B-A3B) when some experts are not exercised by the calibration data. ``awq_lite`` now applies a neutral all-ones ``pre_quant_scale`` to any expert that ends up disabled (no cache-pass tokens, NaN scales, or no search-pass tokens) so its format remains ``nvfp4_awq``, consistent with the rest of the MoE block. A warning is emitted whenever this fallback fires. **Misc** From 1f1aa855befbb97009ef389ddefacecc0fce5050 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 27 Apr 2026 16:36:51 +0000 Subject: [PATCH 3/3] [NVBUG: 6103846] Drop redundant num_cache_steps==0 check Signed-off-by: Chenjie Luo --- modelopt/torch/quantization/model_calib.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 67c63a99ada..0c2033041d6 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1280,12 +1280,13 @@ def postprocess(module, name): " that can be used to forward data through the model many times." ) - if module.awq_lite.num_cache_steps == 0 or not module.awq_lite.is_enabled: - # Expert was either uncalibrated (no cache-pass tokens), had NaN - # in act/weight scales, or saw no search-pass tokens. Max-calibrate - # weights and apply a neutral (all-ones) pre_quant_scale so the - # exporter sees a consistent nvfp4_awq format across all expert - # linears in an MoE group. + if not module.awq_lite.is_enabled: + # Expert is disabled — uncalibrated (no cache-pass tokens, set + # at the pre-search pass above), had NaN in act/weight scales, + # or saw no search-pass tokens. Max-calibrate weights and apply + # a neutral (all-ones) pre_quant_scale so the exporter sees a + # consistent nvfp4_awq format across all expert linears in an + # MoE group. # NOTE: ones-scale must be registered OUTSIDE enable_weight_access_and_writeback # because HF accelerate post_forward drops newly-registered submodule buffers. warnings.warn(