From 0826672e7efdbce4218832c96a5404cd8fb9bc5c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 06:12:29 +0000 Subject: [PATCH] Make awq_lite independent of how the calibration tokens are batched awq_lite averaged |x| per batch in the activation dtype, then averaged those batch means, and summed the per-batch mean squared error of the alpha search. Every batch therefore weighed the same whatever its token count, so a smaller last batch, or an MoE expert that gets few tokens in some batches, changed the scales, and each batch mean was rounded to bf16. It now sums |x| in fp32 and divides by the token count, and sums the squared error, as the reference AWQ does over the whole calibration set. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + modelopt/torch/quantization/model_calib.py | 12 ++++++------ tests/unit/torch/quantization/test_calib.py | 10 ++++++---- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d6397b993a0..f09fd8ce6fd 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -45,6 +45,7 @@ Changelog - The default PTQ recipes now leave multi-token-prediction (MTP) layers unquantized in Megatron vision-language models too (``language_model.mtp.*``), as they already did at the model root. - Fix GPT-OSS draft training with pre-normalization teacher hidden states by enabling the target model's final normalization. +- ``awq_lite`` now weights every calibration token equally in its activation statistics and its alpha search, as the reference AWQ does, so the scales no longer depend on how the tokens are split into batches (a smaller last batch, MoE experts that see few tokens in some batches). AWQ scales calibrated with uneven batches can differ from earlier releases. - Fix vanilla DSpark exports to use the Qwen3 DSpark architecture, head names and configuration expected by vLLM. Re-export existing training checkpoints to use this format. - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index cf14a5b8c3c..741d38d2d85 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1526,8 +1526,8 @@ def get_weight_scale(weight, block_size=None): scale = scale.mean(0).to(torch.float32) return scale - def get_act_scale(x): - return x.abs().contiguous().view(-1, x.shape[-1]).mean(0).to(torch.float32) + def get_act_scale_sum(x): + return x.abs().reshape(-1, x.shape[-1]).sum(0, dtype=torch.float32) def get_scale(x_max, w_max, alpha, tensor_parallel_group=None): scales = ( @@ -1549,7 +1549,7 @@ def update_loss(self, out, out_actual, alpha): out = out[0] if isinstance(out, tuple) else out out = out.to_local() if hasattr(out, "to_local") else out out_actual = out_actual.to_local() if hasattr(out_actual, "to_local") else out_actual - loss = (out - out_actual).float().pow(2).mean() + loss = (out - out_actual).float().pow(2).sum() self.awq_lite.loss[alpha] += loss.to(self.awq_lite.loss[alpha].device) def update_best_params(self): @@ -1586,9 +1586,9 @@ def forward(self, input, *args, **kwargs): # Get local tensor from Dtensor input = input.to_local() if hasattr(input, "to_local") else input - self.awq_lite.act_scale += get_act_scale(self.input_quantizer(input)) + self.awq_lite.act_scale += get_act_scale_sum(self.input_quantizer(input)) self.awq_lite.num_cache_steps += 1 - self.awq_lite.num_tokens += input.numel() / input.shape[-1] + self.awq_lite.num_tokens += input.numel() // input.shape[-1] if self.awq_lite.is_input_quantized: with set_quantizer_by_cfg_context( self.input_quantizer, [{"quantizer_name": "*", "enable": True}] @@ -1658,7 +1658,7 @@ def sync_act_scale_across_dp(module, data_parallel_group): ): # Hack: MoEs forward all tokens through all experts if _if_calib is True module._if_calib = True - module.awq_lite.act_scale = module.awq_lite.act_scale / module.awq_lite.num_cache_steps + module.awq_lite.act_scale = module.awq_lite.act_scale / module.awq_lite.num_tokens has_nan_local = torch.any(torch.isnan(module.awq_lite.act_scale)) or torch.any( torch.isnan(module.awq_lite.weight_scale) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index aa31d0b66d6..0578d4ac043 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -282,8 +282,9 @@ def test_awq_full(): assert torch.allclose(out, out_ref, atol=1e-6) -def test_awq_multi_batch(): - """Test awq multibatch.""" +@pytest.mark.parametrize("batch_sizes", [(2, 2, 2, 2), (1, 3, 4)]) +def test_awq_multi_batch(batch_sizes): + """The batch split must not change the awq result, also when the batches differ in size.""" # awq test model = _SimpleMLP() states = model.state_dict() @@ -291,8 +292,9 @@ def test_awq_multi_batch(): # Multi batch test config = get_awq_config("awq_full") - dataloader = [torch.randn(2, 16, 16) for _ in range(4)] - dataloader_sb = [torch.cat(dataloader, dim=0)] + data = torch.randn(sum(batch_sizes), 16, 16) + dataloader = list(data.split(batch_sizes)) + dataloader_sb = [data] model_s = _SimpleMLP() model_s.load_state_dict(states)