diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d229b16c070..192adf0573e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,6 +52,7 @@ Changelog - ``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``. - 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. +- Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target. diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 2113edea8a7..e4843a3a2ec 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -52,7 +52,9 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: model_module = inspect.getmodule(attention_cls) head = ast.parse(source_code) - bmm_ops = ("matmul", "bmm", "baddbmm") + # No baddbmm: patch() wires args[0] and args[1], which miss its batch1 and batch2 operands + # (torch.baddbmm passes the additive input first, and Bloom passes them as keywords) + bmm_ops = ("matmul", "bmm") sdpa_ops = ("scaled_dot_product_attention",) def is_bmm(node): @@ -72,6 +74,12 @@ def is_sdpa(node): def is_bin_matmul(node): return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) + def walk_post_order(node): + # Source order, children first; expects the q/k score matmul before the matmul consuming it + for child in ast.iter_child_nodes(node): + yield from walk_post_order(child) + yield node + def patch(node, quantizer_names, transpose=False): for index, quantizer_name in enumerate(quantizer_names): if quantizer_name is None: @@ -165,7 +173,7 @@ def patch_binop(node, quantizer_names, transpose=False): bmm_nodes = [] sdpa_nodes = [] bin_matmul_nodes = [] - for node in ast.walk(head): + for node in walk_post_order(head): if is_bmm(node): bmm_nodes.append(node) if is_sdpa(node): @@ -186,16 +194,16 @@ def patch_binop(node, quantizer_names, transpose=False): # after transpose, the quantization will be per-token, i.e., # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) # removing the additional transpose is doable but not trivial - patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) - patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 BMM/Matmul operators with quantizers") if len(bin_matmul_nodes) == 2: patch_binop( - bin_matmul_nodes[1], + bin_matmul_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True, ) - patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) + patch_binop(bin_matmul_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 @ operators with quantizers") if len(sdpa_nodes) == 1: diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 22a86971a6d..ff9e9404ea4 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -52,6 +52,54 @@ def forward(self, hidden_states, **kwargs): return F.scaled_dot_product_attention(q, k, v), None +class FlatMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.matmul(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v) + + +class FlatBMMAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.bmm(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.bmm(probs, v) + + +class FlatBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = q @ k.transpose(-2, -1) + probs = torch.softmax(scores, dim=-1) + return probs @ v + + +class NestedMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.matmul(torch.softmax(torch.matmul(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBMMAttention(nn.Module): + def forward(self, q, k, v): + return torch.bmm(torch.softmax(torch.bmm(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v + + +class BAddBMMAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(bias, q, k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + +class BAddBMMKeywordAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(input=bias, batch1=q, batch2=k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + kv_cache_config = { "quant_cfg": [ {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": 4}, "enable": True}, @@ -115,6 +163,37 @@ def test_kv_quant_hf(model_getter, attn_cls): mtq.unregister(attn_cls) +@pytest.mark.parametrize( + "attn_cls", + [ + FlatMatmulAttention, + FlatBMMAttention, + FlatBinMatmulAttention, + NestedMatmulAttention, + NestedBMMAttention, + NestedBinMatmulAttention, + ], +) +def test_kv_quant_operand_wiring(attn_cls): + """Each q/k/v_bmm_quantizer must calibrate on its own operand, whatever the layout.""" + q, k, v = torch.randn(1, 4, 8), 10 * torch.randn(1, 4, 8), 100 * torch.randn(1, 4, 8) + assert mtq.plugins.register_attention_for_kv_quant(attn_cls) + try: + model = attn_cls() + mtq.quantize(model, kv_cache_config, lambda model: model(q, k, v)) + for name, operand in zip("qkv", (q, k, v)): + quantizer = getattr(model, f"{name}_bmm_quantizer") + assert quantizer.amax == operand.abs().max(), f"{name}_bmm_quantizer saw another tensor" + finally: + mtq.unregister(attn_cls) + + +@pytest.mark.parametrize("attn_cls", [BAddBMMAttention, BAddBMMKeywordAttention]) +def test_kv_quant_rejects_baddbmm(attn_cls): + """baddbmm takes the additive input first, so its operands cannot be wired by position.""" + assert not mtq.plugins.register_attention_for_kv_quant(attn_cls) + + def test_kv_quant_bert(): """Test KV cache quantization on BERT model with decorated attention.""" model_test = get_tiny_bert()