Coverage for transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py: 86%
233 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Joint QKV attention bridge component.
3This module contains the bridge component for attention layers that use a fused qkv matrix.
4"""
5import copy
6from typing import Any, Callable, Dict, Optional
8import einops
9import torch
11from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
12 BaseTensorConversion,
13)
14from transformer_lens.model_bridge.generalized_components.attention import (
15 AttentionBridge,
16)
17from transformer_lens.model_bridge.generalized_components.base import (
18 GeneralizedComponent,
19)
20from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
21from transformer_lens.utilities.quantization import require_readable_weight
24class JointQKVAttentionBridge(AttentionBridge):
25 """Joint QKV attention bridge that wraps a joint qkv linear layer.
27 This component wraps attention layers that use a fused qkv matrix such that
28 the individual activations from the separated q, k, and v matrices are hooked and accessible.
29 """
31 supports_attn_result: bool = True
33 # property_aliases inherited from AttentionBridge (W_Q, W_K, W_V, W_O, b_Q, b_K, b_V, b_O)
35 def __init__(
36 self,
37 name: str,
38 config: Any,
39 split_qkv_matrix: Optional[Callable] = None,
40 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
41 qkv_conversion_rule: Optional[BaseTensorConversion] = None,
42 attn_conversion_rule: Optional[BaseTensorConversion] = None,
43 pattern_conversion_rule: Optional[BaseTensorConversion] = None,
44 requires_position_embeddings: bool = False,
45 requires_attention_mask: bool = False,
46 ):
47 """Initialize the Joint QKV attention bridge.
49 Args:
50 name: The name of this component
51 config: Model configuration (required for auto-conversion detection)
52 split_qkv_matrix: Optional function to split the qkv matrix into q, k, and v linear transformations.
53 If None, uses the default implementation that splits a combined c_attn weight/bias.
54 submodules: Dictionary of submodules to register (e.g., q_proj, k_proj, etc.)
55 qkv_conversion_rule: Optional conversion rule for the individual q, k, and v matrices to convert their output shapes to HookedTransformer format. If None, uses default RearrangeTensorConversion
56 attn_conversion_rule: Optional conversion rule. Passed to parent AttentionBridge. If None, AttentionAutoConversion will be used
57 pattern_conversion_rule: Optional conversion rule for attention patterns. If None,
58 uses AttentionPatternConversion to ensure [n_heads, pos, pos] shape
59 requires_position_embeddings: Whether this attention requires position_embeddings as input
60 requires_attention_mask: Whether this attention requires attention_mask as input
61 """
62 super().__init__(
63 name,
64 config,
65 submodules=submodules,
66 conversion_rule=attn_conversion_rule,
67 pattern_conversion_rule=pattern_conversion_rule,
68 requires_position_embeddings=requires_position_embeddings,
69 requires_attention_mask=requires_attention_mask,
70 )
71 self.split_qkv_matrix = (
72 split_qkv_matrix if split_qkv_matrix is not None else self._default_split_qkv_matrix
73 )
74 if qkv_conversion_rule is not None:
75 self.qkv_conversion_rule = qkv_conversion_rule
76 else:
77 self.qkv_conversion_rule = self._create_qkv_conversion_rule()
78 self.q = LinearBridge(name="q")
79 self.k = LinearBridge(name="k")
80 self.v = LinearBridge(name="v")
81 for submodule_name, submodule in (submodules or {}).items():
82 if not hasattr(self, submodule_name): 82 ↛ 81line 82 didn't jump to line 81 because the condition on line 82 was always true
83 setattr(self, submodule_name, submodule)
84 self.submodules["q"] = self.q
85 self.submodules["k"] = self.k
86 self.submodules["v"] = self.v
87 self.q.hook_out.hook_conversion = self.qkv_conversion_rule
88 self.k.hook_out.hook_conversion = self.qkv_conversion_rule
89 self.v.hook_out.hook_conversion = self.qkv_conversion_rule
91 # Register q, k, v LinearBridges in real_components for weight distribution
92 # This allows set_processed_weights to distribute weights to these submodules
93 self.real_components["q"] = ("q", self.q)
94 self.real_components["k"] = ("k", self.k)
95 self.real_components["v"] = ("v", self.v)
96 if hasattr(self, "o"):
97 self.real_components["o"] = ("o", self.o)
99 self._reference_model: Optional[Any] = None
101 # Exclude stale qkv combined weights from state_dict after splitting.
102 self._register_state_dict_hook(JointQKVAttentionBridge._filter_qkv_state_dict)
103 self.register_load_state_dict_pre_hook(
104 JointQKVAttentionBridge._restore_filtered_qkv_state_dict
105 )
107 def __deepcopy__(self, memo):
108 """Share split_qkv_matrix and config across clones instead of copying.
110 split_qkv_matrix may be a bound method of the architecture adapter,
111 which transitively references the full HF model. Without this override,
112 deepcopy duplicates the entire model per block (~1GB x N_layers).
113 """
114 saved_split_fn = self.split_qkv_matrix
115 saved_config = self.config
117 self.split_qkv_matrix = None # type: ignore[assignment]
118 self.config = None
119 try:
120 # Remove override from defining class (not subclass) to avoid recursion.
121 owner = JointQKVAttentionBridge
122 override = owner.__dict__["__deepcopy__"]
123 del owner.__deepcopy__
124 try:
125 clone = copy.deepcopy(self, memo)
126 finally:
127 owner.__deepcopy__ = override # type: ignore[method-assign]
128 finally:
129 self.split_qkv_matrix = saved_split_fn
130 self.config = saved_config
132 clone.split_qkv_matrix = saved_split_fn
133 clone.config = saved_config
134 return clone
136 @staticmethod
137 def _filter_qkv_state_dict(
138 module: torch.nn.Module,
139 state_dict: Dict[str, Any],
140 prefix: str,
141 local_metadata: Dict[str, Any],
142 ) -> None:
143 """State dict hook that removes stale combined QKV entries."""
144 qkv_prefix = prefix + "qkv."
145 keys_to_remove = [k for k in state_dict if k.startswith(qkv_prefix)]
146 for k in keys_to_remove:
147 del state_dict[k]
149 @staticmethod
150 def _restore_filtered_qkv_state_dict(
151 module: torch.nn.Module,
152 state_dict: Dict[str, Any],
153 prefix: str,
154 local_metadata: Dict[str, Any],
155 strict: bool,
156 missing_keys: list[str],
157 unexpected_keys: list[str],
158 error_msgs: list[str],
159 ) -> None:
160 """Insert current combined weights only to satisfy strict key matching.
162 Production checkpoints restore authoritative values through the unfiltered
163 Hugging Face ``_original_component`` path.
164 """
165 del local_metadata, strict, missing_keys, unexpected_keys, error_msgs
166 qkv = module._modules.get("qkv")
167 if qkv is None: 167 ↛ 168line 167 didn't jump to line 168 because the condition on line 167 was never true
168 return
169 for key, value in qkv.state_dict(prefix=f"{prefix}qkv.").items():
170 state_dict.setdefault(key, value)
172 def _create_qkv_conversion_rule(self) -> BaseTensorConversion:
173 """Create the appropriate conversion rule for the individual q, k, and v matrices.
175 Returns:
176 BaseTensorConversion for individual q, k, and v matrices
177 """
178 assert self.config is not None
180 class ConditionalRearrangeConversion(BaseTensorConversion):
181 def __init__(self, n_heads: int):
182 super().__init__()
183 self.n_heads = n_heads
184 self.pattern = (
185 "batch seq (num_attention_heads d_head) -> batch seq num_attention_heads d_head"
186 )
188 def handle_conversion(self, input_value: torch.Tensor, *full_context) -> torch.Tensor:
189 if input_value.ndim == 4: 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true
190 return input_value
191 elif input_value.ndim == 3: 191 ↛ 196line 191 didn't jump to line 196 because the condition on line 191 was always true
192 return einops.rearrange(
193 input_value, self.pattern, num_attention_heads=self.n_heads
194 )
195 else:
196 raise ValueError(
197 f"Expected 3D or 4D tensor, got {input_value.ndim}D with shape {input_value.shape}"
198 )
200 def revert(self, input_value: torch.Tensor, *full_context) -> torch.Tensor:
201 if input_value.ndim == 3: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 return input_value
203 elif input_value.ndim == 4: 203 ↛ 210line 203 didn't jump to line 210 because the condition on line 203 was always true
204 return einops.rearrange(
205 input_value,
206 "batch seq num_attention_heads d_head -> batch seq (num_attention_heads d_head)",
207 num_attention_heads=self.n_heads,
208 )
209 else:
210 raise ValueError(
211 f"Expected 3D or 4D tensor, got {input_value.ndim}D with shape {input_value.shape}"
212 )
214 return ConditionalRearrangeConversion(self.config.n_heads)
216 def _default_split_qkv_matrix(
217 self, original_attention_component: Any
218 ) -> tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]:
219 """Default implementation to split the QKV matrix into separate linear transformations.
221 This uses the 'qkv' submodule defined in component_mapping to find the combined QKV weights.
222 Assumes combined QKV weights in the format [d_model, 3 * d_model] for weights
223 and [3 * n_head * d_head] for bias.
225 Args:
226 original_attention_component: The original attention layer component
227 Returns:
228 Tuple of nn.Linear modules for Q, K, and V transformations
229 """
230 assert self.config is not None
231 assert original_attention_component is not None
233 # Get the combined QKV component using the 'qkv' submodule name
234 if "qkv" not in self.submodules: 234 ↛ 235line 234 didn't jump to line 235 because the condition on line 234 was never true
235 raise ValueError(
236 "No 'qkv' submodule found in JointQKVAttentionBridge. "
237 "Please define a 'qkv' submodule or provide a custom split_qkv_matrix function."
238 )
240 # Get the actual qkv component name from the bridge
241 qkv_bridge = self.submodules["qkv"]
242 qkv_name = qkv_bridge.name
244 # Ensure qkv_name is not None
245 if qkv_name is None: 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true
246 raise ValueError(
247 "qkv bridge name is None. " "Please provide a custom split_qkv_matrix function."
248 )
250 # Navigate to the component using the name
251 if not hasattr(original_attention_component, qkv_name): 251 ↛ 252line 251 didn't jump to line 252 because the condition on line 251 was never true
252 raise ValueError(
253 f"Cannot find '{qkv_name}' in attention component. "
254 f"Available attributes: {dir(original_attention_component)}. "
255 f"Please provide a custom split_qkv_matrix function."
256 )
258 qkv_component = getattr(original_attention_component, qkv_name)
260 # Before the tensor_split/nn.Parameter below: int8 and uint8 would die
261 # there with an opaque "Only Tensors of floating point ... can require
262 # gradients", and float8 would split SILENTLY into scale-less pieces.
263 qkv_weights = require_readable_weight(
264 qkv_component.weight,
265 operation="split a fused QKV projection at boot",
266 owner=qkv_component,
267 )
269 # Original qkv_weights shape: [d_model, 3 * d_model]
270 # Split into three equal parts along dimension 1 to get Q, K, V weights
271 q_weight, k_weight, v_weight = torch.tensor_split(qkv_weights, 3, dim=1)
273 # Handle bias if it exists
274 has_bias = hasattr(qkv_component, "bias") and qkv_component.bias is not None
275 q_bias: torch.Tensor | None
276 k_bias: torch.Tensor | None
277 v_bias: torch.Tensor | None
278 if has_bias: 278 ↛ 287line 278 didn't jump to line 287 because the condition on line 278 was always true
279 qkv_bias = qkv_component.bias
280 assert isinstance(qkv_bias, torch.Tensor)
282 # Original qkv_bias shape: [3 * n_head * d_head]
283 # Reshape to [3, n_head * d_head] to split by Q, K, V
284 qkv_bias = qkv_bias.reshape(3, self.config.n_heads * self.config.d_head)
285 q_bias, k_bias, v_bias = qkv_bias[0, :], qkv_bias[1, :], qkv_bias[2, :]
286 else:
287 q_bias = k_bias = v_bias = None
289 # Create plain nn.Linear modules that output 3D tensors [batch, seq, d_model]
290 q_linear = torch.nn.Linear(q_weight.shape[0], q_weight.shape[1], bias=has_bias)
291 q_linear.weight = torch.nn.Parameter(q_weight.T)
292 if has_bias and q_bias is not None: 292 ↛ 295line 292 didn't jump to line 295 because the condition on line 292 was always true
293 q_linear.bias = torch.nn.Parameter(q_bias)
295 k_linear = torch.nn.Linear(k_weight.shape[0], k_weight.shape[1], bias=has_bias)
296 k_linear.weight = torch.nn.Parameter(k_weight.T)
297 if has_bias and k_bias is not None: 297 ↛ 300line 297 didn't jump to line 300 because the condition on line 297 was always true
298 k_linear.bias = torch.nn.Parameter(k_bias)
300 v_linear = torch.nn.Linear(v_weight.shape[0], v_weight.shape[1], bias=has_bias)
301 v_linear.weight = torch.nn.Parameter(v_weight.T)
302 if has_bias and v_bias is not None: 302 ↛ 305line 302 didn't jump to line 305 because the condition on line 302 was always true
303 v_linear.bias = torch.nn.Parameter(v_bias)
305 return q_linear, k_linear, v_linear
307 def set_original_component(self, original_component: torch.nn.Module) -> None:
308 """Set the original component that this bridge wraps and initialize LinearBridges for q, k, v, and o transformations.
310 Args:
311 original_component: The original attention layer to wrap
312 """
313 super().set_original_component(original_component)
315 # Capture HF-specific attention flags for faithful reconstruction
316 self._reorder_and_upcast_attn = getattr(
317 original_component, "reorder_and_upcast_attn", False
318 )
320 q_transformation, k_transformation, v_transformation = self.split_qkv_matrix(
321 original_component
322 )
323 self.q.set_original_component(q_transformation)
324 self.k.set_original_component(k_transformation)
325 self.v.set_original_component(v_transformation)
326 if hasattr(self, "o") and hasattr(original_component, "c_proj"):
327 self.o.set_original_component(original_component.c_proj)
329 def forward(self, *args: Any, **kwargs: Any) -> Any:
330 """Forward pass through the qkv linear transformation with hooks.
332 Args:
333 *args: Input arguments, where the first argument should be the input tensor
334 **kwargs: Additional keyword arguments
336 Returns:
337 Output tensor after qkv linear transformation
338 """
339 hooked_input = self._apply_attention_input_hook(*args, **kwargs)
340 if self._is_split_qkv_fork_active():
341 q_output, k_output, v_output = self._split_forward_qkv(hooked_input)
342 else:
343 q_output = self.q(hooked_input)
344 k_output = self.k(hooked_input)
345 v_output = self.v(hooked_input)
346 output = self._reconstruct_attention(q_output, k_output, v_output, **kwargs)
347 output = self._process_output(output)
348 return output
350 def _is_split_qkv_fork_active(self) -> bool:
351 cfg = self.config
352 if cfg is None or not getattr(cfg, "n_heads", 0): 352 ↛ 353line 352 didn't jump to line 353 because the condition on line 352 was never true
353 return False
354 return bool(
355 getattr(cfg, "use_split_qkv_input", False) or getattr(cfg, "use_attn_in", False)
356 )
358 def _split_forward_qkv(
359 self, hidden_states: torch.Tensor
360 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
361 """Fork the residual into independent Q/K/V copies, apply per-head projection.
363 After `split_qkv_matrix` runs in `set_original_component`, q/k/v are
364 separate `nn.Linear` modules whose weights partition the output dim by
365 head (output row h*d_head + i ↔ head h, dim i). Plain nn.Linear applied
366 to a 4D [B, S, H, d_model] copy would broadcast the full weight over
367 every head's copy and then we'd keep only the diagonal — n_heads× extra
368 compute. The per-head einsum in `_project_per_head_qkv` slices W per
369 head directly, producing the same 4D [B, S, H, d_head] result that
370 `_reconstruct_attention` expects.
371 """
372 cfg = self.config
373 assert cfg is not None, "config required for split QKV fork"
374 n_heads = int(cfg.n_heads)
375 n_kv_heads = int(getattr(cfg, "n_key_value_heads", None) or n_heads)
376 d_head = int(getattr(cfg, "d_head", 0) or (int(cfg.d_model) // n_heads))
377 use_split = bool(getattr(cfg, "use_split_qkv_input", False))
378 # #1317: fork pre-LN when available so hook patches match legacy.
379 captured = self._captured_pre_ln_residual
380 source = captured if captured is not None else hidden_states
381 if use_split:
382 q_in = self._fork_and_norm_per_head(source, self.hook_q_input, n_heads)
383 k_in = self._fork_and_norm_per_head(source, self.hook_k_input, n_kv_heads)
384 v_in = self._fork_and_norm_per_head(source, self.hook_v_input, n_kv_heads)
385 else:
386 attn_in = self._fork_and_norm_per_head(source, self.hook_attn_in, n_heads)
387 q_in = attn_in
388 if n_kv_heads != n_heads: 388 ↛ 389line 388 didn't jump to line 389 because the condition on line 388 was never true
389 k_in = attn_in[..., :n_kv_heads, :].contiguous()
390 v_in = attn_in[..., :n_kv_heads, :].contiguous()
391 else:
392 k_in = v_in = attn_in
393 q_4d = self._project_per_head_qkv(self.q, q_in, n_heads, d_head)
394 k_4d = self._project_per_head_qkv(self.k, k_in, n_kv_heads, d_head)
395 v_4d = self._project_per_head_qkv(self.v, v_in, n_kv_heads, d_head)
396 return q_4d, k_4d, v_4d
398 def _process_output(self, output: Any) -> Any:
399 """Process the output from _reconstruct_attention.
401 This override skips the duplicate hook_pattern call since
402 _reconstruct_attention already applies both hook_attn_scores
403 and hook_pattern correctly.
405 Args:
406 output: Output tuple from _reconstruct_attention (attn_output, attn_weights)
408 Returns:
409 Processed output with hook_out applied
410 """
411 attn_pattern = None
412 if isinstance(output, tuple) and len(output) >= 2: 412 ↛ 414line 412 didn't jump to line 414 because the condition on line 412 was always true
413 attn_pattern = output[1]
414 if attn_pattern is not None:
415 self._pattern = attn_pattern
416 if isinstance(output, tuple) and len(output) > 0 and isinstance(output[0], torch.Tensor): 416 ↛ 420line 416 didn't jump to line 420 because the condition on line 416 was always true
417 processed_output = list(output)
418 processed_output[0] = self.hook_hidden_states(output[0])
419 output = tuple(processed_output)
420 if isinstance(output, torch.Tensor): 420 ↛ 421line 420 didn't jump to line 421 because the condition on line 420 was never true
421 output = self.hook_out(output)
422 elif isinstance(output, tuple) and len(output) > 0: 422 ↛ 429line 422 didn't jump to line 429 because the condition on line 422 was always true
423 processed_tuple = list(output)
424 if isinstance(output[0], torch.Tensor): 424 ↛ 426line 424 didn't jump to line 426 because the condition on line 424 was always true
425 processed_tuple[0] = self.hook_out(output[0])
426 if len(processed_tuple) == 1: 426 ↛ 427line 426 didn't jump to line 427 because the condition on line 426 was never true
427 return processed_tuple[0]
428 output = tuple(processed_tuple)
429 return output
431 def _apply_attention_input_hook(self, *args: Any, **kwargs: Any) -> torch.Tensor:
432 """Apply attention input hook to the input tensor.
434 This method extracts the input tensor from args/kwargs and applies the attention
435 input hook in the same way as the super class.
437 Args:
438 *args: Input arguments, where the first argument should be the input tensor
439 **kwargs: Additional keyword arguments that might contain input
441 Returns:
442 Input tensor with attention input hook applied
444 Raises:
445 ValueError: If no input tensor is found in args or kwargs
446 """
447 input_tensor = None
448 if "query_input" in kwargs: 448 ↛ 449line 448 didn't jump to line 449 because the condition on line 448 was never true
449 input_tensor = kwargs["query_input"]
450 elif "hidden_states" in kwargs:
451 input_tensor = kwargs["hidden_states"]
452 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 452 ↛ 455line 452 didn't jump to line 455 because the condition on line 452 was always true
453 input_tensor = args[0]
454 else:
455 raise ValueError("No input tensor found in args or kwargs")
456 return self.hook_in(input_tensor)
458 def _reconstruct_attention(
459 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs
460 ) -> tuple:
461 """Manual attention reconstruction used by the bridge after splitting fused QKV projections."""
462 assert self.original_component is not None
463 assert self.config is not None
464 num_heads = self.config.n_heads
465 num_kv_heads = getattr(self.config, "n_key_value_heads", None) or num_heads
467 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(
468 q, k, v, num_heads, num_kv_heads
469 )
471 # KV cache: extend K/V with cached positions.
472 k, v = self._update_kv_cache(k, v, **kwargs)
474 # GQA/MQA: expand K/V heads to match Q heads
475 if num_kv_heads != num_heads:
476 n_rep = num_heads // num_kv_heads
477 k = k.repeat_interleave(n_rep, dim=1)
478 v = v.repeat_interleave(n_rep, dim=1)
480 # Attention scale: 1/sqrt(d_head) with optional inverse-layer scaling
481 scale = head_dim ** (-0.5)
482 if ( 482 ↛ 487line 482 didn't jump to line 487 because the condition on line 482 was never true
483 hasattr(self.config, "scale_attn_by_inverse_layer_idx")
484 and self.config.scale_attn_by_inverse_layer_idx
485 and self._layer_idx is not None
486 ):
487 scale /= float(self._layer_idx + 1)
489 # When reorder_and_upcast_attn is True, HF computes attention in float32.
490 reorder_and_upcast = getattr(self, "_reorder_and_upcast_attn", False)
491 if reorder_and_upcast: 491 ↛ 492line 491 didn't jump to line 492 because the condition on line 491 was never true
492 q_scores = q.to(torch.float32)
493 k_scores = k.to(torch.float32)
494 else:
495 q_scores = q
496 k_scores = k
498 kv_seq_len = k.shape[-2] # Includes cached positions
499 attn_scores = torch.matmul(q_scores, k_scores.transpose(-2, -1)) * scale
500 attention_mask = kwargs.get("attention_mask", None)
501 attn_scores = self._apply_reconstruct_attention_mask(
502 attn_scores=attn_scores,
503 attention_mask=attention_mask,
504 seq_len=kv_seq_len,
505 q_seq_len=seq_len,
506 )
508 attn_scores = self.hook_attn_scores(attn_scores)
510 attn_weights = self._softmax_dropout_pattern(
511 attn_scores,
512 target_dtype=v.dtype if reorder_and_upcast else None,
513 )
514 attn_output = torch.matmul(attn_weights, v)
515 attn_output = self._reshape_attn_output(
516 attn_output, batch_size, seq_len, num_heads, head_dim
517 )
518 if (
519 bool(getattr(self.config, "use_attn_result", False))
520 and hasattr(self, "o")
521 and self.o.original_component is not None
522 ):
523 # Per-head output pre-sum. Fire hook_z on the pre-projection flat
524 # tensor first so patches at hook_z propagate into the per-head
525 # computation, matching how the default path's `self.o(...)` call
526 # fires o.hook_in before the linear.
527 attn_output = self.o.hook_in(attn_output)
528 z_4d = attn_output.view(batch_size, seq_len, num_heads, head_dim)
529 attn_output = self._compute_per_head_result(z_4d, num_heads, head_dim)
530 else:
531 attn_output = self._apply_output_projection(attn_output)
532 return (attn_output, attn_weights)