From 543083592381420d7caabbe7c8b61f6995e7b636 Mon Sep 17 00:00:00 2001 From: Abhinav Bellapu Date: Fri, 10 Jul 2026 13:29:22 -0700 Subject: [PATCH 1/3] Fix 4-bit split QKV projection --- tests/unit/components/test_attention.py | 74 ++++++++++++++++- .../components/abstract_attention.py | 81 +++++++++---------- 2 files changed, 112 insertions(+), 43 deletions(-) diff --git a/tests/unit/components/test_attention.py b/tests/unit/components/test_attention.py index 683924c44..697f6211e 100644 --- a/tests/unit/components/test_attention.py +++ b/tests/unit/components/test_attention.py @@ -5,13 +5,34 @@ from transformers.utils import is_bitsandbytes_available from transformer_lens.components import Attention +from transformer_lens.components import abstract_attention as abstract_attention_module from transformer_lens.config import HookedTransformerConfig -from transformer_lens.utilities.attention import complex_attn_linear +from transformer_lens.utilities.attention import complex_attn_linear, simple_attn_linear if is_bitsandbytes_available(): from bitsandbytes.nn.modules import Params4bit +class FakeParams4bit: + def __init__(self, projection: torch.Tensor): + self.projection = projection + self.quant_state = object() + + def t(self): + return self.projection + + +class FakeBnb: + @staticmethod + def matmul_4bit(input, weight, bias, quant_state): + return torch.matmul(input, weight) + + +def fake_4bit_weight(weight: torch.Tensor) -> FakeParams4bit: + projection = einops.rearrange(weight, "head d_model d_head -> d_model (head d_head)") + return FakeParams4bit(projection) + + def test_attention_hooked_transformer_config(): cfg = HookedTransformerConfig( n_layers=12, @@ -45,6 +66,57 @@ def test_attention_hooked_transformer_config(): assert torch.all(attn.b_V == 0) +@pytest.mark.parametrize("use_split_qkv_input", [False, True]) +def test_attention_4bit_qkv_projection_matches_unquantized(monkeypatch, use_split_qkv_input): + monkeypatch.setattr(abstract_attention_module, "Params4bit", FakeParams4bit, raising=False) + monkeypatch.setattr(abstract_attention_module, "bnb", FakeBnb, raising=False) + + torch.manual_seed(0) + cfg = HookedTransformerConfig( + n_layers=1, + d_model=6, + n_ctx=4, + d_head=2, + n_heads=3, + load_in_4bit=False, + use_split_qkv_input=use_split_qkv_input, + dtype=torch.float32, + act_fn="relu", + ) + attn = Attention(cfg) + attn.cfg.load_in_4bit = True + + W_Q = torch.randn(cfg.n_heads, cfg.d_model, cfg.d_head) + W_K = torch.randn(cfg.n_heads, cfg.d_model, cfg.d_head) + W_V = torch.randn(cfg.n_heads, cfg.d_model, cfg.d_head) + + for name, weight in (("W_Q", W_Q), ("W_K", W_K), ("W_V", W_V)): + del attn._parameters[name] + setattr(attn, name, fake_4bit_weight(weight)) + + with torch.no_grad(): + attn.b_Q.copy_(torch.randn_like(attn.b_Q)) + attn.b_K.copy_(torch.randn_like(attn.b_K)) + attn.b_V.copy_(torch.randn_like(attn.b_V)) + + if use_split_qkv_input: + query_input = torch.randn(2, 4, cfg.n_heads, cfg.d_model) + key_input = torch.randn(2, 4, cfg.n_heads, cfg.d_model) + value_input = torch.randn(2, 4, cfg.n_heads, cfg.d_model) + expected_fn = complex_attn_linear + else: + query_input = torch.randn(2, 4, cfg.d_model) + key_input = torch.randn(2, 4, cfg.d_model) + value_input = torch.randn(2, 4, cfg.d_model) + expected_fn = simple_attn_linear + + q, k, v = attn.calculate_qkv_matrices(query_input, key_input, value_input) + + assert torch.allclose(q, expected_fn(query_input, W_Q, attn.b_Q)) + assert torch.allclose(k, expected_fn(key_input, W_K, attn.b_K)) + assert torch.allclose(v, expected_fn(value_input, W_V, attn.b_V)) + + @pytest.mark.skipif(not is_bitsandbytes_available(), reason="bitsandbytes is not available") def test_attention_load_in_4bit(): cfg = HookedTransformerConfig( diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 01a0646eb..242fa3548 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/transformer_lens/components/abstract_attention.py @@ -446,59 +446,20 @@ def calculate_qkv_matrices( ) if self.cfg.load_in_4bit: W_Q_4bit = cast(Params4bit, self.W_Q) - q = self.hook_q( - # call bitsandbytes method to dequantize and multiply - bnb.matmul_4bit( - query_input, - W_Q_4bit.t(), - bias=None, - quant_state=W_Q_4bit.quant_state, - ).reshape( - query_input.shape[0], - query_input.shape[1], - self.cfg.n_heads, - self.cfg.d_head, - ) - + self.b_Q - ) + q = self.hook_q(self._project_4bit_qkv(query_input, W_Q_4bit, self.b_Q)) else: q = self.hook_q(attn_fn(query_input, self.W_Q, self.b_Q)) if self.cfg.load_in_4bit: if not isinstance(self.W_K, Params4bit): raise ValueError("W_K must be a Params4bit object if load_in_4bit is True") - k = self.hook_k( - # call bitsandbytes method to dequantize and multiply - bnb.matmul_4bit( - key_input, self.W_K.t(), bias=None, quant_state=self.W_K.quant_state - ).reshape( - key_input.shape[0], - key_input.shape[1], - self.cfg.n_heads, - self.cfg.d_head, - ) - + self.b_K - ) + k = self.hook_k(self._project_4bit_qkv(key_input, self.W_K, self.b_K)) else: k = self.hook_k(attn_fn(key_input, self.W_K, self.b_K)) if self.cfg.load_in_4bit: if not isinstance(self.W_V, Params4bit): raise ValueError("W_V must be a Params4bit object if load_in_4bit is True") - v = self.hook_v( - # call bitsandbytes method to dequantize and multiply - bnb.matmul_4bit( - value_input, - self.W_V.t(), - bias=None, - quant_state=self.W_V.quant_state, - ).reshape( - value_input.shape[0], - value_input.shape[1], - self.cfg.n_heads, - self.cfg.d_head, - ) - + self.b_V - ) + v = self.hook_v(self._project_4bit_qkv(value_input, self.W_V, self.b_V)) else: v = self.hook_v(attn_fn(value_input, self.W_V, self.b_V)) @@ -510,6 +471,42 @@ def calculate_qkv_matrices( return q, k, v + def _project_4bit_qkv( + self, + input: Union[ + Float[torch.Tensor, "batch pos d_model"], + Float[torch.Tensor, "batch pos head_index d_model"], + ], + weight: "Params4bit", + bias: Float[torch.Tensor, "head_index d_head"], + ) -> Float[torch.Tensor, "batch pos head_index d_head"]: + projected = bnb.matmul_4bit( + input, + weight.t(), + bias=None, + quant_state=weight.quant_state, + ) + n_heads, d_head = bias.shape + + if input.ndim == 3: + return projected.reshape(input.shape[0], input.shape[1], n_heads, d_head) + bias + + if input.ndim == 4: + if input.shape[2] != n_heads: + raise ValueError( + "4-bit split QKV inputs must have one input slice per attention head; " + f"got {input.shape[2]} input heads for {n_heads} projection heads." + ) + split_projected = projected.reshape( + input.shape[0], input.shape[1], input.shape[2], n_heads, d_head + ) + return split_projected.diagonal(dim1=2, dim2=3).movedim(-1, 2) + bias + + raise ValueError( + "4-bit QKV projection input must have shape [batch, pos, d_model] or " + "[batch, pos, head_index, d_model]." + ) + def calculate_attention_scores( self, q: Float[torch.Tensor, "batch query_pos head_index d_head"], From 5eb3c5a2394daf3c60f4629958dad7342c5c0e50 Mon Sep 17 00:00:00 2001 From: Abhinav Bellapu Date: Fri, 10 Jul 2026 14:28:11 -0700 Subject: [PATCH 2/3] Retry CI From d3ed3dd2b96469d6130f9c4134ad0a90e0223928 Mon Sep 17 00:00:00 2001 From: Abhinav Bellapu Date: Tue, 14 Jul 2026 15:05:05 -0700 Subject: [PATCH 3/3] Address 4-bit QKV review feedback --- tests/QUARANTINES.md | 4 +- tests/unit/components/test_attention.py | 20 +++++++--- .../components/abstract_attention.py | 40 +++++++++++-------- 3 files changed, 41 insertions(+), 23 deletions(-) diff --git a/tests/QUARANTINES.md b/tests/QUARANTINES.md index 5a587ea38..66c32e4ad 100644 --- a/tests/QUARANTINES.md +++ b/tests/QUARANTINES.md @@ -11,7 +11,7 @@ Rule ([AGENTS.md §10](../AGENTS.md#10-hard-rules)): **never add `xfail` / `skip | Path | Marker | Trigger | |---|---|---| | [`unit/test_lit.py`](unit/test_lit.py) (×18) | `skipif(not LIT_AVAILABLE)` | `pip install lit-nlp` (`lit` group) | -| [`unit/components/test_attention.py`:48](unit/components/test_attention.py) | `skipif(not is_bitsandbytes_available())` | `uv sync --group quantization` | +| [`unit/components/test_attention.py`:130](unit/components/test_attention.py) | `skipif(not is_bitsandbytes_available())` | `uv sync --group quantization` | | [`unit/test_weight_processing.py`:477](unit/test_weight_processing.py) | same | same | | [`unit/factories/test_mlp_factory.py`:40](unit/factories/test_mlp_factory.py) | same | same | @@ -25,7 +25,7 @@ Rule ([AGENTS.md §10](../AGENTS.md#10-hard-rules)): **never add `xfail` / `skip |---|---|---| | [`unit/test_next_sentence_prediction.py`:131](unit/test_next_sentence_prediction.py) | `skipif(not cuda)` | Any CUDA | | [`unit/model_bridge/compatibility/test_next_sentence_prediction.py`:95](unit/model_bridge/compatibility/test_next_sentence_prediction.py) | `skipif(not cuda)` | Any CUDA | -| [`unit/components/test_attention.py`:83](unit/components/test_attention.py) | `skipif(not cuda)` (half/bfloat16) | Any CUDA | +| [`unit/components/test_attention.py`:165](unit/components/test_attention.py) | `skipif(not cuda)` (half/bfloat16) | Any CUDA | | [`acceptance/test_hooked_encoder.py`:227](acceptance/test_hooked_encoder.py) | `skipif(not cuda)` | Any CUDA | | [`acceptance/test_hooked_encoder_decoder.py`:421](acceptance/test_hooked_encoder_decoder.py) | `skipif(not cuda)` | Any CUDA | | [`acceptance/test_multi_gpu.py`:91,105](acceptance/test_multi_gpu.py) | `skipif(device_count < 2)` | 2+ CUDA | diff --git a/tests/unit/components/test_attention.py b/tests/unit/components/test_attention.py index 697f6211e..e80bd937c 100644 --- a/tests/unit/components/test_attention.py +++ b/tests/unit/components/test_attention.py @@ -14,23 +14,33 @@ class FakeParams4bit: - def __init__(self, projection: torch.Tensor): - self.projection = projection + def __init__(self, dequantized: torch.Tensor): + self.data = dequantized self.quant_state = object() def t(self): - return self.projection + return self.data.t() + + +class FakeBnbFunctional: + @staticmethod + def dequantize_4bit(input, quant_state): + return input class FakeBnb: + functional = FakeBnbFunctional() + @staticmethod def matmul_4bit(input, weight, bias, quant_state): + if input.ndim != 3: + raise AssertionError("split QKV projection should dequantize once") return torch.matmul(input, weight) def fake_4bit_weight(weight: torch.Tensor) -> FakeParams4bit: - projection = einops.rearrange(weight, "head d_model d_head -> d_model (head d_head)") - return FakeParams4bit(projection) + dequantized = einops.rearrange(weight, "head d_model d_head -> (head d_head) d_model") + return FakeParams4bit(dequantized) def test_attention_hooked_transformer_config(): diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 242fa3548..f18f1f25d 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/transformer_lens/components/abstract_attention.py @@ -445,20 +445,15 @@ def calculate_qkv_matrices( else simple_attn_linear ) if self.cfg.load_in_4bit: - W_Q_4bit = cast(Params4bit, self.W_Q) - q = self.hook_q(self._project_4bit_qkv(query_input, W_Q_4bit, self.b_Q)) + q = self.hook_q(self._project_4bit_qkv(query_input, self.W_Q, self.b_Q)) else: q = self.hook_q(attn_fn(query_input, self.W_Q, self.b_Q)) if self.cfg.load_in_4bit: - if not isinstance(self.W_K, Params4bit): - raise ValueError("W_K must be a Params4bit object if load_in_4bit is True") k = self.hook_k(self._project_4bit_qkv(key_input, self.W_K, self.b_K)) else: k = self.hook_k(attn_fn(key_input, self.W_K, self.b_K)) if self.cfg.load_in_4bit: - if not isinstance(self.W_V, Params4bit): - raise ValueError("W_V must be a Params4bit object if load_in_4bit is True") v = self.hook_v(self._project_4bit_qkv(value_input, self.W_V, self.b_V)) else: v = self.hook_v(attn_fn(value_input, self.W_V, self.b_V)) @@ -477,18 +472,25 @@ def _project_4bit_qkv( Float[torch.Tensor, "batch pos d_model"], Float[torch.Tensor, "batch pos head_index d_model"], ], - weight: "Params4bit", + weight: Union[nn.Parameter, "Params4bit"], bias: Float[torch.Tensor, "head_index d_head"], ) -> Float[torch.Tensor, "batch pos head_index d_head"]: - projected = bnb.matmul_4bit( - input, - weight.t(), - bias=None, - quant_state=weight.quant_state, - ) + """Project Q/K/V inputs with a 4-bit weight. + + Split inputs dequantize once and reuse the head-wise projection path. + """ + if not isinstance(weight, Params4bit): + raise ValueError("QKV weights must be Params4bit objects if load_in_4bit is True") + n_heads, d_head = bias.shape if input.ndim == 3: + projected = bnb.matmul_4bit( + input, + weight.t(), + bias=None, + quant_state=weight.quant_state, + ) return projected.reshape(input.shape[0], input.shape[1], n_heads, d_head) + bias if input.ndim == 4: @@ -497,10 +499,16 @@ def _project_4bit_qkv( "4-bit split QKV inputs must have one input slice per attention head; " f"got {input.shape[2]} input heads for {n_heads} projection heads." ) - split_projected = projected.reshape( - input.shape[0], input.shape[1], input.shape[2], n_heads, d_head + dequantized_weight = bnb.functional.dequantize_4bit( + weight.data, quant_state=weight.quant_state + ) + split_weight = einops.rearrange( + dequantized_weight, + "(head_index d_head) d_model -> head_index d_model d_head", + head_index=n_heads, + d_head=d_head, ) - return split_projected.diagonal(dim1=2, dim2=3).movedim(-1, 2) + bias + return complex_attn_linear(input, split_weight, bias) raise ValueError( "4-bit QKV projection input must have shape [batch, pos, d_model] or "