diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 4ce0f62a75d..aeae3dd4321 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1316,6 +1316,21 @@ def postprocess(module, name): dtype=w_dtype, device=w_device, ) + # Mirror the calibrated postprocess path, gated on + # is_input_quantized so weight-only AWQ configs (where + # setup() never disabled input_quantizer) stay untouched. + # Collapse any per-channel _amax left over from cache_mode + # max_calibrate into a per-tensor scalar so + # preprocess_linear_fusion's numel==1 assertion passes, and + # re-enable the quantizer (awq_lite.setup disabled it). + if module.awq_lite.is_input_quantized: + if module.input_quantizer.amax is not None: + act_amax = module.input_quantizer.amax + module.input_quantizer._amax_for_smoothing = act_amax.cpu() + module.input_quantizer.reset_amax() + module.input_quantizer.axis = None + module.input_quantizer.amax = act_amax.amax() + module.input_quantizer.enable() else: with enable_weight_access_and_writeback(module, model, name_to_module): postprocess(module, name) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index d2e6fdd03e8..a39ee55d9d2 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -312,6 +312,84 @@ def test_padded_awq(): model(torch.randn(2, 16, 16)) +class _TwoBranchModel(nn.Module): + """Two parallel linears; only the first is exercised by forward_loop.""" + + def __init__(self): + super().__init__() + self.calibrated = nn.Linear(16, 16, bias=False) + self.uncalibrated = nn.Linear(16, 16, bias=False) + + def forward(self, x, branch="calibrated"): + if branch == "calibrated": + return self.calibrated(x) + return self.uncalibrated(x) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 dynamic block quant is CUDA-only") +def test_awq_lite_uncalibrated_linear_keeps_input_quantizer_enabled(): + """Regression test for NVBug 6143871. + + awq_lite.setup() disables the input_quantizer at the start of search. The + calibrated branch re-enables it inside postprocess(); the uncalibrated + branch (no cache-pass tokens, e.g. an MoE expert that never gets routed) + must do the same — otherwise downstream export (set_expert_quantizer_amax + + _export_quantized_weight) drops the input_scale buffer and inference + runtimes that read per-expert input_scale (e.g. TRT-LLM CutlassFusedMoE) + crash with KeyError on '.w1.input_scale'. + + Also asserts the export-critical scalar amax invariant (axis=None, + numel==1) — preprocess_linear_fusion enforces it for fused-expert groups. + """ + torch.manual_seed(0) + model = _TwoBranchModel().cuda() + + def _forward_loop(m): + for _ in range(2): + m(torch.randn(2, 16, 16, device="cuda"), branch="calibrated") + + mtq.quantize(model, mtq.NVFP4_AWQ_LITE_CFG, _forward_loop) + + assert model.calibrated.input_quantizer.is_enabled + assert model.uncalibrated.input_quantizer.is_enabled, ( + "Uncalibrated linear's input_quantizer must remain enabled after " + "awq_lite postprocess so export emits input_scale (NVBug 6143871)." + ) + uncal_q = model.uncalibrated.input_quantizer + # When amax exists (cache-hit but search-miss path), it must be the + # scalar form export expects — preprocess_linear_fusion asserts numel==1. + # When it's None (truly never routed), set_expert_quantizer_amax will + # populate it during export. + if uncal_q.amax is not None: + assert uncal_q.axis is None + assert uncal_q.amax.numel() == 1 + + +def test_awq_lite_uncalibrated_weight_only_keeps_input_quantizer_disabled(): + """Weight-only AWQ companion to NVBug 6143871. + + For weight-only AWQ configs (input_quantizer disabled), awq_lite.setup() + never touches the input_quantizer, so the postprocess uncalibrated branch + must NOT enable it — doing so turns on quantization the user's config had + explicitly opted out of. + """ + torch.manual_seed(0) + model = _TwoBranchModel() + + def _forward_loop(m): + for _ in range(2): + m(torch.randn(2, 16, 16), branch="calibrated") + + mtq.quantize(model, mtq.INT4_AWQ_CFG, _forward_loop) + + assert not model.calibrated.input_quantizer.is_enabled + assert not model.uncalibrated.input_quantizer.is_enabled, ( + "Weight-only AWQ must not flip on the input_quantizer for " + "uncalibrated layers — that would silently quantize activations " + "the user's config left in full precision." + ) + + def test_smoothquant_enable_disable(): torch.manual_seed(1234) model = _SimpleMLP()