Coverage for transformer_lens/model_bridge/generalized_components/block.py: 90%
250 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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 the sequential-block case: attn + mlp with no ln2 would silently
89 # point hook_resid_mid at the wrong tensor. Use ParallelBlockBridge for
90 # 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 # Call parent with merged overrides
105 super().__init__(
106 name,
107 config,
108 submodules=submodules if submodules is not None else {},
109 hook_alias_overrides=merged_overrides if merged_overrides else None,
110 )
112 self._original_block_forward: Optional[Callable[..., Any]] = None
113 self._capture_hooks_wired: bool = False
114 self._capture_hook_handles: list[torch.utils.hooks.RemovableHandle] = []
115 # Fallback for _read_use_hook_mlp_in when block.config is None.
116 self._use_hook_mlp_in: bool = False
117 self.mlp_reads_resid_directly = mlp_reads_resid_directly
118 # Fires on the MLP-branch entry (pre-ln2, or the MLP input on post-norm
119 # blocks) when use_hook_mlp_in is set. See #1317.
120 self.hook_mlp_in = HookPoint()
122 def _wire_ln1_module(self) -> None:
123 """Keep the raw ln1 execution reference outside the ownership tree."""
124 from transformer_lens.model_bridge.generalized_components.attention import (
125 AttentionBridge,
126 )
128 ln1 = self.submodules.get("ln1") if self.submodules else None
129 attn = self.submodules.get("attn") if self.submodules else None
130 if not isinstance(attn, AttentionBridge):
131 return
133 ln1_module = None
134 if (
135 ln1 is not None
136 and getattr(attn, "supports_split_qkv_fork", False)
137 and getattr(ln1, "original_component", None) is not None
138 ):
139 ln1_module = ln1.original_component
141 attn._modules.pop("_ln1_module", None)
142 object.__setattr__(attn, "_ln1_module", ln1_module)
144 def _maybe_wire_capture_hooks(self) -> None:
145 """Install the block's capture hooks (split-qkv fork, hook_mlp_in).
147 Registered on the bridge submodule, not ``original_component`` — the
148 manual bridge forward never calls the raw module. Idempotent.
149 """
150 self._wire_ln1_module()
151 if self._capture_hooks_wired:
152 return
153 from transformer_lens.model_bridge.generalized_components.attention import (
154 AttentionBridge,
155 )
157 ln1 = self.submodules.get("ln1") if self.submodules else None
158 attn = self.submodules.get("attn") if self.submodules else None
159 if (
160 ln1 is not None
161 and isinstance(attn, AttentionBridge)
162 and getattr(attn, "supports_split_qkv_fork", False)
163 and getattr(ln1, "original_component", None) is not None
164 ):
165 attn_ref = cast(AttentionBridge, weakref.proxy(attn))
167 def _capture_pre_ln1(_module: torch.nn.Module, args: tuple) -> None:
168 if args and isinstance(args[0], torch.Tensor): 168 ↛ exitline 168 didn't return from function '_capture_pre_ln1' because the condition on line 168 was always true
169 attn_ref._captured_pre_ln_residual = args[0]
171 handle = ln1.register_forward_pre_hook(_capture_pre_ln1)
172 self._capture_hook_handles.append(handle)
174 # hook_mlp_in must capture the MLP-branch entry point: ln2's input on
175 # pre-norm blocks, the MLP's own input on post-norm blocks (where ln2
176 # follows the MLP and its input is the raw MLP output).
177 if self.mlp_reads_resid_directly:
178 capture_target = self.submodules.get("mlp") if self.submodules else None
179 else:
180 capture_target = self.submodules.get("ln2") if self.submodules else None
181 if (
182 capture_target is not None
183 and getattr(capture_target, "original_component", None) is not None
184 ):
185 hook_mlp_in = self.hook_mlp_in
186 block_ref = weakref.proxy(self)
188 def _capture_mlp_in(_module: torch.nn.Module, args: tuple) -> Any:
189 if not block_ref._read_use_hook_mlp_in():
190 return None
191 if args and isinstance(args[0], torch.Tensor): 191 ↛ 194line 191 didn't jump to line 194 because the condition on line 191 was always true
192 hooked = hook_mlp_in(args[0])
193 return (hooked,) + args[1:]
194 return None
196 handle = capture_target.register_forward_pre_hook(_capture_mlp_in)
197 self._capture_hook_handles.append(handle)
199 self._capture_hooks_wired = True
201 def _teardown_capture_hooks(self) -> None:
202 """Remove the capture hooks installed by _maybe_wire_capture_hooks (and subclass extensions)."""
203 for handle in self._capture_hook_handles:
204 handle.remove()
205 self._capture_hook_handles.clear()
206 self._capture_hooks_wired = False
208 def _read_use_hook_mlp_in(self) -> bool:
209 """Prefer ``block.config.use_hook_mlp_in``; fall back to the block-local flag."""
210 cfg = self.config
211 if cfg is not None and hasattr(cfg, "use_hook_mlp_in"):
212 return bool(cfg.use_hook_mlp_in)
213 return self._use_hook_mlp_in
215 def _clear_attention_capture(self) -> None:
216 """Release the transient residual captured for attention input forks."""
217 from transformer_lens.model_bridge.generalized_components.attention import (
218 AttentionBridge,
219 )
221 attn = self.submodules.get("attn") if self.submodules else None
222 if isinstance(attn, AttentionBridge):
223 attn._captured_pre_ln_residual = None
225 def forward(self, *args: Any, **kwargs: Any) -> Any:
226 """Forward pass through the block bridge.
228 Args:
229 *args: Input arguments
230 **kwargs: Input keyword arguments
232 Returns:
233 The output from the original component
235 Raises:
236 StopAtLayerException: If stop_at_layer is set and this block should stop execution
237 """
238 if self.original_component is None: 238 ↛ 239line 238 didn't jump to line 239 because the condition on line 238 was never true
239 raise RuntimeError(
240 f"Original component not set for {self.name}. Call set_original_component() first."
241 )
243 self._maybe_wire_capture_hooks()
244 self._clear_attention_capture()
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:
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 not (hasattr(self, "_stop_at_layer_idx") and self._stop_at_layer_idx is not None):
325 return
326 if self.name is not None: 326 ↛ 331line 326 didn't jump to line 331 because the condition on line 326 was always true
327 # Anchored alternation: native models name blocks top-level
328 # ("layers.0"), which the old leading-dot pattern missed entirely.
329 match = re.search(r"(?:^|\.)(?:blocks|h|layers)\.(\d+)", self.name)
330 else:
331 match = None
332 if match: 332 ↛ exitline 332 didn't return from function '_check_stop_at_layer' because the condition on line 332 was always true
333 layer_idx = int(match.group(1))
334 if layer_idx == self._stop_at_layer_idx:
335 if len(args) > 0 and isinstance(args[0], torch.Tensor): 335 ↛ 337line 335 didn't jump to line 337 because the condition on line 335 was always true
336 input_tensor = args[0]
337 elif "hidden_states" in kwargs and isinstance(
338 kwargs["hidden_states"], torch.Tensor
339 ):
340 input_tensor = kwargs["hidden_states"]
341 else:
342 raise ValueError(f"Cannot find input tensor to stop at layer {layer_idx}")
343 input_tensor = self.hook_in(input_tensor)
344 raise StopAtLayerException(input_tensor)
346 def _hook_input_hidden_states(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
347 """Apply hook_in to the hidden_states input, whether in args or kwargs."""
348 if len(args) > 0 and isinstance(args[0], torch.Tensor): 348 ↛ 351line 348 didn't jump to line 351 because the condition on line 348 was always true
349 hooked_input = self.hook_in(args[0])
350 args = (hooked_input,) + args[1:]
351 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
352 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
353 return args, kwargs
355 def _filter_kwargs_for_forward(
356 self, kwargs: Dict[str, Any], num_positional_args: int = 0
357 ) -> Dict[str, Any]:
358 """Filter kwargs to only include parameters accepted by original_component.forward().
360 This prevents TypeErrors when the bridge passes parameters (like encoder_attention_mask)
361 that aren't accepted by decoder-only models. It also removes any kwargs that would
362 conflict with positional arguments already being passed.
364 Args:
365 kwargs: The full set of keyword arguments
366 num_positional_args: Number of positional arguments being passed (to avoid conflicts)
368 Returns:
369 Filtered kwargs containing only accepted parameters
370 """
371 if self.original_component is None: 371 ↛ 372line 371 didn't jump to line 372 because the condition on line 371 was never true
372 return kwargs
374 try:
375 # Get the signature of the original component's forward method
376 sig = inspect.signature(self.original_component.forward)
377 param_list = list(sig.parameters.keys())
378 valid_params = set(param_list)
380 # Check if the signature accepts **kwargs (VAR_KEYWORD)
381 accepts_var_keyword = any(
382 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
383 )
385 # If it accepts **kwargs, pass everything through
386 if accepts_var_keyword:
387 return kwargs
389 # Skip params already provided positionally
390 positional_param_names = set(param_list[:num_positional_args])
392 # Filter kwargs: include only if in signature AND not already provided positionally
393 filtered = {
394 k: v
395 for k, v in kwargs.items()
396 if k in valid_params and k not in positional_param_names
397 }
398 return filtered
400 except (ValueError, TypeError):
401 # If we can't inspect the signature, pass through all kwargs
402 # (better to potentially fail than to silently drop important params)
403 return kwargs
406class MLABlockBridge(BlockBridge):
407 """Block wrapping Multi-Head Latent Attention (DeepSeek V2/V3/R1).
409 MLA has no standalone q/k/v projections — Q flows through compressed
410 q_a_proj→q_a_layernorm→q_b_proj, and K/V share a joint kv_a_proj_with_mqa
411 entry point. There is no single HookPoint that represents "input that
412 becomes Q/K/V", so the block-level ``hook_q_input``/``hook_k_input``/
413 ``hook_v_input``/``hook_attn_in`` aliases do not apply. Type-level
414 distinction means a reader of the adapter sees ``MLABlockBridge`` and
415 knows those hooks are absent.
416 """
418 def __init__(
419 self,
420 name: str,
421 config: Optional[Any] = None,
422 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
423 hook_alias_overrides: Optional[Dict[str, str]] = None,
424 ):
425 super().__init__(
426 name,
427 config=config,
428 submodules=submodules,
429 hook_alias_overrides=hook_alias_overrides,
430 )
431 if self.hook_aliases is BlockBridge.hook_aliases: 431 ↛ 433line 431 didn't jump to line 433 because the condition on line 431 was always true
432 self.hook_aliases = dict(self.hook_aliases)
433 for alias in ("hook_q_input", "hook_k_input", "hook_v_input", "hook_attn_in"):
434 self.hook_aliases.pop(alias, None)
437class ParallelBlockBridge(BlockBridge):
438 """Block where attn and MLP both read the pre-attention residual.
440 For GPT-J, NeoX, Pythia, Phi, Cohere, CodeGen, and some Falcon variants,
441 output = resid_pre + attn_out + mlp_out — no distinct post-attention
442 residual exists. Matches legacy HookedTransformer which omits hook_resid_mid
443 when ``cfg.parallel_attn_mlp=True``. Type-level distinction means a reader
444 of the adapter sees ``ParallelBlockBridge`` and knows the hook is absent.
445 """
447 def __init__(
448 self,
449 name: str,
450 config: Optional[Any] = None,
451 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
452 hook_alias_overrides: Optional[Dict[str, str]] = None,
453 ):
454 super().__init__(
455 name,
456 config=config,
457 submodules=submodules,
458 hook_alias_overrides=hook_alias_overrides,
459 )
460 # Ensure instance-level copy before mutating; base may have left the
461 # class-level dict shared when no overrides were passed.
462 if self.hook_aliases is BlockBridge.hook_aliases: 462 ↛ 464line 462 didn't jump to line 464 because the condition on line 462 was always true
463 self.hook_aliases = dict(self.hook_aliases)
464 self.hook_aliases.pop("hook_resid_mid", None)
467class ScaledResidualBlockBridge(BlockBridge):
468 """Block whose sublayer outputs are scaled before the residual add.
470 Granite-family HF blocks compute ``residual + sublayer_out * residual_multiplier``
471 inline, so no submodule output equals the tensor added to the residual stream
472 and the legacy aliases cannot be fixed by re-pointing. ``hook_attn_out`` /
473 ``hook_mlp_out`` become real HookPoints on the block firing on the scaled
474 contribution:
476 - no hooks attached: the forward is untouched (bit-exact);
477 - read-only hooks: they observe ``module_out * scale``, forward stays bit-exact;
478 - a hook that changes the tensor (returned new or mutated in place): the module
479 output is rewritten to ``hooked / scale`` so HF's multiply reconstructs the
480 written value as the contribution (~1-ulp rounding; exact for zero-ablation);
481 - any backward hooks: the rewrite always happens so the autograd graph routes
482 through the HookPoint.
483 """
485 def __init__(
486 self,
487 name: str,
488 config: Optional[Any] = None,
489 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
490 hook_alias_overrides: Optional[Dict[str, str]] = None,
491 mlp_reads_resid_directly: bool = False,
492 residual_contribution_scale: float = 1.0,
493 scaled_attn_submodule: str = "attn",
494 scaled_mlp_submodule: Optional[str] = "mlp",
495 ):
496 """scaled_mlp_submodule=None when no single submodule feeds the MLP-side
497 add (GraniteMoeHybrid sums moe + shared_mlp inline) — hook_mlp_out then
498 stays absent rather than firing with a partial tensor."""
499 super().__init__(
500 name,
501 config=config,
502 submodules=submodules,
503 hook_alias_overrides=hook_alias_overrides,
504 mlp_reads_resid_directly=mlp_reads_resid_directly,
505 )
506 scale = float(residual_contribution_scale)
507 if scale == 0.0:
508 raise ValueError(
509 f"ScaledResidualBlockBridge at '{name}': residual_contribution_scale "
510 f"must be nonzero (the write path divides by it)."
511 )
512 self.residual_contribution_scale = scale
513 self.scaled_attn_submodule = scaled_attn_submodule
514 self.scaled_mlp_submodule = scaled_mlp_submodule
515 if self.hook_aliases is BlockBridge.hook_aliases: 515 ↛ 517line 515 didn't jump to line 517 because the condition on line 515 was always true
516 self.hook_aliases = dict(self.hook_aliases)
517 for alias in ("hook_attn_out", "hook_mlp_out"):
518 self.hook_aliases.pop(alias, None)
519 self.hook_attn_out = HookPoint()
520 if scaled_mlp_submodule is not None:
521 self.hook_mlp_out = HookPoint()
523 def set_original_component(self, original_component: torch.nn.Module) -> None:
524 """Prune contribution HookPoints the bound layer cannot fire.
526 Heterogeneous blocks (GraniteMoeHybrid mamba layers) lack the attn
527 submodule; a HookPoint that exists but never fires is a silent-no-op
528 intervention trap, so the name must be absent instead.
529 """
530 super().set_original_component(original_component)
531 targets = [(self.scaled_attn_submodule, "hook_attn_out")]
532 if self.scaled_mlp_submodule is not None:
533 targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
534 for sub_name, hook_name in targets:
535 sub = self.submodules.get(sub_name) if self.submodules else None
536 remote = getattr(sub, "name", None)
537 first = remote.split(".", 1)[0] if isinstance(remote, str) else None
538 missing = sub is None or (
539 first is not None and getattr(original_component, first, None) is None
540 )
541 if missing:
542 if hasattr(self, hook_name): 542 ↛ 546line 542 didn't jump to line 546 because the condition on line 542 was always true
543 delattr(self, hook_name)
544 # __setattr__ auto-registered the HookPoint; get_hooks() serves
545 # from this registry, so the module deletion alone is not enough.
546 self._hook_registry.pop(hook_name, None)
548 def _maybe_wire_capture_hooks(self) -> None:
549 """Extend the base wiring with the scaled-contribution forward hooks.
551 Shares the base flag and handle list so _teardown_capture_hooks also
552 removes these hooks and re-wiring stays idempotent.
553 """
554 if self._capture_hooks_wired:
555 return
556 super()._maybe_wire_capture_hooks()
557 targets = [(self.scaled_attn_submodule, "hook_attn_out")]
558 if self.scaled_mlp_submodule is not None:
559 targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
560 for sub_name, hook_name in targets:
561 hook_point = getattr(self, hook_name, None)
562 sub = self.submodules.get(sub_name) if self.submodules else None
563 if (
564 hook_point is None
565 or sub is None
566 or getattr(sub, "original_component", None) is None
567 ):
568 continue
569 handle = sub.register_forward_hook(self._make_scaled_contribution_hook(hook_point))
570 self._capture_hook_handles.append(handle)
572 def _make_scaled_contribution_hook(
573 self, hook_point: HookPoint
574 ) -> Callable[[torch.nn.Module, tuple, Any], Any]:
575 """Build a forward hook exposing ``output * scale`` through hook_point."""
576 scale = self.residual_contribution_scale
578 def _hook(_module: torch.nn.Module, _args: tuple, output: Any) -> Any:
579 if not hook_point.has_hooks(dir="both"):
580 return None
581 is_tuple = isinstance(output, tuple)
582 out = output[0] if is_tuple else output
583 if not isinstance(out, torch.Tensor): 583 ↛ 584line 583 didn't jump to line 584 because the condition on line 583 was never true
584 return None
585 scaled = out * scale
586 hooked = hook_point(scaled)
587 # Compare against a freshly computed reference: an in-place mutation
588 # alters `scaled` itself, so an identity check would miss it.
589 if not hook_point.has_hooks(dir="bwd") and torch.equal(hooked, out * scale):
590 return None
591 new = hooked / scale
592 return ((new,) + output[1:]) if is_tuple else new
594 return _hook
597class DelegatedAttentionBlockBridge(BlockBridge):
598 """Block whose attention is delegated wholesale to HF (no split-qkv fork).
600 For architectures with heterogeneous per-layer attention structure — e.g.
601 Gemma 4, where KV-shared layers have no ``k_proj``/``v_proj`` at all and
602 K==V layers have no ``v_proj`` — there is no uniform HookPoint that
603 represents "input that becomes Q/K/V", so the block-level ``hook_q_input``/
604 ``hook_k_input``/``hook_v_input``/``hook_attn_in`` aliases do not apply.
605 Type-level distinction means a reader of the adapter sees
606 ``DelegatedAttentionBlockBridge`` and knows those hooks are absent.
607 """
609 # Tell the component benchmark this block's attention is delegated wholesale
610 # to HF and cannot be tested in isolation (requires model-specific kwargs).
611 maintain_native_attention: bool = True
613 def __init__(
614 self,
615 name: str,
616 config: Optional[Any] = None,
617 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
618 hook_alias_overrides: Optional[Dict[str, str]] = None,
619 ):
620 super().__init__(
621 name,
622 config=config,
623 submodules=submodules,
624 hook_alias_overrides=hook_alias_overrides,
625 )
626 if self.hook_aliases is BlockBridge.hook_aliases:
627 self.hook_aliases = dict(self.hook_aliases)
628 for alias in ("hook_q_input", "hook_k_input", "hook_v_input", "hook_attn_in"):
629 self.hook_aliases.pop(alias, None)