Coverage for transformer_lens/model_bridge/generalized_components/bloom_attention.py: 87%
106 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""BLOOM-specific attention bridge component.
3BLOOM attention requires special arguments (residual, alibi, attention_mask) that standard
4JointQKVAttentionBridge doesn't handle. This custom component passes these arguments through.
5"""
6from typing import Any, Callable, Dict, Mapping, Optional
8import torch
10from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
11 BaseTensorConversion,
12)
13from transformer_lens.model_bridge.generalized_components.base import (
14 GeneralizedComponent,
15)
16from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
17 JointQKVAttentionBridge,
18)
21class BloomAttentionBridge(JointQKVAttentionBridge):
22 """Attention bridge for BLOOM models that handles residual connections and ALiBi.
24 BLOOM attention has a unique forward signature that requires:
25 - residual: The residual connection tensor from before the attention layer
26 - alibi: ALiBi positional encoding bias
27 - attention_mask: Attention mask for padding/causality
29 This bridge ensures these arguments are properly passed through to the original component.
30 """
32 def __init__(
33 self,
34 name: str,
35 config: Any,
36 split_qkv_matrix: Optional[Callable] = None,
37 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
38 qkv_conversion_rule: Optional[BaseTensorConversion] = None,
39 attn_conversion_rule: Optional[BaseTensorConversion] = None,
40 pattern_conversion_rule: Optional[BaseTensorConversion] = None,
41 ):
42 """Initialize the BLOOM attention bridge.
44 Args:
45 name: The name of this component
46 config: Model configuration
47 split_qkv_matrix: Function to split the qkv matrix into q, k, and v
48 submodules: Dictionary of submodules to register
49 qkv_conversion_rule: Optional conversion rule for q, k, v matrices
50 attn_conversion_rule: Optional conversion rule for attention output
51 pattern_conversion_rule: Optional conversion rule for attention patterns
52 """
53 # BLOOM attention doesn't require attention_mask as a constructor arg,
54 # but it DOES require it in forward(), so we don't set requires_attention_mask=True
55 super().__init__(
56 name=name,
57 config=config,
58 split_qkv_matrix=split_qkv_matrix,
59 submodules=submodules,
60 qkv_conversion_rule=qkv_conversion_rule,
61 attn_conversion_rule=attn_conversion_rule,
62 pattern_conversion_rule=pattern_conversion_rule,
63 requires_position_embeddings=False,
64 requires_attention_mask=False,
65 )
67 def forward(self, *args: Any, **kwargs: Any) -> Any:
68 """Forward pass through BLOOM attention with hooks.
70 Uses the parent's hooked Q/K/V split path so that hook_q, hook_k, hook_v,
71 hook_attn_scores, and hook_pattern all fire correctly. ALiBi bias and
72 attention masking are handled in _reconstruct_attention.
74 BLOOM attention requires these arguments:
75 - hidden_states (first positional arg)
76 - residual (second positional arg)
77 - alibi, attention_mask, layer_past, etc. (keyword args)
79 Args:
80 *args: Input arguments (hidden_states, residual)
81 **kwargs: Additional keyword arguments including alibi, attention_mask
83 Returns:
84 Output from BLOOM attention (tuple of hidden_states and optionally attention_weights)
85 """
86 if self.original_component is None: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true
87 raise RuntimeError(
88 f"Original component not set for {self.name}. Call set_original_component() first."
89 )
91 # Extract hidden_states (first positional arg) and residual (second positional arg)
92 if len(args) > 0 and isinstance(args[0], torch.Tensor): 92 ↛ 94line 92 didn't jump to line 94 because the condition on line 92 was always true
93 hidden_states = args[0]
94 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
95 hidden_states = kwargs["hidden_states"]
96 else:
97 raise ValueError("Could not find hidden_states in args or kwargs")
99 residual = args[1] if len(args) > 1 and isinstance(args[1], torch.Tensor) else None
101 # Apply input hook
102 hooked_input = self.hook_in(hidden_states)
104 # Run through split Q/K/V projections (these fire hook_q, hook_k, hook_v),
105 # via the per-head fork when use_split_qkv_input / use_attn_in is set so
106 # those gated hooks fire here as they do on every other joint-QKV bridge.
107 if self._is_split_qkv_fork_active():
108 q_output, k_output, v_output = self._split_forward_qkv(hooked_input)
109 else:
110 q_output = self.q(hooked_input)
111 k_output = self.k(hooked_input)
112 v_output = self.v(hooked_input)
114 # Reconstruct attention with ALiBi (fires hook_attn_scores, hook_pattern)
115 attn_output, attn_weights = self._reconstruct_attention(
116 q_output, k_output, v_output, **kwargs
117 )
119 # BLOOM's original attention applies dropout_add(dense_output, residual, ...)
120 # inside the attention module, not in the block. We must replicate this.
121 if residual is not None: 121 ↛ 131line 121 didn't jump to line 131 because the condition on line 121 was always true
122 assert self.original_component is not None
123 hidden_dropout = getattr(self.original_component, "hidden_dropout", 0.0)
124 if self.training: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true
125 attn_output = torch.nn.functional.dropout(
126 attn_output, p=hidden_dropout, training=True
127 )
128 attn_output = attn_output + residual
130 # Apply output hook
131 output = (attn_output, attn_weights)
132 output = self._process_output(output)
134 return output
136 def _reconstruct_attention(
137 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any
138 ) -> tuple:
139 """Reconstruct attention using BLOOM's ALiBi-based score computation.
141 BLOOM fuses the ALiBi positional bias into scores via baddbmm.
142 """
143 assert self.original_component is not None
144 assert self.config is not None
145 num_heads = self.config.n_heads
147 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(q, k, v, num_heads)
149 # KV cache: extend K/V with cached positions.
150 k, v = self._update_kv_cache(k, v, **kwargs)
152 kv_seq_len = k.shape[-2] # Includes cached positions
153 # Reshape to [batch*heads, seq, head_dim] for baddbmm
154 q_bh = q.reshape(batch_size * num_heads, seq_len, head_dim)
155 k_bh = k.reshape(batch_size * num_heads, kv_seq_len, head_dim)
156 v_bh = v.reshape(batch_size * num_heads, kv_seq_len, head_dim)
158 inv_norm_factor = head_dim ** (-0.5)
160 alibi = kwargs.get("alibi", None)
161 if alibi is not None:
162 # Resize alibi to match kv_seq_len (may differ after cache update).
163 alibi_kv_len = alibi.shape[-1]
164 if alibi_kv_len < kv_seq_len:
165 # ALiBi is slope * position — recompute for the extended length.
166 if alibi.ndim == 3 and alibi.shape[1] == 1: 166 ↛ 176line 166 didn't jump to line 176 because the condition on line 166 was always true
167 slopes = alibi[:, 0, 1:2] # [batch*heads, 1]
168 if slopes.numel() > 0 and slopes.abs().sum() > 0: 168 ↛ 176line 168 didn't jump to line 176 because the condition on line 168 was always true
169 positions = torch.arange(
170 kv_seq_len, device=alibi.device, dtype=alibi.dtype
171 ).unsqueeze(0)
172 alibi = slopes * positions # [batch*heads, kv_seq_len]
173 alibi = alibi.unsqueeze(1) # [batch*heads, 1, kv_seq_len]
174 elif alibi_kv_len > kv_seq_len: 174 ↛ 175line 174 didn't jump to line 175 because the condition on line 174 was never true
175 alibi = alibi[..., :kv_seq_len]
176 attn_scores = alibi.baddbmm(
177 batch1=q_bh,
178 batch2=k_bh.transpose(-1, -2),
179 beta=1.0,
180 alpha=inv_norm_factor,
181 )
182 else:
183 attn_scores = torch.bmm(q_bh, k_bh.transpose(-1, -2)) * inv_norm_factor
185 attn_scores = attn_scores.view(batch_size, num_heads, seq_len, -1)
187 attention_mask = kwargs.get("attention_mask", None)
188 if attention_mask is not None:
189 causal_mask = attention_mask[:, :, :, : attn_scores.shape[-1]]
190 attn_scores = attn_scores + causal_mask
192 attn_scores = self.hook_attn_scores(attn_scores)
194 # Softmax in float32 for numerical stability (matches HF BLOOM)
195 attn_weights = self._softmax_dropout_pattern(
196 attn_scores, target_dtype=q.dtype, upcast_to_fp32=True
197 )
199 # bmm in [batch*heads, seq, seq] format for BLOOM compatibility
200 attn_weights_bh = attn_weights.reshape(batch_size * num_heads, seq_len, -1)
201 attn_output = torch.bmm(attn_weights_bh, v_bh)
203 attn_output = attn_output.view(batch_size, num_heads, seq_len, head_dim)
204 attn_output = self._reshape_attn_output(
205 attn_output, batch_size, seq_len, num_heads, head_dim
206 )
207 if (
208 bool(getattr(self.config, "use_attn_result", False))
209 and hasattr(self, "o")
210 and self.o.original_component is not None
211 ):
212 # Fire hook_z on the flat pre-projection tensor first, so patches at
213 # hook_z reach the per-head computation (same order as the parent).
214 attn_output = self.o.hook_in(attn_output)
215 z_4d = attn_output.view(batch_size, seq_len, num_heads, head_dim)
216 attn_output = self._compute_per_head_result(z_4d, num_heads, head_dim)
217 else:
218 attn_output = self._apply_output_projection(attn_output)
220 return (attn_output, attn_weights)
222 def set_processed_weights(
223 self, weights: Mapping[str, torch.Tensor | None], verbose: bool = False
224 ) -> None:
225 """Set processed weights and recombine Q/K/V back into combined QKV.
227 BloomAttentionBridge's forward() delegates to the original HF attention
228 component which uses the combined query_key_value weight. After weight
229 processing (fold_ln etc.) modifies the split Q/K/V weights, we must
230 recombine them back into the interleaved QKV format so the original
231 component uses the processed weights.
232 """
233 # First, let the parent distribute weights to Q/K/V/O submodules
234 super().set_processed_weights(dict(weights), verbose=verbose) # type: ignore[arg-type]
236 if self.original_component is None: 236 ↛ 237line 236 didn't jump to line 237 because the condition on line 236 was never true
237 return
239 # Get the processed Q/K/V weights from split components
240 assert self.q.original_component is not None
241 assert self.k.original_component is not None
242 assert self.v.original_component is not None
243 q_weight: torch.Tensor = self.q.original_component.weight.data # type: ignore[union-attr, assignment]
244 k_weight: torch.Tensor = self.k.original_component.weight.data # type: ignore[union-attr, assignment]
245 v_weight: torch.Tensor = self.v.original_component.weight.data # type: ignore[union-attr, assignment]
247 assert self.config is not None
248 n_heads: int = self.config.n_heads
249 d_head: int = self.config.d_head
250 d_model = int(q_weight.shape[1])
252 # Reverse the split: recombine into interleaved QKV format
253 # [n_heads*d_head, d_model] -> [d_model, n_heads, d_head]
254 W_Q = q_weight.T.reshape(d_model, n_heads, d_head)
255 W_K = k_weight.T.reshape(d_model, n_heads, d_head)
256 W_V = v_weight.T.reshape(d_model, n_heads, d_head)
258 # Stack into [d_model, n_heads, 3, d_head] (interleaved format)
259 W_combined = torch.stack([W_Q, W_K, W_V], dim=2)
261 # Reshape to [d_model, 3*n_heads*d_head] and transpose to nn.Linear format
262 qkv_weight = W_combined.reshape(d_model, 3 * n_heads * d_head).T
264 # Update the original component's combined QKV weight
265 self.original_component.query_key_value.weight = torch.nn.Parameter( # type: ignore[union-attr]
266 qkv_weight
267 )
269 # Also recombine biases
270 q_bias = self.q.original_component.bias # type: ignore[union-attr]
271 if q_bias is not None: 271 ↛ exitline 271 didn't return from function 'set_processed_weights' because the condition on line 271 was always true
272 assert self.k.original_component is not None
273 assert self.v.original_component is not None
274 k_bias = self.k.original_component.bias.data # type: ignore[union-attr]
275 v_bias = self.v.original_component.bias.data # type: ignore[union-attr]
277 # [n_heads*d_head] -> [n_heads, d_head]
278 b_Q = q_bias.data.reshape(n_heads, d_head) # type: ignore[union-attr, operator]
279 b_K = k_bias.reshape(n_heads, d_head) # type: ignore[operator]
280 b_V = v_bias.reshape(n_heads, d_head) # type: ignore[operator]
282 # Stack into [n_heads, 3, d_head] and flatten
283 qkv_bias = torch.stack([b_Q, b_K, b_V], dim=1).reshape(3 * n_heads * d_head)
284 self.original_component.query_key_value.bias = torch.nn.Parameter( # type: ignore[union-attr]
285 qkv_bias
286 )