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 683924c44..e80bd937c 100644 --- a/tests/unit/components/test_attention.py +++ b/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( diff --git a/transformer_lens/components/abstract_attention.py b/transformer_lens/components/abstract_attention.py index 01a0646eb..f18f1f25d 100644 --- a/transformer_lens/components/abstract_attention.py +++ b/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"],