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)