Coverage for transformer_lens/model_bridge/generalized_components/block.py: 90%
260 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Block bridge component.
3This module contains the bridge component for transformer blocks.
4"""
5from __future__ import annotations
7import inspect
8import re
9import weakref
10from typing import Any, Callable, Dict, Optional, cast
12import torch
14from transformer_lens.hook_points import HookPoint
15from transformer_lens.model_bridge.exceptions import StopAtLayerException
16from transformer_lens.model_bridge.generalized_components.base import (
17 GeneralizedComponent,
18)
20# Layer-type variant submodule names. Tuple for deterministic iteration order.
21# Extend here when adding new hybrid variant types.
22VARIANT_SUBMODULE_NAMES: tuple[str, ...] = ("attn", "linear_attn", "mamba", "mixer", "ssm")
23_VARIANT_SUBMODULE_SET: frozenset[str] = frozenset(VARIANT_SUBMODULE_NAMES)
25# Infrastructure modules excluded from submodule introspection.
26_BLOCK_INTERNAL_MODULES: frozenset[str] = frozenset({"hook_in", "hook_out", "_original_component"})
28# Norm-module prefixes excluded from layer_types() labels.
29_NORM_PREFIXES: tuple[str, ...] = ("ln", "layer_norm", "norm", "rms")
32class BlockBridge(GeneralizedComponent):
33 """Bridge component for transformer blocks.
35 This component provides standardized input/output hooks and monkey-patches
36 HuggingFace blocks to insert hooks at positions matching HookedTransformer.
37 """
39 is_list_item: bool = True
40 hook_out_is_single_residual_stream: bool = True
41 # hook_mlp_in is a direct HookPoint on this class (not aliased) so it can
42 # fire on the MLP-branch entry (pre-ln2, or the MLP input on post-norm
43 # blocks); see __init__. The normalized mlp input stays at block.mlp.hook_in.
44 hook_aliases = {
45 "hook_resid_pre": "hook_in",
46 "hook_resid_mid": "ln2.hook_in",
47 "hook_resid_post": "hook_out",
48 "hook_attn_in": "attn.hook_attn_in",
49 "hook_attn_out": "attn.hook_out",
50 "hook_q_input": "attn.hook_q_input",
51 "hook_k_input": "attn.hook_k_input",
52 "hook_v_input": "attn.hook_v_input",
53 "hook_mlp_out": "mlp.hook_out",
54 }
56 def __init__(
57 self,
58 name: str,
59 config: Optional[Any] = None,
60 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
61 hook_alias_overrides: Optional[Dict[str, str]] = None,
62 mlp_reads_resid_directly: bool = False,
63 ):
64 """Initialize the block bridge.
66 Args:
67 name: The name of the component in the model
68 config: Optional configuration (unused for BlockBridge)
69 submodules: Dictionary of submodules to register
70 hook_alias_overrides: Optional dictionary to override default hook aliases.
71 For example, {"hook_attn_out": "ln1_post.hook_out"} will make hook_attn_out
72 point to ln1_post.hook_out instead of the default attn.hook_out.
73 mlp_reads_resid_directly: True for post-norm blocks where the MLP consumes
74 the mid-residual with no pre-MLP norm (OLMo 2 layout). Moves the
75 hook_mlp_in capture from ln2 (whose input there is the raw MLP output)
76 to the MLP itself.
77 """
78 # ln1_post/ln2_post redirect attn_out/mlp_out to match HookedTransformer's
79 # placement (hook fires after the post-norm, not before).
80 auto_overrides = {}
81 if submodules is not None:
82 if "ln1_post" in submodules:
83 auto_overrides["hook_attn_out"] = "ln1_post.hook_out"
84 if "ln2_post" in submodules:
85 auto_overrides["hook_mlp_out"] = "ln2_post.hook_out"
86 merged_overrides = {**auto_overrides, **(hook_alias_overrides or {})}
88 # Guard against the bug where a sequential block (attn + mlp) with no ln2
89 # silently points hook_resid_mid at the wrong tensor. Use
90 # ParallelBlockBridge for parallel-residual architectures.
91 # Skip the check on generic-container / attn-only uses (no mlp).
92 has_attn_like = submodules is not None and any(
93 k in submodules for k in _VARIANT_SUBMODULE_SET
94 )
95 has_mlp = submodules is not None and "mlp" in submodules
96 has_ln2 = submodules is not None and "ln2" in submodules
97 if has_attn_like and has_mlp and not has_ln2 and type(self) is BlockBridge: 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true
98 raise ValueError(
99 f"BlockBridge at '{name}': 'ln2' submodule not declared. "
100 f"Either declare ln2, or use ParallelBlockBridge for a "
101 f"parallel-residual architecture."
102 )
104 super().__init__(
105 name,
106 config,
107 submodules=submodules if submodules is not None else {},
108 hook_alias_overrides=merged_overrides if merged_overrides else None,
109 )
111 self._original_block_forward: Optional[Callable[..., Any]] = None
112 self._capture_hooks_wired: bool = False
113 self._capture_hook_handles: list[torch.utils.hooks.RemovableHandle] = []
114 # Fallback for _read_use_hook_mlp_in when block.config is None.
115 self._use_hook_mlp_in: bool = False
116 self.mlp_reads_resid_directly = mlp_reads_resid_directly
117 # Fires on the MLP-branch entry (pre-ln2, or the MLP input on post-norm
118 # blocks) when use_hook_mlp_in is set. See #1317.
119 self.hook_mlp_in = HookPoint()
121 def _wire_ln1_module(self) -> None:
122 """Keep the raw ln1 execution reference outside the ownership tree."""
123 from transformer_lens.model_bridge.generalized_components.attention import (
124 AttentionBridge,
125 )
127 ln1 = self.submodules.get("ln1") if self.submodules else None
128 attn = self.submodules.get("attn") if self.submodules else None
129 if not isinstance(attn, AttentionBridge):
130 return
132 ln1_module = None
133 if (
134 ln1 is not None
135 and getattr(attn, "supports_split_qkv_fork", False)
136 and getattr(ln1, "original_component", None) is not None
137 ):
138 ln1_module = ln1.original_component
140 attn._modules.pop("_ln1_module", None)
141 object.__setattr__(attn, "_ln1_module", ln1_module)
143 def _maybe_wire_capture_hooks(self) -> None:
144 """Install the block's capture hooks (split-qkv fork, hook_mlp_in).
146 Registered on the bridge submodule, not ``original_component`` — the
147 manual bridge forward never calls the raw module. Idempotent.
148 """
149 self._wire_ln1_module()
150 if self._capture_hooks_wired:
151 return
152 from transformer_lens.model_bridge.generalized_components.attention import (
153 AttentionBridge,
154 )
156 ln1 = self.submodules.get("ln1") if self.submodules else None
157 attn = self.submodules.get("attn") if self.submodules else None
158 if (
159 ln1 is not None
160 and isinstance(attn, AttentionBridge)
161 and getattr(attn, "supports_split_qkv_fork", False)
162 and getattr(ln1, "original_component", None) is not None
163 ):
164 attn_ref = cast(AttentionBridge, weakref.proxy(attn))
166 def _capture_pre_ln1(_module: torch.nn.Module, args: tuple) -> None:
167 if args and isinstance(args[0], torch.Tensor): 167 ↛ exitline 167 didn't return from function '_capture_pre_ln1' because the condition on line 167 was always true
168 attn_ref._captured_pre_ln_residual = args[0]
170 handle = ln1.register_forward_pre_hook(_capture_pre_ln1)
171 self._capture_hook_handles.append(handle)
173 # hook_mlp_in must capture the MLP-branch entry point: ln2's input on
174 # pre-norm blocks, the MLP's own input on post-norm blocks (where ln2
175 # follows the MLP and its input is the raw MLP output).
176 if self.mlp_reads_resid_directly:
177 capture_target = self.submodules.get("mlp") if self.submodules else None
178 else:
179 capture_target = self.submodules.get("ln2") if self.submodules else None
180 if (
181 capture_target is not None
182 and getattr(capture_target, "original_component", None) is not None
183 ):
184 hook_mlp_in = self.hook_mlp_in
185 block_ref = weakref.proxy(self)
187 def _capture_mlp_in(_module: torch.nn.Module, args: tuple) -> Any:
188 if not block_ref._read_use_hook_mlp_in():
189 return None
190 if args and isinstance(args[0], torch.Tensor): 190 ↛ 193line 190 didn't jump to line 193 because the condition on line 190 was always true
191 hooked = hook_mlp_in(args[0])
192 return (hooked,) + args[1:]
193 return None
195 handle = capture_target.register_forward_pre_hook(_capture_mlp_in)
196 self._capture_hook_handles.append(handle)
198 self._capture_hooks_wired = True
200 def _teardown_capture_hooks(self) -> None:
201 """Remove the capture hooks installed by _maybe_wire_capture_hooks (and subclass extensions)."""
202 for handle in self._capture_hook_handles:
203 handle.remove()
204 self._capture_hook_handles.clear()
205 self._capture_hooks_wired = False
207 def _read_use_hook_mlp_in(self) -> bool:
208 """Prefer ``block.config.use_hook_mlp_in``; fall back to the block-local flag."""
209 cfg = self.config
210 if cfg is not None and hasattr(cfg, "use_hook_mlp_in"):
211 return bool(cfg.use_hook_mlp_in)
212 return self._use_hook_mlp_in
214 def _clear_attention_capture(self) -> None:
215 """Release the transient residual captured for attention input forks."""
216 from transformer_lens.model_bridge.generalized_components.attention import (
217 AttentionBridge,
218 )
220 attn = self.submodules.get("attn") if self.submodules else None
221 if isinstance(attn, AttentionBridge):
222 attn._captured_pre_ln_residual = None
224 def forward(self, *args: Any, **kwargs: Any) -> Any:
225 """Forward pass through the block bridge.
227 Args:
228 *args: Input arguments
229 **kwargs: Input keyword arguments
231 Returns:
232 The output from the original component
234 Raises:
235 StopAtLayerException: If stop_at_layer is set and this block should stop execution
236 """
237 if self.original_component is None: 237 ↛ 238line 237 didn't jump to line 238 because the condition on line 237 was never true
238 raise RuntimeError(
239 f"Original component not set for {self.name}. Call set_original_component() first."
240 )
242 self._maybe_wire_capture_hooks()
243 self._clear_attention_capture()
244 args, kwargs = self._maybe_inject_start_residual(args, kwargs)
245 self._check_stop_at_layer(*args, **kwargs)
246 args, kwargs = self._hook_input_hidden_states(args, kwargs)
248 # Filter kwargs to only include parameters accepted by the original component
249 # This prevents errors when passing encoder-specific params to decoder-only models
250 filtered_kwargs = self._filter_kwargs_for_forward(kwargs, len(args))
252 try:
253 output = self.original_component(*args, **filtered_kwargs)
254 finally:
255 self._clear_attention_capture()
256 force_tuple_for_bare_tensor = self._is_standalone_hidden_state_call(args, filtered_kwargs)
257 return self._apply_output_hook(
258 output, force_tuple_for_bare_tensor=force_tuple_for_bare_tensor
259 )
261 def _apply_output_hook(
262 self,
263 output: Any,
264 wrap_single_element: bool = True,
265 force_tuple_for_bare_tensor: bool = False,
266 ) -> Any:
267 """Hook the primary tensor in the output and return the result.
269 Args:
270 output: Raw output from the original component (tensor or tuple).
271 wrap_single_element: If True, single-element tuples stay as tuples after
272 hooking (default, required by most HF models). If False, single-element
273 tuples are unwrapped to a bare tensor (Bloom convention).
274 force_tuple_for_bare_tensor: If True, bare tensor outputs are wrapped into
275 a one-element tuple after hooking. This keeps standalone BlockBridge
276 calls compatible with HF block APIs that expose tuple-like block outputs,
277 while preserving tensor outputs during newer HF parent-model execution.
278 """
279 if isinstance(output, tuple) and len(output) > 0:
280 first = output[0]
281 if isinstance(first, torch.Tensor): 281 ↛ 286line 281 didn't jump to line 286 because the condition on line 281 was always true
282 first = self.hook_out(first)
283 if len(output) == 1:
284 return (first,) if wrap_single_element else first
285 output = (first,) + output[1:]
286 return output
287 if isinstance(output, torch.Tensor):
288 output = self.hook_out(output)
289 if force_tuple_for_bare_tensor and wrap_single_element:
290 return (output,)
291 return output
292 return output
294 @staticmethod
295 def _is_standalone_hidden_state_call(args: tuple, kwargs: dict) -> bool:
296 """Return True for direct block(hidden_states) style calls.
298 Transformers versions differ on whether parent model loops expect block
299 outputs as tuples or tensors. We preserve the original tensor return during
300 full-model execution, but expose tuple-like output for standalone component
301 calls so `output[0]` does not accidentally drop the batch dimension.
302 """
303 if len(args) == 1 and isinstance(args[0], torch.Tensor) and not kwargs:
304 return True
305 return (
306 len(args) == 0
307 and set(kwargs.keys()) == {"hidden_states"}
308 and isinstance(kwargs["hidden_states"], torch.Tensor)
309 )
311 def _extract_layer_idx(self) -> Optional[int]:
312 """Parse this block's layer index from its name (TL/GPT-2/LLaMA patterns)."""
313 if self.name is None: 313 ↛ 314line 313 didn't jump to line 314 because the condition on line 313 was never true
314 return None
315 match = re.search(r"(?:^|\.)(?:blocks|h|layers)\.(\d+)", self.name)
316 return int(match.group(1)) if match else None
318 def _check_stop_at_layer(self, *args: Any, **kwargs: Any) -> None:
319 """Check if execution should stop before this block. Raises StopAtLayerException.
321 The _stop_at_layer_idx attribute is set by the bridge's forward method.
322 Supports TL/GPT-2/LLaMA naming patterns for layer index extraction.
323 """
324 if getattr(self, "_stop_at_layer_idx", None) is None:
325 return
326 layer_idx = self._extract_layer_idx()
327 if layer_idx is not None and layer_idx == self._stop_at_layer_idx:
328 if len(args) > 0 and isinstance(args[0], torch.Tensor): 328 ↛ 330line 328 didn't jump to line 330 because the condition on line 328 was always true
329 input_tensor = args[0]
330 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
331 input_tensor = kwargs["hidden_states"]
332 else:
333 raise ValueError(f"Cannot find input tensor to stop at layer {layer_idx}")
334 input_tensor = self.hook_in(input_tensor)
335 raise StopAtLayerException(input_tensor)
337 def _maybe_inject_start_residual(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
338 """If this is the start_at_layer block, swap in the caller's residual.
340 Mirror of ``_check_stop_at_layer``: the bridge's forward stashes the
341 residual-stream input on the block via ``_start_residual`` and sets
342 ``_start_at_layer_idx``. This block replaces its incoming hidden states
343 with that residual; ``_hook_input_hidden_states`` then fires ``hook_in``
344 on it, so ``hook_resid_pre`` reflects the injected value.
345 """
346 if getattr(self, "_start_at_layer_idx", None) is None:
347 return args, kwargs
348 if self._extract_layer_idx() != self._start_at_layer_idx:
349 return args, kwargs
350 residual = getattr(self, "_start_residual", None)
351 if residual is None: 351 ↛ 352line 351 didn't jump to line 352 because the condition on line 351 was never true
352 return args, kwargs
353 if len(args) > 0 and isinstance(args[0], torch.Tensor): 353 ↛ 355line 353 didn't jump to line 355 because the condition on line 353 was always true
354 args = (residual,) + args[1:]
355 elif "hidden_states" in kwargs:
356 kwargs = {**kwargs, "hidden_states": residual}
357 return args, kwargs
359 def _hook_input_hidden_states(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
360 """Apply hook_in to the hidden_states input, whether in args or kwargs."""
361 if len(args) > 0 and isinstance(args[0], torch.Tensor): 361 ↛ 364line 361 didn't jump to line 364 because the condition on line 361 was always true
362 hooked_input = self.hook_in(args[0])
363 args = (hooked_input,) + args[1:]
364 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
365 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
366 return args, kwargs
368 def _filter_kwargs_for_forward(
369 self, kwargs: Dict[str, Any], num_positional_args: int = 0
370 ) -> Dict[str, Any]:
371 """Filter kwargs to only include parameters accepted by original_component.forward().
373 This prevents TypeErrors when the bridge passes parameters (like encoder_attention_mask)
374 that aren't accepted by decoder-only models. It also removes any kwargs that would
375 conflict with positional arguments already being passed.
377 Args:
378 kwargs: The full set of keyword arguments
379 num_positional_args: Number of positional arguments being passed (to avoid conflicts)
381 Returns:
382 Filtered kwargs containing only accepted parameters
383 """
384 if self.original_component is None: 384 ↛ 385line 384 didn't jump to line 385 because the condition on line 384 was never true
385 return kwargs
387 try:
388 # Get the signature of the original component's forward method
389 sig = inspect.signature(self.original_component.forward)
390 param_list = list(sig.parameters.keys())
391 valid_params = set(param_list)
393 # Check if the signature accepts **kwargs (VAR_KEYWORD)
394 accepts_var_keyword = any(
395 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
396 )
398 # If it accepts **kwargs, pass everything through
399 if accepts_var_keyword:
400 return kwargs
402 # Skip params already provided positionally
403 positional_param_names = set(param_list[:num_positional_args])
405 # Filter kwargs: include only if in signature AND not already provided positionally
406 filtered = {
407 k: v
408 for k, v in kwargs.items()
409 if k in valid_params and k not in positional_param_names
410 }
411 return filtered
413 except (ValueError, TypeError):
414 # If we can't inspect the signature, pass through all kwargs
415 # (better to potentially fail than to silently drop important params)
416 return kwargs
419class MLABlockBridge(BlockBridge):
420 """Block wrapping Multi-Head Latent Attention (DeepSeek V2/V3/R1).
422 MLA has no standalone q/k/v projections — Q flows through compressed
423 q_a_proj→q_a_layernorm→q_b_proj, and K/V share a joint kv_a_proj_with_mqa
424 entry point. There is no single HookPoint that represents "input that
425 becomes Q/K/V", so the block-level ``hook_q_input``/``hook_k_input``/
426 ``hook_v_input``/``hook_attn_in`` aliases do not apply. Type-level
427 distinction means a reader of the adapter sees ``MLABlockBridge`` and
428 knows those hooks are absent.
429 """
431 def __init__(
432 self,
433 name: str,
434 config: Optional[Any] = None,
435 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
436 hook_alias_overrides: Optional[Dict[str, str]] = None,
437 ):
438 super().__init__(
439 name,
440 config=config,
441 submodules=submodules,
442 hook_alias_overrides=hook_alias_overrides,
443 )
444 if self.hook_aliases is BlockBridge.hook_aliases: 444 ↛ 446line 444 didn't jump to line 446 because the condition on line 444 was always true
445 self.hook_aliases = dict(self.hook_aliases)
446 for alias in ("hook_q_input", "hook_k_input", "hook_v_input", "hook_attn_in"):
447 self.hook_aliases.pop(alias, None)
450class ParallelBlockBridge(BlockBridge):
451 """Block where attn and MLP both read the pre-attention residual.
453 For GPT-J, NeoX, Pythia, Phi, Cohere, CodeGen, and some Falcon variants,
454 output = resid_pre + attn_out + mlp_out — no distinct post-attention
455 residual exists. Matches legacy HookedTransformer which omits hook_resid_mid
456 when ``cfg.parallel_attn_mlp=True``. Type-level distinction means a reader
457 of the adapter sees ``ParallelBlockBridge`` and knows the hook is absent.
458 """
460 def __init__(
461 self,
462 name: str,
463 config: Optional[Any] = None,
464 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
465 hook_alias_overrides: Optional[Dict[str, str]] = None,
466 ):
467 super().__init__(
468 name,
469 config=config,
470 submodules=submodules,
471 hook_alias_overrides=hook_alias_overrides,
472 )
473 # Ensure instance-level copy before mutating; base may have left the
474 # class-level dict shared when no overrides were passed.
475 if self.hook_aliases is BlockBridge.hook_aliases: 475 ↛ 477line 475 didn't jump to line 477 because the condition on line 475 was always true
476 self.hook_aliases = dict(self.hook_aliases)
477 self.hook_aliases.pop("hook_resid_mid", None)
480class ScaledResidualBlockBridge(BlockBridge):
481 """Block whose sublayer outputs are scaled before the residual add.
483 Granite-family HF blocks compute ``residual + sublayer_out * residual_multiplier``
484 inline, so no submodule output equals the tensor added to the residual stream
485 and the legacy aliases cannot be fixed by re-pointing. ``hook_attn_out`` /
486 ``hook_mlp_out`` become real HookPoints on the block firing on the scaled
487 contribution:
489 - no hooks attached: the forward is untouched (bit-exact);
490 - read-only hooks: they observe ``module_out * scale``, forward stays bit-exact;
491 - a hook that changes the tensor (returned new or mutated in place): the module
492 output is rewritten to ``hooked / scale`` so HF's multiply reconstructs the
493 written value as the contribution (~1-ulp rounding; exact for zero-ablation);
494 - any backward hooks: the rewrite always happens so the autograd graph routes
495 through the HookPoint.
496 """
498 def __init__(
499 self,
500 name: str,
501 config: Optional[Any] = None,
502 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
503 hook_alias_overrides: Optional[Dict[str, str]] = None,
504 mlp_reads_resid_directly: bool = False,
505 residual_contribution_scale: float = 1.0,
506 scaled_attn_submodule: str = "attn",
507 scaled_mlp_submodule: Optional[str] = "mlp",
508 ):
509 """scaled_mlp_submodule=None when no single submodule feeds the MLP-side
510 add (GraniteMoeHybrid sums moe + shared_mlp inline) — hook_mlp_out then
511 stays absent rather than firing with a partial tensor."""
512 super().__init__(
513 name,
514 config=config,
515 submodules=submodules,
516 hook_alias_overrides=hook_alias_overrides,
517 mlp_reads_resid_directly=mlp_reads_resid_directly,
518 )
519 scale = float(residual_contribution_scale)
520 if scale == 0.0:
521 raise ValueError(
522 f"ScaledResidualBlockBridge at '{name}': residual_contribution_scale "
523 f"must be nonzero (the write path divides by it)."
524 )
525 self.residual_contribution_scale = scale
526 self.scaled_attn_submodule = scaled_attn_submodule
527 self.scaled_mlp_submodule = scaled_mlp_submodule
528 if self.hook_aliases is BlockBridge.hook_aliases: 528 ↛ 530line 528 didn't jump to line 530 because the condition on line 528 was always true
529 self.hook_aliases = dict(self.hook_aliases)
530 for alias in ("hook_attn_out", "hook_mlp_out"):
531 self.hook_aliases.pop(alias, None)
532 self.hook_attn_out = HookPoint()
533 if scaled_mlp_submodule is not None:
534 self.hook_mlp_out = HookPoint()
536 def set_original_component(self, original_component: torch.nn.Module) -> None:
537 """Prune contribution HookPoints the bound layer cannot fire.
539 Heterogeneous blocks (GraniteMoeHybrid mamba layers) lack the attn
540 submodule; a HookPoint that exists but never fires is a silent-no-op
541 intervention trap, so the name must be absent instead.
542 """
543 super().set_original_component(original_component)
544 targets = [(self.scaled_attn_submodule, "hook_attn_out")]
545 if self.scaled_mlp_submodule is not None:
546 targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
547 for sub_name, hook_name in targets:
548 sub = self.submodules.get(sub_name) if self.submodules else None
549 remote = getattr(sub, "name", None)
550 first = remote.split(".", 1)[0] if isinstance(remote, str) else None
551 missing = sub is None or (
552 first is not None and getattr(original_component, first, None) is None
553 )
554 if missing:
555 if hasattr(self, hook_name): 555 ↛ 559line 555 didn't jump to line 559 because the condition on line 555 was always true
556 delattr(self, hook_name)
557 # __setattr__ auto-registered the HookPoint; get_hooks() serves
558 # from this registry, so the module deletion alone is not enough.
559 self._hook_registry.pop(hook_name, None)
561 def _maybe_wire_capture_hooks(self) -> None:
562 """Extend the base wiring with the scaled-contribution forward hooks.
564 Shares the base flag and handle list so _teardown_capture_hooks also
565 removes these hooks and re-wiring stays idempotent.
566 """
567 if self._capture_hooks_wired:
568 return
569 super()._maybe_wire_capture_hooks()
570 targets = [(self.scaled_attn_submodule, "hook_attn_out")]
571 if self.scaled_mlp_submodule is not None:
572 targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
573 for sub_name, hook_name in targets:
574 hook_point = getattr(self, hook_name, None)
575 sub = self.submodules.get(sub_name) if self.submodules else None
576 if (
577 hook_point is None
578 or sub is None
579 or getattr(sub, "original_component", None) is None
580 ):
581 continue
582 handle = sub.register_forward_hook(self._make_scaled_contribution_hook(hook_point))
583 self._capture_hook_handles.append(handle)
585 def _make_scaled_contribution_hook(
586 self, hook_point: HookPoint
587 ) -> Callable[[torch.nn.Module, tuple, Any], Any]:
588 """Build a forward hook exposing ``output * scale`` through hook_point."""
589 scale = self.residual_contribution_scale
591 def _hook(_module: torch.nn.Module, _args: tuple, output: Any) -> Any:
592 if not hook_point.has_hooks(dir="both"):
593 return None
594 is_tuple = isinstance(output, tuple)
595 out = output[0] if is_tuple else output
596 if not isinstance(out, torch.Tensor): 596 ↛ 597line 596 didn't jump to line 597 because the condition on line 596 was never true
597 return None
598 scaled = out * scale
599 hooked = hook_point(scaled)
600 # Compare against a freshly computed reference: an in-place mutation
601 # alters `scaled` itself, so an identity check would miss it.
602 if not hook_point.has_hooks(dir="bwd") and torch.equal(hooked, out * scale):
603 return None
604 new = hooked / scale
605 return ((new,) + output[1:]) if is_tuple else new
607 return _hook
610class DelegatedAttentionBlockBridge(BlockBridge):
611 """Block whose attention is delegated wholesale to HF (no split-qkv fork).
613 For architectures with heterogeneous per-layer attention structure — e.g.
614 Gemma 4, where KV-shared layers have no ``k_proj``/``v_proj`` at all and
615 K==V layers have no ``v_proj`` — there is no uniform HookPoint that
616 represents "input that becomes Q/K/V", so the block-level ``hook_q_input``/
617 ``hook_k_input``/``hook_v_input``/``hook_attn_in`` aliases do not apply.
618 Type-level distinction means a reader of the adapter sees
619 ``DelegatedAttentionBlockBridge`` and knows those hooks are absent.
620 """
622 # Tell the component benchmark this block's attention is delegated wholesale
623 # to HF and cannot be tested in isolation (requires model-specific kwargs).
624 maintain_native_attention: bool = True
626 def __init__(
627 self,
628 name: str,
629 config: Optional[Any] = None,
630 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
631 hook_alias_overrides: Optional[Dict[str, str]] = None,
632 ):
633 super().__init__(
634 name,
635 config=config,
636 submodules=submodules,
637 hook_alias_overrides=hook_alias_overrides,
638 )
639 if self.hook_aliases is BlockBridge.hook_aliases:
640 self.hook_aliases = dict(self.hook_aliases)
641 for alias in ("hook_q_input", "hook_k_input", "hook_v_input", "hook_attn_in"):
642 self.hook_aliases.pop(alias, None)