From 02c9a50c6bd1c3edaf4597ec5ae73e84c5880d7f 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 eeb20408d313..6fefe7c877ba 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -97,6 +97,7 @@ Changelog - Layerwise-calibration resume (``checkpoint_dir``) no longer runs code from a modified checkpoint directory; an unusable file now fails with an error instead. - 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. - ``mtq.fold_weight`` now folds sequential weight quantizers (e.g. W4A8 INT4 + FP8 weights) into the weights; it used to leave them unfolded and enabled. - 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``. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index d9a342897553..7633522c3e4c 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1738,8 +1738,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 = x_max.pow(alpha) @@ -1759,7 +1759,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): @@ -1796,9 +1796,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}] @@ -1868,7 +1868,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)) if use_weight_scale: diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 69c56ba87706..23bfa0dbe7b0 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -360,8 +360,9 @@ def test_awq_full(use_weight_scale): 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() @@ -369,8 +370,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)