From aab14ded73b2ce1cf1f688102cf4083dd7dc3fe7 Mon Sep 17 00:00:00 2001 From: realAsma Date: Sun, 19 Jul 2026 01:31:23 +0000 Subject: [PATCH 1/2] fix(autoquant): score grouped QKV at attention output Signed-off-by: realAsma --- modelopt/torch/quantization/algorithms.py | 4 +- .../unit/torch/quantization/test_autoquant.py | 80 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index c7803ecf838..5b8d65680a5 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -445,6 +445,7 @@ def attrs(self) -> list[str]: _LINEAR_ATTN_QKVZ_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_qkv|in_proj_z)$") _LINEAR_ATTN_BA_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_a|in_proj_b)$") +_ATTN_QKV_RULE = r"^(.*?)\.(q_proj|k_proj|v_proj)$" def _linear_attn_qkvz_group_key(_model, name: str) -> str | None: @@ -471,7 +472,7 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC): method_name: str | None = None quant_grouping_rules = [ - r"^(.*?)\.(q_proj|k_proj|v_proj)$", # q_proj, k_proj, v_proj for llama like models + _ATTN_QKV_RULE, # q_proj, k_proj, v_proj for llama like models # gate_proj, up_proj, down_proj for Qwen3 like MoE models r"^(.*?\.mlp\.experts)\.\d+\.(gate_proj|up_proj|down_proj)$", r"^(.*?\.mixer\.experts)\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts @@ -1025,6 +1026,7 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher): method_name = "gradient" score_module_rules = [ + _ATTN_QKV_RULE, # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts) r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$", r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$", # NemotronH MoE experts diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index b85feb32649..0f268ebb310 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -56,6 +56,37 @@ def forward(self, x): return x +class _ParallelAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_proj = torch.nn.Linear(4, 4, bias=False) + self.k_proj = torch.nn.Linear(4, 4, bias=False) + self.v_proj = torch.nn.Linear(4, 4, bias=False) + + weight = torch.tensor( + [ + [0.11, -0.37, 0.53, 0.89], + [-0.29, 0.47, 0.71, -0.97], + [0.19, 0.41, -0.67, 0.83], + [0.31, -0.59, 0.73, -0.91], + ] + ) + for projection in (self.q_proj, self.k_proj, self.v_proj): + projection.weight.data.copy_(weight) + + def forward(self, x): + return self.q_proj(x) + self.k_proj(x) + self.v_proj(x) + + +class _ParallelAttentionModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.attn = _ParallelAttention() + + def forward(self, x): + return self.attn(x) + + class TransformerBlock(torch.nn.Module): def __init__(self): super().__init__() @@ -419,6 +450,55 @@ def test_auto_quantize_disabled_layers_no_poison(): assert QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG) in hparam.choices +def test_auto_quantize_qkv_score_uses_attention_output(): + model = _ParallelAttentionModel() + data = torch.tensor([[[0.13, -0.31, 0.79, -0.43], [0.61, 0.17, -0.23, 0.97]]]) + output_grad = torch.tensor([[[0.7, -1.1, 0.3, 1.3], [-0.5, 0.9, 1.7, -0.2]]]) + + model, search_history = mtq.auto_quantize( + model, + constraints={"effective_bits": 16.0}, + quantization_formats=[mtq.INT8_DEFAULT_CFG], + data_loader=[data], + forward_step=lambda model, batch: model(batch), + loss_func=lambda output, batch: (output * output_grad).sum(), + num_calib_steps=1, + num_score_steps=1, + ) + + hparam = model.attn.q_proj.get_hparam("quant_recipe") + assert model.attn.k_proj.get_hparam("quant_recipe") is hparam + assert model.attn.v_proj.get_hparam("quant_recipe") is hparam + assert hparam.score_modules == [model.attn] + + no_quant_recipe = QuantRecipe(None) + int8_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG) + hparam.active = no_quant_recipe + attention_output = model.attn(data) + projection_outputs = [projection(data) for projection in hparam.quant_modules] + hparam.active = int8_recipe + quantized_attention_output = model.attn(data) + quantized_projection_outputs = [projection(data) for projection in hparam.quant_modules] + + expected_score = ((output_grad * (quantized_attention_output - attention_output)) ** 2).sum() + old_projection_score = sum( + ((output_grad * (quantized - output)) ** 2).sum() + for output, quantized in zip(projection_outputs, quantized_projection_outputs) + ) + qkv_stats = next( + stats + for stats in search_history["candidate_stats"].values() + if set(stats["module_names"]) == {"attn.q_proj", "attn.k_proj", "attn.v_proj"} + ) + actual_score = qkv_stats["scores"][qkv_stats["formats"].index(int8_recipe)] + + assert actual_score == pytest.approx(expected_score.item()) + assert actual_score != pytest.approx(old_projection_score.item()) + assert qkv_stats["costs"][qkv_stats["formats"].index(no_quant_recipe)] == sum( + projection.weight.numel() for projection in hparam.quant_modules + ) + + INT4INT8_AWQ_CFG = { "quant_cfg": [ {"quantizer_name": "*", "enable": False}, From 73ad76f6f0fee8a18dbc9ca9168694bac670601f Mon Sep 17 00:00:00 2001 From: realAsma Date: Sun, 19 Jul 2026 12:13:05 +0000 Subject: [PATCH 2/2] fix(autoquant): support keyword-bound scoring inputs Signed-off-by: realAsma --- modelopt/torch/quantization/algorithms.py | 6 +++--- tests/unit/torch/quantization/test_autoquant.py | 9 +++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/modelopt/torch/quantization/algorithms.py b/modelopt/torch/quantization/algorithms.py index 5b8d65680a5..f98e12b7708 100644 --- a/modelopt/torch/quantization/algorithms.py +++ b/modelopt/torch/quantization/algorithms.py @@ -1107,12 +1107,12 @@ def forward_backward_step(model, data): @torch.enable_grad() def _estimate_auto_quantize_scores(self, is_param_grad_enabled): # TODO: remove the no-quant recipe - def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): + def auto_quantize_score_estimate_forward(module, *args, **kwargs): for hparam in module._hparams_for_scoring: if hparam.is_configurable: hparam.active = QuantRecipe(quant_cfg=None) - output = module._forward_original(input, *args, **kwargs) + output = module._forward_original(*args, **kwargs) # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass. # With gradient checkpointing, gradients are computed in the local forward pass during backward pass @@ -1130,7 +1130,7 @@ def auto_quantize_score_estimate_forward(module, input, *args, **kwargs): if recipe == QuantRecipe(quant_cfg=None): continue hparam.active = recipe - output_diff = module._forward_original(input, *args, **kwargs) + output_diff = module._forward_original(*args, **kwargs) if isinstance(output_diff, tuple): output_diff = output_diff[0] - output[0] diff --git a/tests/unit/torch/quantization/test_autoquant.py b/tests/unit/torch/quantization/test_autoquant.py index 0f268ebb310..79841f7cb2a 100644 --- a/tests/unit/torch/quantization/test_autoquant.py +++ b/tests/unit/torch/quantization/test_autoquant.py @@ -83,8 +83,8 @@ def __init__(self): super().__init__() self.attn = _ParallelAttention() - def forward(self, x): - return self.attn(x) + def forward(self, x, input_as_keyword=False): + return self.attn(x=x) if input_as_keyword else self.attn(x) class TransformerBlock(torch.nn.Module): @@ -450,7 +450,8 @@ def test_auto_quantize_disabled_layers_no_poison(): assert QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG) in hparam.choices -def test_auto_quantize_qkv_score_uses_attention_output(): +@pytest.mark.parametrize("input_as_keyword", [False, True]) +def test_auto_quantize_qkv_score_uses_attention_output(input_as_keyword): model = _ParallelAttentionModel() data = torch.tensor([[[0.13, -0.31, 0.79, -0.43], [0.61, 0.17, -0.23, 0.97]]]) output_grad = torch.tensor([[[0.7, -1.1, 0.3, 1.3], [-0.5, 0.9, 1.7, -0.2]]]) @@ -460,7 +461,7 @@ def test_auto_quantize_qkv_score_uses_attention_output(): constraints={"effective_bits": 16.0}, quantization_formats=[mtq.INT8_DEFAULT_CFG], data_loader=[data], - forward_step=lambda model, batch: model(batch), + forward_step=lambda model, batch: model(batch, input_as_keyword=input_as_keyword), loss_func=lambda output, batch: (output * output_grad).sum(), num_calib_steps=1, num_score_steps=1,