[NVBUG: 6103846] Fix nvfp4_awq export for uncalibrated MoE experts

CHANGELOG.rst

@@ -47,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**
 

modelopt/torch/quantization/model_calib.py

@@ -1270,11 +1270,30 @@ 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 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(
+                    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 +1308,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)