diff --git a/CHANGELOG.rst b/CHANGELOG.rst index eeb20408d31..6fefe7c877b 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 d9a34289755..7633522c3e4 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 69c56ba8770..23bfa0dbe7b 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)