Coverage for transformer_lens/model_bridge/generalized_components/alibi_joint_qkv_attention.py: 84%
71 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""ALiBi joint QKV attention bridge component.
3Handles models that use ALiBi (Attention with Linear Biases) with fused QKV projections.
4Splits fused QKV, reimplements attention with ALiBi bias and hooks at each stage.
5"""
7from __future__ import annotations
9from typing import Any, Dict, Optional
11import torch
13from transformer_lens.model_bridge.generalized_components.alibi_utils import (
14 build_alibi_tensor,
15)
16from transformer_lens.model_bridge.generalized_components.base import (
17 GeneralizedComponent,
18)
19from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
20 JointQKVAttentionBridge,
21)
24class ALiBiJointQKVAttentionBridge(JointQKVAttentionBridge):
25 """Attention bridge for models using ALiBi position encoding with fused QKV.
27 Splits fused QKV, reimplements attention with ALiBi bias fused into scores,
28 and fires hooks at each stage (hook_q, hook_k, hook_v, hook_attn_scores,
29 hook_pattern). ALiBi bias is added to raw attention scores before scaling.
30 """
32 def __init__(
33 self,
34 name: str,
35 config: Any,
36 split_qkv_matrix: Any = None,
37 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
38 **kwargs: Any,
39 ):
40 super().__init__(
41 name=name,
42 config=config,
43 split_qkv_matrix=split_qkv_matrix,
44 submodules=submodules,
45 requires_position_embeddings=False,
46 requires_attention_mask=False,
47 **kwargs,
48 )
50 def forward(self, *args: Any, **kwargs: Any) -> Any:
51 """Forward pass: split QKV, apply ALiBi, fire hooks."""
52 if self.original_component is None: 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true
53 raise RuntimeError(
54 f"Original component not set for {self.name}. "
55 "Call set_original_component() first."
56 )
58 if len(args) > 0 and isinstance(args[0], torch.Tensor): 58 ↛ 60line 58 didn't jump to line 60 because the condition on line 58 was always true
59 hidden_states = args[0]
60 elif "hidden_states" in kwargs:
61 hidden_states = kwargs.pop("hidden_states")
62 else:
63 raise ValueError("Could not find hidden_states in args or kwargs")
65 hooked_input = self.hook_in(hidden_states)
67 # Split fused QKV via parent's split mechanism
68 q_output = self.q(hooked_input)
69 k_output = self.k(hooked_input)
70 v_output = self.v(hooked_input)
72 attn_output, attn_weights = self._reconstruct_attention(
73 q_output, k_output, v_output, **kwargs
74 )
76 output = self.hook_out(attn_output)
77 return output, attn_weights
79 def _reconstruct_attention(
80 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any
81 ) -> tuple[torch.Tensor, torch.Tensor]:
82 """Reconstruct attention with ALiBi bias fused into scores."""
83 num_heads = self.config.n_heads if self.config else 32
84 num_kv_heads = getattr(self.config, "n_key_value_heads", None) or num_heads
85 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(
86 q, k, v, num_heads, num_kv_heads
87 )
89 # Without this, cached decoding attends to the current token alone.
90 # HF caches pre-expansion K/V, so update before the GQA expand; the
91 # ALiBi bias and mask below arrive sized to the full key length.
92 k, v = self._update_kv_cache(k, v, **kwargs)
94 # GQA/MQA: expand K/V heads to match Q heads
95 if num_kv_heads != num_heads:
96 n_rep = num_heads // num_kv_heads
97 k = k.repeat_interleave(n_rep, dim=1)
98 v = v.repeat_interleave(n_rep, dim=1)
100 inv_norm_factor = head_dim**-0.5
102 # Raw attention scores
103 attn_scores = torch.matmul(q, k.transpose(-2, -1))
104 attn_scores = attn_scores.view(batch_size, num_heads, seq_len, -1)
106 # Upcast for numerical stability
107 input_dtype = attn_scores.dtype
108 if input_dtype in (torch.float16, torch.bfloat16): 108 ↛ 109line 108 didn't jump to line 109 because the condition on line 108 was never true
109 attn_scores = attn_scores.to(torch.float32)
111 # Add ALiBi bias
112 alibi = kwargs.get("alibi", None)
113 if alibi is not None:
114 kv_len = attn_scores.shape[-1]
115 alibi_view = alibi.view(batch_size, num_heads, 1, -1)
116 if alibi_view.shape[-1] > kv_len: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true
117 alibi_view = alibi_view[..., :kv_len]
118 attn_scores = attn_scores + alibi_view
120 # Scale after ALiBi (matches HF Falcon)
121 attn_scores = attn_scores * inv_norm_factor
123 # Add attention mask
124 attention_mask = kwargs.get("attention_mask", None)
125 if attention_mask is not None: 125 ↛ 129line 125 didn't jump to line 129 because the condition on line 125 was always true
126 attention_mask = self._normalize_compatibility_mask_sentinel(attention_mask)
127 attn_scores = attn_scores + attention_mask[:, :, :, : attn_scores.shape[-1]]
129 attn_scores = self.hook_attn_scores(attn_scores)
131 attn_weights = self._softmax_dropout_pattern(
132 attn_scores, upcast_to_fp32=True, target_dtype=q.dtype
133 )
135 # Weighted sum
136 attn_output = torch.matmul(
137 attn_weights.view(batch_size * num_heads, seq_len, -1),
138 v.reshape(batch_size * num_heads, -1, head_dim),
139 )
140 attn_output = attn_output.view(batch_size, num_heads, seq_len, head_dim)
141 attn_output = self._reshape_attn_output(
142 attn_output, batch_size, seq_len, num_heads, head_dim
143 )
144 attn_output = self._apply_output_projection(attn_output)
146 return attn_output, attn_weights
148 def get_random_inputs(
149 self,
150 batch_size: int = 2,
151 seq_len: int = 8,
152 device: Optional[torch.device] = None,
153 dtype: Optional[torch.dtype] = None,
154 ) -> Dict[str, Any]:
155 """Generate test inputs including ALiBi tensor and attention mask."""
156 if device is None: 156 ↛ 158line 156 didn't jump to line 158 because the condition on line 156 was always true
157 device = torch.device("cpu")
158 if dtype is None: 158 ↛ 161line 158 didn't jump to line 161 because the condition on line 158 was always true
159 dtype = torch.float32
161 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 2048
162 num_heads = self.config.n_heads if self.config and hasattr(self.config, "n_heads") else 32
164 attention_mask = torch.ones(batch_size, seq_len, device=device, dtype=torch.long)
165 # HF Falcon passes alibi as [batch*heads, 1, seq] — reshape to match
166 alibi_4d = build_alibi_tensor(attention_mask, num_heads, dtype)
167 alibi = alibi_4d.reshape(batch_size * num_heads, 1, seq_len)
169 # Causal mask: [batch, 1, seq, seq]
170 causal_mask = torch.triu(
171 torch.full((seq_len, seq_len), torch.finfo(dtype).min, device=device, dtype=dtype),
172 diagonal=1,
173 )
174 causal_mask = causal_mask.unsqueeze(0).unsqueeze(0).expand(batch_size, 1, -1, -1)
176 return {
177 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype),
178 "alibi": alibi,
179 "attention_mask": causal_mask,
180 }