Coverage for transformer_lens/model_bridge/generalized_components/attention.py: 88%
494 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Attention bridge component.
3This module contains the bridge component for attention layers.
4"""
5import logging
6import warnings
7from typing import Any, Dict, Optional, Tuple
9import einops
10import torch
11from transformers.pytorch_utils import Conv1D
13logger = logging.getLogger(__name__)
15from transformer_lens.conversion_utils.conversion_steps.attention_auto_conversion import (
16 AttentionAutoConversion,
17)
18from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
19 BaseTensorConversion,
20)
21from transformer_lens.hook_points import HookPoint
22from transformer_lens.model_bridge.generalized_components.base import (
23 GeneralizedComponent,
24)
25from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config
26from transformer_lens.utilities.quantization import require_readable_weight
29class PerLayerGeometryError(ValueError):
30 """A weight accessor refused a head split on per-layer attention geometry.
32 Dedicated type so callers with a raw-weight fallback (the centering
33 benchmark) catch exactly this and nothing else — a generic ValueError
34 catch would silently mask unrelated accessor regressions.
35 """
38class AttentionBridge(GeneralizedComponent):
39 """Bridge component for attention layers.
41 This component handles the conversion between Hugging Face attention layers
42 and TransformerLens attention components.
43 """
45 hook_aliases = {
46 "hook_q": "q.hook_out",
47 "hook_k": "k.hook_out",
48 "hook_v": "v.hook_out",
49 "hook_z": "o.hook_in",
50 }
52 # Override to False on variants without a pre-LN fork (e.g. MLA); skips
53 # the split-qkv HookPoints and the BlockBridge pre-ln1 capture.
54 supports_split_qkv_fork: bool = True
55 # Reconstructed variants can opt in independently of Q/K/V input forks.
56 supports_attn_result: bool = False
57 property_aliases = {
58 "W_Q": "q.weight",
59 "W_K": "k.weight",
60 "W_V": "v.weight",
61 "W_O": "o.weight",
62 "b_Q": "q.bias",
63 "b_K": "k.bias",
64 "b_V": "v.bias",
65 "b_O": "o.bias",
66 }
68 def __init__(
69 self,
70 name: Optional[str],
71 config: Any,
72 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
73 conversion_rule: Optional[BaseTensorConversion] = None,
74 pattern_conversion_rule: Optional[BaseTensorConversion] = None,
75 maintain_native_attention: bool = False,
76 requires_position_embeddings: bool = False,
77 requires_attention_mask: bool = False,
78 attention_mask_4d: bool = False,
79 requires_relative_position_bias: bool = False,
80 is_cross_attention: bool = False,
81 is_causal: bool = True,
82 optional: bool = False,
83 fused_qkv: bool = False,
84 ):
85 """Initialize the attention bridge.
87 Args:
88 name: The name of this component, or None when projections live on the parent
89 config: Model configuration (required for auto-conversion detection)
90 submodules: Dictionary of submodules to register (e.g., q_proj, k_proj, etc.)
91 conversion_rule: Optional conversion rule. If None, AttentionAutoConversion will be used
92 pattern_conversion_rule: Optional conversion rule for attention patterns. If None,
93 uses AttentionPatternConversion to ensure [n_heads, pos, pos] shape
94 maintain_native_attention: If True, preserve the original HF attention implementation
95 without wrapping. Use for models with custom attention
96 (e.g., attention sinks, specialized RoPE). Defaults to False.
97 requires_position_embeddings: If True, this attention requires position_embeddings argument
98 (e.g., Gemma-3 with dual RoPE). Defaults to False.
99 requires_attention_mask: If True, this attention requires attention_mask argument
100 (e.g., GPTNeoX/Pythia). Defaults to False.
101 attention_mask_4d: If True, generate 4D attention_mask [batch, 1, tgt_len, src_len]
102 instead of 2D [batch, seq_len]. Required for OPT. Defaults to False.
103 requires_relative_position_bias: T5/mT5-style relative attention; supplies a
104 zero ``position_bias`` so HF's forward skips its ``cache_position[-1]`` fallback.
105 is_cross_attention: Encoder-decoder cross-attention; supplies ``key_value_states``.
106 is_causal: If True, apply a causal (lower-triangular) mask when reconstructing
107 attention. Set False for bidirectional encoders (e.g. T5Gemma's encoder).
108 """
109 if conversion_rule is None: 109 ↛ 111line 109 didn't jump to line 111 because the condition on line 109 was always true
110 conversion_rule = AttentionAutoConversion(config)
111 super().__init__(
112 name,
113 config=config,
114 submodules=submodules or {},
115 conversion_rule=conversion_rule,
116 optional=optional,
117 )
118 if fused_qkv:
119 # A combined QKV projection has no q/k/v submodules for the class
120 # aliases to resolve against; leaving them in place produces dead
121 # aliases and resolution-audit noise (raven Wqkv, OpenELM qkv_proj).
122 self.hook_aliases = {
123 key: value
124 for key, value in type(self).hook_aliases.items()
125 if key not in {"hook_q", "hook_k", "hook_v"}
126 }
127 self.hook_attn_scores = HookPoint()
128 self.hook_pattern = HookPoint()
129 self.hook_hidden_states = HookPoint()
130 # Per-head attention output, pre-sum across heads.
131 # Shape [batch, pos, n_heads, d_model] when fired. Gated at fire time
132 # by cfg.use_attn_result; the HookPoint exists unconditionally so
133 # run_with_cache key lookups never miss.
134 self.hook_result = HookPoint()
135 # Pre-ln1 fork hooks ([B, S, H, D]) gated by use_split_qkv_input /
136 # use_attn_in; fall back to post-ln1 if BlockBridge can't wire ln1. See #1317.
137 if self.supports_split_qkv_fork:
138 self.hook_attn_in = HookPoint()
139 self.hook_q_input = HookPoint()
140 self.hook_k_input = HookPoint()
141 self.hook_v_input = HookPoint()
142 self._captured_pre_ln_residual: Optional[torch.Tensor] = None
143 self._ln1_module: Optional[torch.nn.Module] = None
144 if (
145 hasattr(config, "positional_embedding_type")
146 and config.positional_embedding_type == "rotary"
147 ):
148 self.hook_rot_k = HookPoint()
149 self.hook_rot_q = HookPoint()
150 self.hook_hidden_states.hook_conversion = conversion_rule
151 if pattern_conversion_rule is not None: 151 ↛ 152line 151 didn't jump to line 152 because the condition on line 151 was never true
152 self.hook_pattern.hook_conversion = pattern_conversion_rule
153 self._attn_scores = None
154 self._pattern = None
155 self._hf_forward_wrapped = False
156 self.maintain_native_attention = maintain_native_attention
157 self.requires_position_embeddings = requires_position_embeddings
158 self.requires_attention_mask = requires_attention_mask
159 self.attention_mask_4d = attention_mask_4d
160 self.requires_relative_position_bias = requires_relative_position_bias
161 self.is_cross_attention = is_cross_attention
162 self.is_causal = is_causal
163 self._layer_idx: Optional[int] = None
165 def set_original_component(self, original_component: torch.nn.Module) -> None:
166 """Set original component and capture layer index for KV caching."""
167 super().set_original_component(original_component)
168 layer_idx_raw = getattr(original_component, "layer_idx", None)
169 if layer_idx_raw is not None:
170 self._layer_idx = int(layer_idx_raw)
172 def _apply_ln1_per_head(self, x: torch.Tensor) -> torch.Tensor:
173 """Apply ln1 to [B, S, H, D] with H folded into the batch. Identity if ln1 unwired.
175 Routes through the raw HF norm to avoid refiring ln1's internal hooks
176 per-head — deliberate divergence from legacy's *Pre sub-hook firing.
177 """
178 if self._ln1_module is None: 178 ↛ 179line 178 didn't jump to line 179 because the condition on line 178 was never true
179 return x
180 b, s, h, d = x.shape
181 return self._ln1_module(x.reshape(b * s * h, d)).reshape(b, s, h, d)
183 def _fork_and_norm_per_head(
184 self, source: torch.Tensor, hook: HookPoint, n_heads: int
185 ) -> torch.Tensor:
186 """Repeat residual to [B, S, H, D], fire ``hook``, re-LN iff source is pre-LN."""
187 forked = einops.repeat(source, "b s d -> b s h d", h=n_heads).contiguous()
188 forked = hook(forked)
189 if self._captured_pre_ln_residual is not None:
190 forked = self._apply_ln1_per_head(forked)
191 return forked
193 def setup_hook_compatibility(self) -> None:
194 """Setup hook compatibility transformations to match HookedTransformer behavior.
196 This sets up hook conversions that ensure Bridge hooks have the same shapes
197 as HookedTransformer hooks. This includes reshaping Q/K/V/Z hooks from
198 [batch, seq, d_model] to [batch, seq, n_heads, d_head] format.
200 This is called during Bridge.__init__ and should always be run.
201 Note: This method is idempotent - can be called multiple times safely.
202 """
203 if self._hf_forward_wrapped:
204 return
205 if hasattr(self.config, "n_heads"): 205 ↛ 207line 205 didn't jump to line 207 because the condition on line 205 was always true
206 self._setup_qkv_hook_reshaping()
207 self._hf_forward_wrapped = True
209 def get_random_inputs(
210 self,
211 batch_size: int = 2,
212 seq_len: int = 8,
213 device: Optional[torch.device] = None,
214 dtype: Optional[torch.dtype] = None,
215 ) -> Dict[str, Any]:
216 """Get random inputs for testing this attention component.
218 Generates appropriate inputs based on the attention's requirements
219 (position_embeddings, attention_mask, etc.).
221 Args:
222 batch_size: Batch size for the test inputs
223 seq_len: Sequence length for the test inputs
224 device: Device to create tensors on (defaults to CPU)
225 dtype: Dtype for generated tensors (defaults to float32)
227 Returns:
228 Dictionary of keyword arguments to pass to forward()
229 """
230 if device is None:
231 device = torch.device("cpu")
232 if dtype is None:
233 dtype = torch.float32
234 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768
235 inputs: Dict[str, Any] = {
236 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype)
237 }
238 if self.requires_position_embeddings:
239 if self.config: 239 ↛ 247line 239 didn't jump to line 247 because the condition on line 239 was always true
240 if hasattr(self.config, "d_head"): 240 ↛ 242line 240 didn't jump to line 242 because the condition on line 240 was always true
241 d_head = self.config.d_head
242 elif hasattr(self.config, "head_dim"):
243 d_head = self.config.head_dim
244 else:
245 d_head = 64
246 else:
247 d_head = 64
248 rotary_pct = get_rotary_pct_from_config(self.config)
249 rotary_ndims = int(rotary_pct * d_head)
250 cos = torch.ones(1, seq_len, rotary_ndims, device=device, dtype=dtype)
251 sin = torch.zeros(1, seq_len, rotary_ndims, device=device, dtype=dtype)
252 inputs["position_embeddings"] = (cos, sin)
253 # For models with internal rotary embeddings (e.g., GPT-J), the HF attention
254 # forward expects position_ids to index into embed_positions. Models using
255 # requires_position_embeddings get (cos, sin) tuples instead.
256 if ( 256 ↛ 262line 256 didn't jump to line 262 because the condition on line 256 was never true
257 self.config
258 and hasattr(self.config, "positional_embedding_type")
259 and self.config.positional_embedding_type == "rotary"
260 and not self.requires_position_embeddings
261 ):
262 inputs["position_ids"] = (
263 torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
264 )
265 if self.requires_attention_mask: 265 ↛ 266line 265 didn't jump to line 266 because the condition on line 265 was never true
266 if self.attention_mask_4d:
267 # Generate 4D attention mask [batch, 1, tgt_len, src_len] for models like OPT
268 inputs["attention_mask"] = torch.ones(
269 batch_size, 1, seq_len, seq_len, device=device
270 )
271 else:
272 # Generate 2D attention mask [batch, seq_len] for most models
273 inputs["attention_mask"] = torch.ones(batch_size, seq_len, device=device)
274 if self.requires_relative_position_bias: 274 ↛ 276line 274 didn't jump to line 276 because the condition on line 274 was never true
275 # Zero bias short-circuits HF's None-cache_position fallback in T5Attention.
276 n_heads = self.config.n_heads if self.config and hasattr(self.config, "n_heads") else 1
277 inputs["position_bias"] = torch.zeros(
278 1, n_heads, seq_len, seq_len, device=device, dtype=dtype
279 )
280 if self.is_cross_attention:
281 inputs["key_value_states"] = torch.randn(
282 batch_size, seq_len, d_model, device=device, dtype=dtype
283 )
284 return inputs
286 def _setup_qkv_hook_reshaping(self) -> None:
287 """Setup hook reshaping for Q/K/V/Z to match HookedTransformer shapes.
289 Reshapes hooks from [batch, seq, d_model] to [batch, seq, n_heads, d_head] format.
290 For models with Grouped Query Attention (GQA), K and V use n_kv_heads instead of n_heads.
292 Sets up conversions for:
293 - q.hook_out (aliased as hook_q)
294 - k.hook_out (aliased as hook_k) - uses n_kv_heads if GQA
295 - v.hook_out (aliased as hook_v) - uses n_kv_heads if GQA
296 - o.hook_in (aliased as hook_z)
297 """
299 class ReshapeForAttentionHeads(BaseTensorConversion):
300 """Reshape tensors to split attention heads for Q/K/V/Z compatibility."""
302 # Marks the [batch, pos, head, d_head] layout so consumers (e.g. run_with_cache's
303 # pos_slice) know the position axis is two from the end, not one.
304 splits_attention_heads = True
306 def __init__(self, n_heads: int, d_head: int):
307 super().__init__()
308 self.n_heads = n_heads
309 self.d_head = d_head
311 def handle_conversion(self, input_value, *full_context):
312 """Convert from [batch, seq, d_model] to [batch, seq, n_heads, d_head]."""
313 if len(input_value.shape) == 3: 313 ↛ 317line 313 didn't jump to line 317 because the condition on line 313 was always true
314 b, s, d = input_value.shape
315 if d == self.n_heads * self.d_head:
316 return input_value.view(b, s, self.n_heads, self.d_head)
317 return input_value
319 def revert(self, input_value, *full_context):
320 """Revert from [batch, seq, n_heads, d_head] to [batch, seq, d_model]."""
321 if len(input_value.shape) == 4: 321 ↛ 326line 321 didn't jump to line 326 because the condition on line 321 was always true
322 b, s, n_h, d_h = input_value.shape
323 if n_h == self.n_heads and d_h == self.d_head: 323 ↛ 326line 323 didn't jump to line 326 because the condition on line 323 was always true
324 # reshape (not view) — callers may pass non-contiguous tensors
325 return input_value.reshape(b, s, n_h * d_h)
326 return input_value
328 if self.config is None: 328 ↛ 329line 328 didn't jump to line 329 because the condition on line 328 was never true
329 raise RuntimeError(f"Config not set for {self.name}")
331 # Get n_heads (try n_heads first, then n_head)
332 if hasattr(self.config, "n_heads"): 332 ↛ 334line 332 didn't jump to line 334 because the condition on line 332 was always true
333 n_heads = self.config.n_heads
334 elif hasattr(self.config, "n_head"):
335 n_heads = self.config.n_head
336 else:
337 # Can't setup reshaping without knowing number of heads
338 return
340 # Get d_head (try d_head first, then compute from d_model or n_embd)
341 if hasattr(self.config, "d_head"):
342 d_head = self.config.d_head
343 elif hasattr(self.config, "d_model"): 343 ↛ 344line 343 didn't jump to line 344 because the condition on line 343 was never true
344 d_head = self.config.d_model // n_heads
345 elif hasattr(self.config, "n_embd"): 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 d_head = self.config.n_embd // n_heads
347 else:
348 # Can't setup reshaping without knowing head dimension
349 return
350 n_kv_heads = n_heads
351 if hasattr(self.config, "n_key_value_heads") and self.config.n_key_value_heads is not None:
352 n_kv_heads = self.config.n_key_value_heads
354 # Per-layer d_head makes width division ambiguous (gemma4: 2048 = 8x256
355 # = 4x512), so dividing by the majority scalar mis-factorizes silently.
356 layer_d_head: Optional[int] = d_head
357 per_layer_hd = getattr(self.config, "per_layer_head_dim", None)
358 if per_layer_hd:
359 entry: Any = None
360 if self._layer_idx is not None and self._layer_idx < len(per_layer_hd):
361 entry = per_layer_hd[self._layer_idx]
362 if not (isinstance(entry, int) and entry > 0):
363 entry = getattr(self.original_component, "head_dim", None)
364 layer_d_head = entry if isinstance(entry, int) and entry > 0 else None
366 def conversion_dims(
367 proj: Any, cfg_heads: int, width_attr: str = "out_features"
368 ) -> Tuple[int, int]:
369 """(n_heads, d_head) from the BOUND module width, not the cfg scalars.
371 OpenELM/laguna vary head counts per layer; cfg-only dims no-op there,
372 flipping hook shapes mid-model. MLA excluded — its o-width is v_head_dim.
373 """
374 if getattr(self, "_v_head_dim", None): 374 ↛ 375line 374 didn't jump to line 375 because the condition on line 374 was never true
375 return cfg_heads, d_head
376 component = getattr(proj, "original_component", None)
377 width = getattr(component, width_attr, None)
378 if width and layer_d_head and width % layer_d_head == 0:
379 return width // layer_d_head, layer_d_head
380 return cfg_heads, layer_d_head or d_head
382 if hasattr(self, "q") and self.q is not None and hasattr(self.q, "hook_out"):
383 q_reshape = ReshapeForAttentionHeads(*conversion_dims(self.q, n_heads))
384 self.q.hook_out.hook_conversion = q_reshape
385 if hasattr(self, "k") and self.k is not None and hasattr(self.k, "hook_out"):
386 k_reshape = ReshapeForAttentionHeads(*conversion_dims(self.k, n_kv_heads))
387 self.k.hook_out.hook_conversion = k_reshape
388 # Subclasses that de-alias hook_k onto their own HookPoint (K-scaling
389 # architectures) need the same conversion, whichever order runs first.
390 if "hook_k" not in self.hook_aliases and isinstance(
391 getattr(self, "hook_k", None), HookPoint
392 ):
393 self.hook_k.hook_conversion = k_reshape
394 if hasattr(self, "v") and self.v is not None and hasattr(self.v, "hook_out"):
395 v_reshape = ReshapeForAttentionHeads(*conversion_dims(self.v, n_kv_heads))
396 self.v.hook_out.hook_conversion = v_reshape
397 if hasattr(self, "o") and self.o is not None and hasattr(self.o, "hook_in"): 397 ↛ 403line 397 didn't jump to line 403 because the condition on line 397 was always true
398 z_reshape = ReshapeForAttentionHeads(
399 *conversion_dims(self.o, n_heads, width_attr="in_features")
400 )
401 self.o.hook_in.hook_conversion = z_reshape
403 class TransposeRotaryHeads(BaseTensorConversion):
404 """Transpose rotary hook tensors from HF format to HookedTransformer format."""
406 def handle_conversion(self, input_value, *full_context):
407 """Convert from [batch, n_heads, seq, d_head] to [batch, seq, n_heads, d_head]."""
408 if len(input_value.shape) == 4: 408 ↛ 410line 408 didn't jump to line 410 because the condition on line 408 was always true
409 return input_value.transpose(1, 2)
410 return input_value
412 def revert(self, input_value, *full_context):
413 """Revert from [batch, seq, n_heads, d_head] to [batch, n_heads, seq, d_head]."""
414 if len(input_value.shape) == 4: 414 ↛ 416line 414 didn't jump to line 416 because the condition on line 414 was always true
415 return input_value.transpose(1, 2)
416 return input_value
418 if hasattr(self, "hook_rot_q"):
419 self.hook_rot_q.hook_conversion = TransposeRotaryHeads()
420 if hasattr(self, "hook_rot_k"):
421 self.hook_rot_k.hook_conversion = TransposeRotaryHeads()
423 def _update_kv_cache(
424 self, k: torch.Tensor, v: torch.Tensor, **kwargs: Any
425 ) -> tuple[torch.Tensor, torch.Tensor]:
426 """Update KV cache if provided, returning the (possibly extended) K and V.
428 Call this after K/V projections and any positional embeddings (e.g. RoPE)
429 have been applied, but before computing attention scores. If no cache is
430 present in kwargs, K and V are returned unchanged.
431 """
432 past_key_values = kwargs.get("past_key_values", None)
433 if past_key_values is None:
434 # GPT-NeoX/GPT-J/Bloom/Falcon/MPT/CodeGen/GPTBigCode still name the
435 # cache `layer_past`; missing it leaves the cache empty, so every
436 # decode step attends to itself alone and generation ignores the prompt.
437 past_key_values = kwargs.get("layer_past", None)
438 if past_key_values is None:
439 return k, v
440 layer_idx = getattr(self, "_layer_idx", None)
441 if layer_idx is None:
442 logger.warning(
443 "%s: past_key_values provided but _layer_idx is None "
444 "(HF component missing layer_idx attribute). KV cache update "
445 "skipped — cached generation will ignore earlier tokens.",
446 self.name,
447 )
448 return k, v
449 # Some architectures (e.g. HRM-Text's recurrent stacks) pass a cycle_offset
450 # so each stack invocation writes to a unique cache slot. Non-participating
451 # models simply leave it unset.
452 cycle_offset = kwargs.get("cycle_offset", 0)
453 k, v = past_key_values.update(k, v, layer_idx + cycle_offset)
454 return k, v
456 def _reshape_qkv_to_heads(
457 self,
458 q: torch.Tensor,
459 k: torch.Tensor,
460 v: torch.Tensor,
461 num_heads: int,
462 num_kv_heads: int | None = None,
463 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]:
464 """Reshape Q/K/V from [batch, seq, hidden] or [batch, seq, heads, head_dim]
465 to [batch, heads, seq, head_dim]. Returns (q, k, v, batch_size, seq_len, head_dim).
467 Args:
468 num_kv_heads: If provided and differs from num_heads (GQA), K/V use
469 this head count for the 3D reshape path.
470 """
471 if num_kv_heads is None:
472 num_kv_heads = num_heads
473 if q.ndim == 3:
474 batch_size, seq_len, q_hidden = q.shape
475 head_dim: int = q_hidden // num_heads
476 q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
477 k = k.view(batch_size, seq_len, num_kv_heads, head_dim).transpose(1, 2)
478 v = v.view(batch_size, seq_len, num_kv_heads, head_dim).transpose(1, 2)
479 elif q.ndim == 4: 479 ↛ 486line 479 didn't jump to line 486 because the condition on line 479 was always true
480 batch_size, seq_len = q.shape[0], q.shape[1]
481 head_dim = q.shape[-1]
482 q = q.transpose(1, 2)
483 k = k.transpose(1, 2)
484 v = v.transpose(1, 2)
485 else:
486 raise ValueError(f"Unexpected Q tensor shape: {q.shape}. Expected 3D or 4D.")
487 return q, k, v, batch_size, seq_len, head_dim
489 def _apply_attn_dropout(self, attn_weights: torch.Tensor) -> torch.Tensor:
490 """Apply attention dropout from the original HF component if present."""
491 if self.original_component is not None: 491 ↛ 497line 491 didn't jump to line 497 because the condition on line 491 was always true
492 dropout_fn = getattr(self.original_component, "attn_dropout", None)
493 if dropout_fn is None:
494 dropout_fn = getattr(self.original_component, "attention_dropout", None)
495 if dropout_fn is not None and callable(dropout_fn):
496 attn_weights = dropout_fn(attn_weights)
497 return attn_weights
499 def _apply_output_projection(self, attn_output: torch.Tensor) -> torch.Tensor:
500 """Apply the output projection (self.o) if present."""
501 if hasattr(self, "o") and self.o is not None:
502 attn_output = self.o(attn_output)
503 return attn_output
505 def _softmax_dropout_pattern(
506 self,
507 attn_scores: torch.Tensor,
508 target_dtype: torch.dtype | None = None,
509 upcast_to_fp32: bool = False,
510 ) -> torch.Tensor:
511 """Apply softmax, dropout, and hook_pattern to attention scores.
513 Args:
514 attn_scores: Raw attention scores [batch, heads, q_seq, kv_seq].
515 target_dtype: If set, cast weights to this dtype after softmax.
516 upcast_to_fp32: If True, compute softmax in float32 for numerical
517 stability, then cast to target_dtype.
518 """
519 if upcast_to_fp32:
520 attn_weights = self._masked_softmax(attn_scores, dtype=torch.float32)
521 if target_dtype is not None: 521 ↛ 527line 521 didn't jump to line 527 because the condition on line 521 was always true
522 attn_weights = attn_weights.to(target_dtype)
523 else:
524 attn_weights = self._masked_softmax(attn_scores)
525 if target_dtype is not None:
526 attn_weights = attn_weights.to(target_dtype)
527 attn_weights = self._scrub_compatibility_pattern_nans(attn_weights)
528 attn_weights = self._apply_attn_dropout(attn_weights)
529 attn_weights = self.hook_pattern(attn_weights)
530 return attn_weights
532 def _masked_softmax(
533 self, attn_scores: torch.Tensor, dtype: torch.dtype | None = None
534 ) -> torch.Tensor:
535 """Softmax over keys that zeroes fully masked rows in compatibility mode.
537 Compatibility mode masks with -inf, so a query with no visible key (e.g.
538 the leading pads of a left-padded row) softmaxes to NaN. Zeroing that NaN
539 afterwards only fixes the forward: softmax's backward still multiplies by
540 the NaN output and leaks it into Q/K and the residual stream.
541 """
542 if not self.compatibility_mode:
543 return torch.nn.functional.softmax(attn_scores, dim=-1, dtype=dtype)
544 fully_masked = torch.isneginf(attn_scores).all(dim=-1, keepdim=True)
545 pattern = torch.nn.functional.softmax(
546 attn_scores.masked_fill(fully_masked, 0.0), dim=-1, dtype=dtype
547 )
548 return pattern.masked_fill(fully_masked, 0.0)
550 def _scrub_compatibility_pattern_nans(self, pattern: torch.Tensor) -> torch.Tensor:
551 """Match HookedTransformer's NaN scrub of the attention pattern."""
552 if self.compatibility_mode:
553 pattern = torch.where(torch.isnan(pattern), torch.zeros_like(pattern), pattern)
554 return pattern
556 def _normalize_compatibility_mask_sentinel(self, attention_mask: torch.Tensor) -> torch.Tensor:
557 """Normalize additive mask sentinels before dtype conversion or addition."""
558 if self.compatibility_mode and attention_mask.is_floating_point():
559 attention_mask = attention_mask.masked_fill(
560 attention_mask <= torch.finfo(attention_mask.dtype).min, -torch.inf
561 )
562 return attention_mask
564 def _reshape_attn_output(
565 self,
566 attn_output: torch.Tensor,
567 batch_size: int,
568 seq_len: int,
569 num_heads: int,
570 head_dim: int,
571 ) -> torch.Tensor:
572 """Reshape attention output from [batch, heads, seq, dim] to [batch, seq, hidden]."""
573 attn_output = attn_output.transpose(1, 2).contiguous()
574 attn_output = attn_output.view(batch_size, seq_len, num_heads * head_dim)
575 return attn_output
577 def _apply_reconstruct_attention_mask(
578 self,
579 attn_scores: torch.Tensor,
580 attention_mask: torch.Tensor | None,
581 seq_len: int,
582 q_seq_len: int | None = None,
583 ) -> torch.Tensor:
584 """Apply causal and optional attention masking to reconstructed scores.
586 HuggingFace-style 4D masks already encode causal semantics, so they are
587 treated as authoritative. Lower-rank masks do not, so the local causal
588 mask is still applied before adding the caller-provided padding mask.
590 Args:
591 attn_scores: Attention scores [batch, heads, q_seq_len, kv_seq_len].
592 attention_mask: Optional mask from the caller.
593 seq_len: The KV sequence length (total positions including cache).
594 q_seq_len: The query sequence length. When using KV cache this is
595 shorter than seq_len. Defaults to seq_len when not provided.
596 """
597 if q_seq_len is None:
598 q_seq_len = seq_len
599 min_dtype = torch.finfo(attn_scores.dtype).min
600 mask_value = -torch.inf if self.compatibility_mode else min_dtype
601 use_direct_hf_mask = attention_mask is not None and attention_mask.ndim >= 4
602 # Bidirectional attention (encoders) and cross-attention have no causal
603 # structure, so only synthesize the triangular mask for causal self-attention.
604 apply_causal = self.is_causal and not self.is_cross_attention
605 if not use_direct_hf_mask and apply_causal:
606 # Rectangular causal mask: query i attends to KV 0..(offset+i)
607 # where offset = kv_seq_len - q_seq_len (cached positions).
608 causal_mask = torch.ones(
609 q_seq_len, seq_len, device=attn_scores.device, dtype=torch.bool
610 )
611 causal_mask = torch.tril(causal_mask, diagonal=seq_len - q_seq_len)
612 attn_scores = attn_scores.masked_fill(~causal_mask, mask_value)
614 if attention_mask is not None:
615 if attention_mask.shape[-1] != seq_len: 615 ↛ 616line 615 didn't jump to line 616 because the condition on line 615 was never true
616 attention_mask = attention_mask[..., :seq_len]
617 if attention_mask.ndim >= 3 and attention_mask.shape[-2] != q_seq_len:
618 # Extra query rows mean a full-sequence mask on a cached decode step,
619 # where the live queries are the LAST rows; taking the first hands
620 # every step position 0 (Baichuan-13B fuses ALiBi slopes in here).
621 attention_mask = attention_mask[..., -q_seq_len:, :]
623 if attention_mask.dtype == torch.bool:
624 attention_mask = torch.where(
625 attention_mask,
626 torch.zeros((), dtype=attn_scores.dtype, device=attn_scores.device),
627 torch.full((), mask_value, dtype=attn_scores.dtype, device=attn_scores.device),
628 )
629 else:
630 attention_mask = self._normalize_compatibility_mask_sentinel(attention_mask)
631 attention_mask = attention_mask.to(dtype=attn_scores.dtype)
632 attn_scores = attn_scores + attention_mask
634 return attn_scores
636 def _get_n_heads(self, use_kv: bool = False) -> int:
637 """Resolve the number of attention heads from config.
639 Args:
640 use_kv: If True, return n_key_value_heads (for GQA) when available.
641 """
642 assert self.config is not None, "config required to resolve n_heads"
643 if use_kv:
644 if hasattr(self.config, "n_key_value_heads") and self.config.n_key_value_heads:
645 return self.config.n_key_value_heads
646 if hasattr(self.config, "n_heads"):
647 return self.config.n_heads
648 return self.config.n_head
650 def _weight_layout_in_out(self, proj: Any) -> Optional[bool]:
651 """Whether ``proj``'s wrapped module stores its weight as [in, out].
653 Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores
654 [out_features, in_features]. Returns None when the wrapped module is
655 neither, so callers can fall back to a shape heuristic.
656 """
657 component = getattr(proj, "original_component", None)
658 if isinstance(component, Conv1D):
659 return True
660 if isinstance(component, torch.nn.Linear):
661 return False
662 return None
664 def _reshape_weight_to_3d(
665 self,
666 weight: torch.Tensor,
667 n_heads: int,
668 pattern: str = "qkv",
669 in_out_layout: Optional[bool] = None,
670 ) -> torch.Tensor:
671 """Reshape a 2D weight to 3D by splitting heads.
673 Args:
674 weight: 2D weight tensor
675 n_heads: Number of heads to split into
676 pattern: "qkv" for [n_heads, d_model, d_head], "o" for [n_heads, d_head, d_model]
677 in_out_layout: True if the weight is stored [in, out] (Conv1D), False
678 for [out, in] (nn.Linear). None falls back to a shape heuristic,
679 which cannot distinguish the two when the weight is square.
680 """
681 if pattern == "o":
682 if in_out_layout is None:
683 # Heuristic: assumes [in, out] whenever the head split fits dim 0.
684 in_out_layout = weight.shape[0] == n_heads * (
685 weight.shape[1] // n_heads
686 if weight.shape[1] % n_heads == 0
687 else weight.shape[0] // n_heads
688 )
689 mat = weight if in_out_layout else weight.T
690 self._check_head_split_width(mat, n_heads)
691 return einops.rearrange(
692 mat, "(n_heads d_head) d_model -> n_heads d_head d_model", n_heads=n_heads
693 )
694 # QKV pattern
695 if in_out_layout is None:
696 in_out_layout = weight.shape[0] % n_heads != 0
697 mat = weight.T if in_out_layout else weight
698 self._check_head_split_width(mat, n_heads)
699 return einops.rearrange(
700 mat, "(n_heads d_head) d_model -> n_heads d_model d_head", n_heads=n_heads
701 )
703 def _check_head_split_width(self, mat: torch.Tensor, n_heads: int) -> None:
704 """Refuse a head split whose width disagrees with the model's geometry.
706 einops only needs divisibility, so per-layer geometry (gemma4 KV-shared
707 layers) would factorize into wrong-shaped heads silently. MLA's
708 bind-time ``_v_head_dim`` is a legitimate second width (o_proj).
709 """
710 d_head = getattr(self.config, "d_head", None)
711 if not d_head:
712 return
713 allowed_dims = {d_head}
714 v_head_dim = getattr(self, "_v_head_dim", None)
715 if v_head_dim:
716 allowed_dims.add(v_head_dim)
717 if mat.shape[0] % n_heads == 0 and mat.shape[0] // n_heads in allowed_dims:
718 return
719 raise PerLayerGeometryError(
720 f"Cannot split {tuple(mat.shape)} weight on '{self.name}' into "
721 f"{n_heads} heads of d_head in {sorted(allowed_dims)}: this "
722 "layer's attention geometry differs from the config-level scalars "
723 "(per-layer num_key_value_heads/head_dim, e.g. gemma4 KV-shared "
724 "or K==V layers). Read the wrapped module's weight directly for "
725 "this layer instead."
726 )
728 def _project_per_head_qkv(
729 self,
730 linear_bridge: "GeneralizedComponent",
731 input_4d: torch.Tensor,
732 n_heads: int,
733 d_head: int,
734 ) -> torch.Tensor:
735 """Per-head Q/K/V projection over a 4D residual fork.
737 Plain nn.Linear applied to [batch, pos, H, d_model] broadcasts the
738 same weight across heads' copies — which for the split-qkv fork means
739 head h's copy sees every head's W rows, not just head h's. This routes
740 head h's copy through head h's W slice only via a per-head einsum.
742 Fires `linear_bridge.hook_out` on the flat 3D tensor so the hook sees
743 the same shape as the default path and downstream code receives a
744 consistent 4D `[B, S, H, d_head]` regardless of whether the user's
745 hook modified the tensor (which would otherwise trigger the
746 `hook_conversion.revert` 4D→3D flatten).
747 """
748 component = linear_bridge.original_component
749 assert component is not None, "LinearBridge.original_component not set"
750 weight = require_readable_weight(
751 component.weight,
752 operation="project per head (use_split_qkv_input / use_attn_in)",
753 owner=component,
754 )
755 bias = component.bias
756 w3d = einops.rearrange(
757 weight,
758 "(n_heads d_head) d_model -> n_heads d_model d_head",
759 n_heads=n_heads,
760 d_head=d_head,
761 )
762 out = torch.einsum("bshd,hde->bshe", input_4d, w3d)
763 if bias is not None:
764 b2d = einops.rearrange(bias, "(n_heads d_head) -> n_heads d_head", n_heads=n_heads)
765 assert isinstance(b2d, torch.Tensor)
766 out = out + b2d
767 # Flatten to 3D for hook_out (matches default-path shape); the
768 # hook_conversion reshapes to 4D for the user's fwd_hook, then reverts
769 # to 3D if the hook returned a modified tensor. Return 4D always.
770 b, s = out.shape[0], out.shape[1]
771 out_flat = out.reshape(b, s, n_heads * d_head)
772 out_flat = linear_bridge.hook_out(out_flat)
773 return out_flat.reshape(b, s, n_heads, d_head)
775 def _compute_per_head_result(
776 self,
777 z_4d: torch.Tensor,
778 n_heads: int,
779 d_head: int,
780 ) -> torch.Tensor:
781 """Per-head attention output pre-sum across heads.
783 Computes (z[..., h, :] @ W_O_per_head[h]) for each head h, fires
784 hook_result on the resulting [batch, pos, n_heads, d_model], then sums
785 across heads and adds b_O. Distributive over weight folding
786 (`sum_h z_h @ W_O_h + b_O == z_flat @ W_O.T + b_O`), so compat-mode and
787 raw-weight paths produce identical logits.
788 """
789 o = self.o.original_component
790 weight = o.weight
791 bias = getattr(o, "bias", None)
792 # HF Conv1D (GPT-2, GPT-J, CodeGen) stores weight as [in, out]; nn.Linear
793 # stores [out, in]. When W_O is square (d_model == n_heads*d_head, which
794 # is the common case), shape alone is ambiguous — dispatch on module
795 # type instead.
796 weight_is_in_out = type(o).__name__ == "Conv1D"
797 if weight_is_in_out:
798 w_per_head = einops.rearrange(
799 weight,
800 "(n_heads d_head) d_model -> n_heads d_head d_model",
801 n_heads=n_heads,
802 d_head=d_head,
803 )
804 else:
805 w_per_head = einops.rearrange(
806 weight,
807 "d_model (n_heads d_head) -> n_heads d_head d_model",
808 n_heads=n_heads,
809 d_head=d_head,
810 )
811 per_head = torch.einsum("bshd,hdm->bshm", z_4d, w_per_head)
812 per_head = self.hook_result(per_head)
813 summed = per_head.sum(dim=-2)
814 if bias is not None:
815 summed = summed + bias
816 return summed
818 def forward(self, *args: Any, **kwargs: Any) -> Any:
819 """Simplified forward pass - minimal wrapping around original component.
821 This does minimal wrapping: hook_in → delegate to HF → hook_out.
822 This ensures we match HuggingFace's exact output without complex intermediate processing.
824 Args:
825 *args: Input arguments to pass to the original component
826 **kwargs: Input keyword arguments to pass to the original component
828 Returns:
829 The output from the original component, with only input/output hooks applied
830 """
831 if self.original_component is None: 831 ↛ 832line 831 didn't jump to line 832 because the condition on line 831 was never true
832 raise RuntimeError(
833 f"Original component not set for {self.name}. Call set_original_component() first."
834 )
835 if "query_input" in kwargs:
836 hooked = self.hook_in(kwargs["query_input"])
837 kwargs["query_input"] = hooked
838 elif "hidden_states" in kwargs:
839 hooked = self.hook_in(kwargs["hidden_states"])
840 kwargs["hidden_states"] = hooked
841 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 841 ↛ 847line 841 didn't jump to line 847 because the condition on line 841 was always true
842 hooked = self.hook_in(args[0])
843 args = (hooked,) + args[1:]
844 # Modules exposing the pattern seam (NativeAttention) run the pattern
845 # through hook_pattern INSIDE the computation, so hook edits re-weight
846 # the output; post-hoc firing below cannot (see _fire_post_hoc_hook).
847 pattern_hooked_inside = bool(getattr(self.original_component, "accepts_pattern_fn", False))
848 if pattern_hooked_inside:
849 kwargs["pattern_fn"] = self.hook_pattern
850 # Relative-position-bias modules (T5 family) return position_bias, not
851 # attention weights, in output[1]: force the real weights into the
852 # output and capture Q/K to rebuild the pre-softmax scores.
853 reroute_position_bias = self.requires_relative_position_bias
854 caller_output_attentions = bool(kwargs.get("output_attentions", False))
855 qk_capture: Dict[str, torch.Tensor] = {}
856 capture_handles: list = []
857 if reroute_position_bias:
858 kwargs["output_attentions"] = True
859 capture_handles = self._capture_qk_outputs(qk_capture)
860 # try/finally so the captured tensor (and its autograd graph) is
861 # released even if original_component raises.
862 try:
863 output = self.original_component(*args, **kwargs)
864 finally:
865 self._captured_pre_ln_residual = None
866 for handle in capture_handles:
867 handle.remove()
868 if isinstance(output, tuple) and len(output) >= 2:
869 # output[0] is attention output
870 # output[1] may be attention weights (pattern) or position_bias (T5)
871 # Additional elements may include position_bias, attention weights, etc.
872 attn_output = self.hook_out(output[0])
873 if reroute_position_bias:
874 output = (attn_output,) + self._fire_position_bias_module_hooks(
875 output[1:], qk_capture, caller_output_attentions
876 )
877 else:
878 second_element = output[1]
879 # Fire hook_pattern if the second element is attention weights (4D tensor)
880 if isinstance(second_element, torch.Tensor) and second_element.dim() == 4:
881 # This looks like attention weights [batch, heads, seq, seq]
882 if not pattern_hooked_inside:
883 second_element = self._fire_post_hoc_hook(
884 self.hook_pattern, second_element, "hook_pattern"
885 )
886 # Also store for potential hook_attn_scores (before softmax)
887 # Note: Most HF implementations return post-softmax weights
888 self._fire_post_hoc_hook(
889 self.hook_attn_scores, second_element, "hook_attn_scores"
890 )
891 output = (attn_output, second_element) + output[2:]
892 elif isinstance(output, tuple) and len(output) == 1: 892 ↛ 893line 892 didn't jump to line 893 because the condition on line 892 was never true
893 output = (self.hook_out(output[0]),)
894 else:
895 output = self.hook_out(output)
896 return output
898 def _capture_qk_outputs(self, store: Dict[str, torch.Tensor]) -> list:
899 """Capture the Q/K projection outputs during the wrapped module's forward.
901 The relative-position-bias path rebuilds pre-softmax scores from these
902 (the wrapped module exposes neither the scores nor a seam to hook them).
903 Returns the handles; the caller removes them after the forward.
904 """
905 handles = []
906 for key in ("q", "k"):
907 proj = self._modules.get(key)
908 if proj is None: 908 ↛ anywhereline 908 didn't jump anywhere: it always raised an exception.
909 continue
911 def make_capture(name: str):
912 def capture(module: Any, inputs: Any, out: Any) -> None:
913 if isinstance(out, torch.Tensor): 913 ↛ exitline 913 didn't return from function 'capture' because the condition on line 913 was always true
914 store[name] = out
916 return capture
918 handles.append(proj.register_forward_hook(make_capture(key)))
919 return handles
921 def _fire_position_bias_module_hooks(
922 self,
923 rest: Tuple[Any, ...],
924 qk_capture: Dict[str, torch.Tensor],
925 caller_output_attentions: bool,
926 ) -> Tuple[Any, ...]:
927 """Fire pattern/scores hooks for a module whose output[1] is position_bias.
929 ``rest`` is the wrapped module's output after the attention output:
930 (position_bias, attn_weights?) with the weights present because forward
931 forced ``output_attentions``. The forced element is stripped again
932 unless the caller asked for it, so upstream tuple indexing (the HF
933 stack reads fixed positions) stays intact. The bias itself gets no
934 hooks: it is a broadcast [1, heads, q, k] additive term, not a pattern.
935 """
936 position_bias = rest[0]
937 weights = rest[1] if len(rest) >= 2 else None
938 if not (isinstance(weights, torch.Tensor) and weights.dim() == 4): 938 ↛ 939line 938 didn't jump to line 939 because the condition on line 938 was never true
939 return rest
940 scores = self._reconstruct_bias_scores(qk_capture, position_bias, weights)
941 if scores is not None: 941 ↛ 943line 941 didn't jump to line 943 because the condition on line 941 was always true
942 self._fire_post_hoc_hook(self.hook_attn_scores, scores, "hook_attn_scores")
943 weights = self._fire_post_hoc_hook(self.hook_pattern, weights, "hook_pattern")
944 if caller_output_attentions:
945 return (position_bias, weights) + rest[2:]
946 return (position_bias,) + rest[2:]
948 def _reconstruct_bias_scores(
949 self,
950 qk_capture: Dict[str, torch.Tensor],
951 position_bias: Any,
952 attn_weights: torch.Tensor,
953 ) -> Optional[torch.Tensor]:
954 """Rebuild pre-softmax scores as q @ k^T + position_bias.
956 T5-family attention applies no 1/sqrt(d) scaling, and the returned bias
957 already folds in the attention mask, so this reproduces the exact tensor
958 HF softmaxed. Returns None when Q or K was not captured this forward
959 (e.g. cross-attention K served from an encoder-decoder cache) or the
960 captured shapes disagree with the weights (e.g. cached decode steps).
961 """
962 q = qk_capture.get("q")
963 k = qk_capture.get("k")
964 if q is None or k is None or q.dim() != 3 or k.dim() != 3: 964 ↛ 965line 964 didn't jump to line 965 because the condition on line 964 was never true
965 return None
966 batch, n_heads, q_len, k_len = attn_weights.shape
967 if ( 967 ↛ 975line 967 didn't jump to line 975 because the condition on line 967 was never true
968 q.shape[0] != batch
969 or q.shape[1] != q_len
970 or k.shape[0] != batch
971 or k.shape[1] != k_len
972 or q.shape[2] != k.shape[2]
973 or q.shape[2] % n_heads != 0
974 ):
975 return None
976 q4 = q.view(batch, q_len, n_heads, -1).transpose(1, 2)
977 k4 = k.view(batch, k_len, n_heads, -1).transpose(1, 2)
978 scores = torch.matmul(q4, k4.transpose(3, 2))
979 if isinstance(position_bias, torch.Tensor): 979 ↛ 981line 979 didn't jump to line 981 because the condition on line 979 was always true
980 scores = scores + position_bias
981 return scores
983 def _fire_post_hoc_hook(
984 self, hook: HookPoint, tensor: torch.Tensor, hook_label: str
985 ) -> torch.Tensor:
986 """Fire a hook on a tensor the wrapped HF module already consumed.
988 An edit returned from such a hook cannot flow back into the attention
989 output, so it earns a warning instead of silently doing nothing.
990 Read-only hooks (e.g. caching) return the tensor unchanged and stay
991 silent.
992 """
993 result = hook(tensor)
994 if isinstance(result, torch.Tensor) and result is not tensor:
995 try:
996 unchanged = result.shape == tensor.shape and torch.equal(result, tensor)
997 except RuntimeError:
998 unchanged = False
999 if not unchanged:
1000 warnings.warn(
1001 f"A hook on {self.name or type(self).__name__}.{hook_label} returned a "
1002 "modified tensor, but this attention delegates its computation to the "
1003 "wrapped HF module, which has already consumed the original values. The "
1004 "edit is discarded and does not affect the model's output.",
1005 UserWarning,
1006 stacklevel=2,
1007 )
1008 return result
1010 @property
1011 def W_Q(self) -> torch.Tensor:
1012 """Get W_Q in 3D format [n_heads, d_model, d_head].
1014 Gated query projections retain their live query-and-gate parameter;
1015 this analysis view selects the query rows interleaved within each head.
1016 """
1017 weight = require_readable_weight(
1018 self.q.weight, operation=f"read W_Q from {self.name}", owner=self.q
1019 )
1020 if weight.ndim == 2 and self.config is not None: 1020 ↛ 1041line 1020 didn't jump to line 1041 because the condition on line 1020 was always true
1021 n_heads = self._get_n_heads()
1022 in_out_layout = self._weight_layout_in_out(self.q)
1023 if getattr(self.config, "gated_q_proj", False):
1024 d_head = int(self.config.d_head)
1025 gated_width = n_heads * d_head * 2
1026 if in_out_layout is True: 1026 ↛ 1027line 1026 didn't jump to line 1027 because the condition on line 1026 was never true
1027 output_first_weight = weight.T
1028 elif in_out_layout is False or weight.shape[0] == gated_width: 1028 ↛ 1030line 1028 didn't jump to line 1030 because the condition on line 1028 was always true
1029 output_first_weight = weight
1030 elif weight.shape[1] == gated_width:
1031 output_first_weight = weight.T
1032 else:
1033 output_first_weight = None
1035 if output_first_weight is not None and output_first_weight.shape[0] == gated_width: 1035 ↛ 1040line 1035 didn't jump to line 1040 because the condition on line 1035 was always true
1036 # Preserve the live query-gate projection; W_Q is an analysis-only query view.
1037 per_head_weight = output_first_weight.reshape(n_heads, d_head * 2, -1)
1038 return per_head_weight[:, :d_head, :].transpose(-1, -2)
1040 return self._reshape_weight_to_3d(weight, n_heads, in_out_layout=in_out_layout)
1041 return weight
1043 @property
1044 def W_K(self) -> torch.Tensor:
1045 """Get W_K in 3D format [n_heads, d_model, d_head] (uses n_kv_heads for GQA)."""
1046 weight = require_readable_weight(
1047 self.k.weight, operation=f"read W_K from {self.name}", owner=self.k
1048 )
1049 if weight.ndim == 2 and self.config is not None: 1049 ↛ 1055line 1049 didn't jump to line 1055 because the condition on line 1049 was always true
1050 return self._reshape_weight_to_3d(
1051 weight,
1052 self._get_n_heads(use_kv=True),
1053 in_out_layout=self._weight_layout_in_out(self.k),
1054 )
1055 return weight
1057 @property
1058 def W_V(self) -> torch.Tensor:
1059 """Get W_V in 3D format [n_heads, d_model, d_head] (uses n_kv_heads for GQA)."""
1060 weight = require_readable_weight(
1061 self.v.weight, operation=f"read W_V from {self.name}", owner=self.v
1062 )
1063 if weight.ndim == 2 and self.config is not None: 1063 ↛ 1069line 1063 didn't jump to line 1069 because the condition on line 1063 was always true
1064 return self._reshape_weight_to_3d(
1065 weight,
1066 self._get_n_heads(use_kv=True),
1067 in_out_layout=self._weight_layout_in_out(self.v),
1068 )
1069 return weight
1071 @property
1072 def W_O(self) -> torch.Tensor:
1073 """Get W_O in 3D format [n_heads, d_head, d_model]."""
1074 weight = require_readable_weight(
1075 self.o.weight, operation=f"read W_O from {self.name}", owner=self.o
1076 )
1077 if weight.ndim == 2 and self.config is not None: 1077 ↛ 1084line 1077 didn't jump to line 1084 because the condition on line 1077 was always true
1078 return self._reshape_weight_to_3d(
1079 weight,
1080 self._get_n_heads(),
1081 pattern="o",
1082 in_out_layout=self._weight_layout_in_out(self.o),
1083 )
1084 return weight
1086 def _reshape_bias(
1087 self, bias: Optional[torch.Tensor], use_kv: bool = False
1088 ) -> Optional[torch.Tensor]:
1089 """Reshape 1D bias to [n_heads, d_head]."""
1090 if bias is not None and bias.ndim == 1 and self.config is not None:
1091 n_heads = self._get_n_heads(use_kv=use_kv)
1092 return einops.rearrange(bias, "(n_heads d_head) -> n_heads d_head", n_heads=n_heads)
1093 return bias
1095 @property
1096 def b_Q(self) -> Optional[torch.Tensor]:
1097 """Get b_Q in 2D format [n_heads, d_head]."""
1098 return self._reshape_bias(self.q.bias)
1100 @property
1101 def b_K(self) -> Optional[torch.Tensor]:
1102 """Get b_K in 2D format [n_heads, d_head] (uses n_kv_heads for GQA)."""
1103 return self._reshape_bias(self.k.bias, use_kv=True)
1105 @property
1106 def b_V(self) -> Optional[torch.Tensor]:
1107 """Get b_V in 2D format [n_heads, d_head] (uses n_kv_heads for GQA)."""
1108 return self._reshape_bias(self.v.bias, use_kv=True)
1110 @property
1111 def b_O(self) -> Optional[torch.Tensor]:
1112 """Get b_O bias from linear bridge."""
1113 return self.o.bias