Fix 4-bit split QKV projection

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 |

tests/unit/components/test_attention.py

@@ -5,13 +5,44 @@
 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, dequantized: torch.Tensor):
+        self.data = dequantized
+        self.quant_state = object()
+
+    def t(self):
+        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:
+    dequantized = einops.rearrange(weight, "head d_model d_head -> (head d_head) d_model")
+    return FakeParams4bit(dequantized)
+
+
 def test_attention_hooked_transformer_config():
     cfg = HookedTransformerConfig(
         n_layers=12,
@@ -45,6 +76,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(

transformer_lens/components/abstract_attention.py

@@ -445,60 +445,16 @@ 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(
-                # 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, 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(
-                # 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 +466,55 @@ 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: Union[nn.Parameter, "Params4bit"],
+        bias: Float[torch.Tensor, "head_index d_head"],
+    ) -> Float[torch.Tensor, "batch pos head_index d_head"]:
+        """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:
+            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."
+                )
+            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 complex_attn_linear(input, split_weight, 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"],