Coverage for transformer_lens/model_bridge/generalized_components/block.py: 90%
264 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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 when the MLP consumes the mid-residual with
74 no pre-MLP norm. Moves the hook_mlp_in capture from ln2 to the MLP
75 itself. Implied for post-norm-only blocks (ln2_post without ln2).
76 """
77 has_ln2 = submodules is not None and "ln2" in submodules
78 # OLMo 2 layout: ln2_post without ln2 means the MLP reads the mid-residual directly.
79 post_norm_only_mlp = submodules is not None and "ln2_post" in submodules and not has_ln2
80 # ln1_post/ln2_post redirect attn_out/mlp_out to match HookedTransformer's
81 # placement (hook fires after the post-norm, not before).
82 auto_overrides = {}
83 if submodules is not None:
84 if "ln1_post" in submodules:
85 auto_overrides["hook_attn_out"] = "ln1_post.hook_out"
86 if "ln2_post" in submodules:
87 auto_overrides["hook_mlp_out"] = "ln2_post.hook_out"
88 if post_norm_only_mlp:
89 auto_overrides["hook_resid_mid"] = "mlp.hook_in"
90 merged_overrides = {**auto_overrides, **(hook_alias_overrides or {})}
92 # Guard against the bug where a sequential block (attn + mlp) with no MLP-side
93 # norm silently points hook_resid_mid at the wrong tensor. Use
94 # ParallelBlockBridge for parallel-residual architectures.
95 # Skip the check on generic-container / attn-only uses (no mlp).
96 has_attn_like = submodules is not None and any(
97 k in submodules for k in _VARIANT_SUBMODULE_SET
98 )
99 has_mlp = submodules is not None and "mlp" in submodules
100 if (
101 has_attn_like
102 and has_mlp
103 and not has_ln2
104 and not post_norm_only_mlp
105 and type(self) is BlockBridge
106 ):
107 raise ValueError(
108 f"BlockBridge at '{name}': neither 'ln2' nor 'ln2_post' declared. "
109 f"Declare ln2 (pre-MLP norm) or ln2_post (post-MLP norm), or use "
110 f"ParallelBlockBridge for a parallel-residual architecture."
111 )
113 super().__init__(
114 name,
115 config,
116 submodules=submodules if submodules is not None else {},
117 hook_alias_overrides=merged_overrides if merged_overrides else None,
118 )
120 self._original_block_forward: Optional[Callable[..., Any]] = None
121 self._capture_hooks_wired: bool = False
122 self._capture_hook_handles: list[torch.utils.hooks.RemovableHandle] = []
123 # Fallback for _read_use_hook_mlp_in when block.config is None.
124 self._use_hook_mlp_in: bool = False
125 self.mlp_reads_resid_directly = mlp_reads_resid_directly or post_norm_only_mlp
126 # Fires on the MLP-branch entry (pre-ln2, or the MLP input on post-norm
127 # blocks) when use_hook_mlp_in is set. See #1317.
128 self.hook_mlp_in = HookPoint()
130 def _wire_ln1_module(self) -> None:
131 """Keep the raw ln1 execution reference outside the ownership tree."""
132 from transformer_lens.model_bridge.generalized_components.attention import (
133 AttentionBridge,
134 )
136 ln1 = self.submodules.get("ln1") if self.submodules else None
137 attn = self.submodules.get("attn") if self.submodules else None
138 if not isinstance(attn, AttentionBridge):
139 return
141 ln1_module = None
142 if (
143 ln1 is not None
144 and getattr(attn, "supports_split_qkv_fork", False)
145 and getattr(ln1, "original_component", None) is not None
146 ):
147 ln1_module = ln1.original_component
149 attn._modules.pop("_ln1_module", None)
150 object.__setattr__(attn, "_ln1_module", ln1_module)
152 def _maybe_wire_capture_hooks(self) -> None:
153 """Install the block's capture hooks (split-qkv fork, hook_mlp_in).
155 Registered on the bridge submodule, not ``original_component`` — the
156 manual bridge forward never calls the raw module. Idempotent.
157 """
158 self._wire_ln1_module()
159 if self._capture_hooks_wired:
160 return
161 from transformer_lens.model_bridge.generalized_components.attention import (
162 AttentionBridge,
163 )
165 ln1 = self.submodules.get("ln1") if self.submodules else None
166 attn = self.submodules.get("attn") if self.submodules else None
167 if (
168 ln1 is not None
169 and isinstance(attn, AttentionBridge)
170 and getattr(attn, "supports_split_qkv_fork", False)
171 and getattr(ln1, "original_component", None) is not None
172 ):
173 attn_ref = cast(AttentionBridge, weakref.proxy(attn))
175 def _capture_pre_ln1(_module: torch.nn.Module, args: tuple) -> None:
176 if args and isinstance(args[0], torch.Tensor): 176 ↛ exitline 176 didn't return from function '_capture_pre_ln1' because the condition on line 176 was always true
177 attn_ref._captured_pre_ln_residual = args[0]
179 handle = ln1.register_forward_pre_hook(_capture_pre_ln1)
180 self._capture_hook_handles.append(handle)
182 # hook_mlp_in must capture the MLP-branch entry point: ln2's input on
183 # pre-norm blocks, the MLP's own input on post-norm blocks (where ln2
184 # follows the MLP and its input is the raw MLP output).
185 if self.mlp_reads_resid_directly:
186 capture_target = self.submodules.get("mlp") if self.submodules else None
187 else:
188 capture_target = self.submodules.get("ln2") if self.submodules else None
189 if (
190 capture_target is not None
191 and getattr(capture_target, "original_component", None) is not None
192 ):
193 hook_mlp_in = self.hook_mlp_in
194 block_ref = weakref.proxy(self)
196 def _capture_mlp_in(_module: torch.nn.Module, args: tuple) -> Any:
197 if not block_ref._read_use_hook_mlp_in():
198 return None
199 if args and isinstance(args[0], torch.Tensor): 199 ↛ 202line 199 didn't jump to line 202 because the condition on line 199 was always true
200 hooked = hook_mlp_in(args[0])
201 return (hooked,) + args[1:]
202 return None
204 handle = capture_target.register_forward_pre_hook(_capture_mlp_in)
205 self._capture_hook_handles.append(handle)
207 self._capture_hooks_wired = True
209 def _teardown_capture_hooks(self) -> None:
210 """Remove the capture hooks installed by _maybe_wire_capture_hooks (and subclass extensions)."""
211 for handle in self._capture_hook_handles:
212 handle.remove()
213 self._capture_hook_handles.clear()
214 self._capture_hooks_wired = False
216 def _read_use_hook_mlp_in(self) -> bool:
217 """Prefer ``block.config.use_hook_mlp_in``; fall back to the block-local flag."""
218 cfg = self.config
219 if cfg is not None and hasattr(cfg, "use_hook_mlp_in"):
220 return bool(cfg.use_hook_mlp_in)
221 return self._use_hook_mlp_in
223 def _clear_attention_capture(self) -> None:
224 """Release the transient residual captured for attention input forks."""
225 from transformer_lens.model_bridge.generalized_components.attention import (
226 AttentionBridge,
227 )
229 attn = self.submodules.get("attn") if self.submodules else None
230 if isinstance(attn, AttentionBridge):
231 attn._captured_pre_ln_residual = None
233 def forward(self, *args: Any, **kwargs: Any) -> Any:
234 """Forward pass through the block bridge.
236 Args:
237 *args: Input arguments
238 **kwargs: Input keyword arguments
240 Returns:
241 The output from the original component
243 Raises:
244 StopAtLayerException: If stop_at_layer is set and this block should stop execution
245 """
246 if self.original_component is None: 246 ↛ 247line 246 didn't jump to line 247 because the condition on line 246 was never true
247 raise RuntimeError(
248 f"Original component not set for {self.name}. Call set_original_component() first."
249 )
251 self._maybe_wire_capture_hooks()
252 self._clear_attention_capture()
253 args, kwargs = self._maybe_inject_start_residual(args, kwargs)
254 self._check_stop_at_layer(*args, **kwargs)
255 args, kwargs = self._hook_input_hidden_states(args, kwargs)
257 # Filter kwargs to only include parameters accepted by the original component
258 # This prevents errors when passing encoder-specific params to decoder-only models
259 filtered_kwargs = self._filter_kwargs_for_forward(kwargs, len(args))
261 try:
262 output = self.original_component(*args, **filtered_kwargs)
263 finally:
264 self._clear_attention_capture()
265 force_tuple_for_bare_tensor = self._is_standalone_hidden_state_call(args, filtered_kwargs)
266 return self._apply_output_hook(
267 output, force_tuple_for_bare_tensor=force_tuple_for_bare_tensor
268 )
270 def _apply_output_hook(
271 self,
272 output: Any,
273 wrap_single_element: bool = True,
274 force_tuple_for_bare_tensor: bool = False,
275 ) -> Any:
276 """Hook the primary tensor in the output and return the result.
278 Args:
279 output: Raw output from the original component (tensor or tuple).
280 wrap_single_element: If True, single-element tuples stay as tuples after
281 hooking (default, required by most HF models). If False, single-element
282 tuples are unwrapped to a bare tensor (Bloom convention).
283 force_tuple_for_bare_tensor: If True, bare tensor outputs are wrapped into
284 a one-element tuple after hooking. This keeps standalone BlockBridge
285 calls compatible with HF block APIs that expose tuple-like block outputs,
286 while preserving tensor outputs during newer HF parent-model execution.
287 """
288 if isinstance(output, tuple) and len(output) > 0:
289 first = output[0]
290 if isinstance(first, torch.Tensor): 290 ↛ 295line 290 didn't jump to line 295 because the condition on line 290 was always true
291 first = self.hook_out(first)
292 if len(output) == 1:
293 return (first,) if wrap_single_element else first
294 output = (first,) + output[1:]
295 return output
296 if isinstance(output, torch.Tensor):
297 output = self.hook_out(output)
298 if force_tuple_for_bare_tensor and wrap_single_element:
299 return (output,)
300 return output
301 return output
303 @staticmethod
304 def _is_standalone_hidden_state_call(args: tuple, kwargs: dict) -> bool:
305 """Return True for direct block(hidden_states) style calls.
307 Transformers versions differ on whether parent model loops expect block
308 outputs as tuples or tensors. We preserve the original tensor return during
309 full-model execution, but expose tuple-like output for standalone component
310 calls so `output[0]` does not accidentally drop the batch dimension.
311 """
312 if len(args) == 1 and isinstance(args[0], torch.Tensor) and not kwargs:
313 return True
314 return (
315 len(args) == 0
316 and set(kwargs.keys()) == {"hidden_states"}
317 and isinstance(kwargs["hidden_states"], torch.Tensor)
318 )
320 def _extract_layer_idx(self) -> Optional[int]:
321 """Parse this block's layer index from its name (TL/GPT-2/LLaMA patterns)."""
322 if self.name is None: 322 ↛ 323line 322 didn't jump to line 323 because the condition on line 322 was never true
323 return None
324 match = re.search(r"(?:^|\.)(?:blocks|h|layers)\.(\d+)", self.name)
325 return int(match.group(1)) if match else None
327 def _check_stop_at_layer(self, *args: Any, **kwargs: Any) -> None:
328 """Check if execution should stop before this block. Raises StopAtLayerException.
330 The _stop_at_layer_idx attribute is set by the bridge's forward method.
331 Supports TL/GPT-2/LLaMA naming patterns for layer index extraction.
332 """
333 if getattr(self, "_stop_at_layer_idx", None) is None:
334 return
335 layer_idx = self._extract_layer_idx()
336 if layer_idx is not None and layer_idx == self._stop_at_layer_idx:
337 if len(args) > 0 and isinstance(args[0], torch.Tensor): 337 ↛ 339line 337 didn't jump to line 339 because the condition on line 337 was always true
338 input_tensor = args[0]
339 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
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 _maybe_inject_start_residual(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
347 """If this is the start_at_layer block, swap in the caller's residual.
349 Mirror of ``_check_stop_at_layer``: the bridge's forward stashes the
350 residual-stream input on the block via ``_start_residual`` and sets
351 ``_start_at_layer_idx``. This block replaces its incoming hidden states
352 with that residual; ``_hook_input_hidden_states`` then fires ``hook_in``
353 on it, so ``hook_resid_pre`` reflects the injected value.
354 """
355 if getattr(self, "_start_at_layer_idx", None) is None:
356 return args, kwargs
357 if self._extract_layer_idx() != self._start_at_layer_idx:
358 return args, kwargs
359 residual = getattr(self, "_start_residual", None)
360 if residual is None: 360 ↛ 361line 360 didn't jump to line 361 because the condition on line 360 was never true
361 return args, kwargs
362 if len(args) > 0 and isinstance(args[0], torch.Tensor): 362 ↛ 364line 362 didn't jump to line 364 because the condition on line 362 was always true
363 args = (residual,) + args[1:]
364 elif "hidden_states" in kwargs:
365 kwargs = {**kwargs, "hidden_states": residual}
366 return args, kwargs
368 def _hook_input_hidden_states(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
369 """Apply hook_in to the hidden_states input, whether in args or kwargs."""
370 if len(args) > 0 and isinstance(args[0], torch.Tensor): 370 ↛ 373line 370 didn't jump to line 373 because the condition on line 370 was always true
371 hooked_input = self.hook_in(args[0])
372 args = (hooked_input,) + args[1:]
373 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
374 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
375 return args, kwargs
377 def _filter_kwargs_for_forward(
378 self, kwargs: Dict[str, Any], num_positional_args: int = 0
379 ) -> Dict[str, Any]:
380 """Filter kwargs to only include parameters accepted by original_component.forward().
382 This prevents TypeErrors when the bridge passes parameters (like encoder_attention_mask)
383 that aren't accepted by decoder-only models. It also removes any kwargs that would
384 conflict with positional arguments already being passed.
386 Args:
387 kwargs: The full set of keyword arguments
388 num_positional_args: Number of positional arguments being passed (to avoid conflicts)
390 Returns:
391 Filtered kwargs containing only accepted parameters
392 """
393 if self.original_component is None: 393 ↛ 394line 393 didn't jump to line 394 because the condition on line 393 was never true
394 return kwargs
396 try:
397 # Get the signature of the original component's forward method
398 sig = inspect.signature(self.original_component.forward)
399 param_list = list(sig.parameters.keys())
400 valid_params = set(param_list)
402 # Check if the signature accepts **kwargs (VAR_KEYWORD)
403 accepts_var_keyword = any(
404 p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()
405 )
407 # If it accepts **kwargs, pass everything through
408 if accepts_var_keyword:
409 return kwargs
411 # Skip params already provided positionally
412 positional_param_names = set(param_list[:num_positional_args])
414 # Filter kwargs: include only if in signature AND not already provided positionally
415 filtered = {
416 k: v
417 for k, v in kwargs.items()
418 if k in valid_params and k not in positional_param_names
419 }
420 return filtered
422 except (ValueError, TypeError):
423 # If we can't inspect the signature, pass through all kwargs
424 # (better to potentially fail than to silently drop important params)
425 return kwargs
428class MLABlockBridge(BlockBridge):
429 """Block wrapping Multi-Head Latent Attention (DeepSeek V2/V3/R1).
431 MLA has no standalone q/k/v projections — Q flows through compressed
432 q_a_proj→q_a_layernorm→q_b_proj, and K/V share a joint kv_a_proj_with_mqa
433 entry point. There is no single HookPoint that represents "input that
434 becomes Q/K/V", so the block-level ``hook_q_input``/``hook_k_input``/
435 ``hook_v_input``/``hook_attn_in`` aliases do not apply. Type-level
436 distinction means a reader of the adapter sees ``MLABlockBridge`` and
437 knows those hooks are absent.
438 """
440 def __init__(
441 self,
442 name: str,
443 config: Optional[Any] = None,
444 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
445 hook_alias_overrides: Optional[Dict[str, str]] = None,
446 ):
447 super().__init__(
448 name,
449 config=config,
450 submodules=submodules,
451 hook_alias_overrides=hook_alias_overrides,
452 )
453 if self.hook_aliases is BlockBridge.hook_aliases: 453 ↛ 455line 453 didn't jump to line 455 because the condition on line 453 was always true
454 self.hook_aliases = dict(self.hook_aliases)
455 for alias in ("hook_q_input", "hook_k_input", "hook_v_input", "hook_attn_in"):
456 self.hook_aliases.pop(alias, None)
459class ParallelBlockBridge(BlockBridge):
460 """Block where attn and MLP both read the pre-attention residual.
462 For GPT-J, NeoX, Pythia, Phi, Cohere, CodeGen, and some Falcon variants,
463 output = resid_pre + attn_out + mlp_out — no distinct post-attention
464 residual exists. Matches legacy HookedTransformer which omits hook_resid_mid
465 when ``cfg.parallel_attn_mlp=True``. Type-level distinction means a reader
466 of the adapter sees ``ParallelBlockBridge`` and knows the hook is absent.
467 """
469 def __init__(
470 self,
471 name: str,
472 config: Optional[Any] = None,
473 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
474 hook_alias_overrides: Optional[Dict[str, str]] = None,
475 ):
476 super().__init__(
477 name,
478 config=config,
479 submodules=submodules,
480 hook_alias_overrides=hook_alias_overrides,
481 )
482 # Ensure instance-level copy before mutating; base may have left the
483 # class-level dict shared when no overrides were passed.
484 if self.hook_aliases is BlockBridge.hook_aliases: 484 ↛ 486line 484 didn't jump to line 486 because the condition on line 484 was always true
485 self.hook_aliases = dict(self.hook_aliases)
486 self.hook_aliases.pop("hook_resid_mid", None)
489class ScaledResidualBlockBridge(BlockBridge):
490 """Block whose sublayer outputs are scaled before the residual add.
492 Granite-family HF blocks compute ``residual + sublayer_out * residual_multiplier``
493 inline, so no submodule output equals the tensor added to the residual stream
494 and the legacy aliases cannot be fixed by re-pointing. ``hook_attn_out`` /
495 ``hook_mlp_out`` become real HookPoints on the block firing on the scaled
496 contribution:
498 - no hooks attached: the forward is untouched (bit-exact);
499 - read-only hooks: they observe ``module_out * scale``, forward stays bit-exact;
500 - a hook that changes the tensor (returned new or mutated in place): the module
501 output is rewritten to ``hooked / scale`` so HF's multiply reconstructs the
502 written value as the contribution (~1-ulp rounding; exact for zero-ablation);
503 - any backward hooks: the rewrite always happens so the autograd graph routes
504 through the HookPoint.
505 """
507 def __init__(
508 self,
509 name: str,
510 config: Optional[Any] = None,
511 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
512 hook_alias_overrides: Optional[Dict[str, str]] = None,
513 mlp_reads_resid_directly: bool = False,
514 residual_contribution_scale: float = 1.0,
515 scaled_attn_submodule: str = "attn",
516 scaled_mlp_submodule: Optional[str] = "mlp",
517 ):
518 """scaled_mlp_submodule=None when no single submodule feeds the MLP-side
519 add (GraniteMoeHybrid sums moe + shared_mlp inline) — hook_mlp_out then
520 stays absent rather than firing with a partial tensor."""
521 super().__init__(
522 name,
523 config=config,
524 submodules=submodules,
525 hook_alias_overrides=hook_alias_overrides,
526 mlp_reads_resid_directly=mlp_reads_resid_directly,
527 )
528 scale = float(residual_contribution_scale)
529 if scale == 0.0:
530 raise ValueError(
531 f"ScaledResidualBlockBridge at '{name}': residual_contribution_scale "
532 f"must be nonzero (the write path divides by it)."
533 )
534 self.residual_contribution_scale = scale
535 self.scaled_attn_submodule = scaled_attn_submodule
536 self.scaled_mlp_submodule = scaled_mlp_submodule
537 attn = self.submodules.get(scaled_attn_submodule) if self.submodules else None
538 if attn is not None: 538 ↛ 540line 538 didn't jump to line 540 because the condition on line 538 was always true
539 attn._result_contribution_scale = scale
540 if self.hook_aliases is BlockBridge.hook_aliases: 540 ↛ 542line 540 didn't jump to line 542 because the condition on line 540 was always true
541 self.hook_aliases = dict(self.hook_aliases)
542 for alias in ("hook_attn_out", "hook_mlp_out"):
543 self.hook_aliases.pop(alias, None)
544 self.hook_attn_out = HookPoint()
545 if scaled_mlp_submodule is not None:
546 self.hook_mlp_out = HookPoint()
548 def set_original_component(self, original_component: torch.nn.Module) -> None:
549 """Prune contribution HookPoints the bound layer cannot fire.
551 Heterogeneous blocks (GraniteMoeHybrid mamba layers) lack the attn
552 submodule; a HookPoint that exists but never fires is a silent-no-op
553 intervention trap, so the name must be absent instead.
554 """
555 super().set_original_component(original_component)
556 targets = [(self.scaled_attn_submodule, "hook_attn_out")]
557 if self.scaled_mlp_submodule is not None:
558 targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
559 for sub_name, hook_name in targets:
560 sub = self.submodules.get(sub_name) if self.submodules else None
561 remote = getattr(sub, "name", None)
562 first = remote.split(".", 1)[0] if isinstance(remote, str) else None
563 missing = sub is None or (
564 first is not None and getattr(original_component, first, None) is None
565 )
566 if missing:
567 if hasattr(self, hook_name): 567 ↛ 571line 567 didn't jump to line 571 because the condition on line 567 was always true
568 delattr(self, hook_name)
569 # __setattr__ auto-registered the HookPoint; get_hooks() serves
570 # from this registry, so the module deletion alone is not enough.
571 self._hook_registry.pop(hook_name, None)
573 def _maybe_wire_capture_hooks(self) -> None:
574 """Extend the base wiring with the scaled-contribution forward hooks.
576 Shares the base flag and handle list so _teardown_capture_hooks also
577 removes these hooks and re-wiring stays idempotent.
578 """
579 if self._capture_hooks_wired:
580 return
581 super()._maybe_wire_capture_hooks()
582 targets = [(self.scaled_attn_submodule, "hook_attn_out")]
583 if self.scaled_mlp_submodule is not None:
584 targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
585 for sub_name, hook_name in targets:
586 hook_point = getattr(self, hook_name, None)
587 sub = self.submodules.get(sub_name) if self.submodules else None
588 if (
589 hook_point is None
590 or sub is None
591 or getattr(sub, "original_component", None) is None
592 ):
593 continue
594 handle = sub.register_forward_hook(self._make_scaled_contribution_hook(hook_point))
595 self._capture_hook_handles.append(handle)
597 def _make_scaled_contribution_hook(
598 self, hook_point: HookPoint
599 ) -> Callable[[torch.nn.Module, tuple, Any], Any]:
600 """Build a forward hook exposing ``output * scale`` through hook_point."""
601 scale = self.residual_contribution_scale
603 def _hook(_module: torch.nn.Module, _args: tuple, output: Any) -> Any:
604 if not hook_point.has_hooks(dir="both"):
605 return None
606 is_tuple = isinstance(output, tuple)
607 out = output[0] if is_tuple else output
608 if not isinstance(out, torch.Tensor): 608 ↛ 609line 608 didn't jump to line 609 because the condition on line 608 was never true
609 return None
610 new = self._hook_scaled_output(out, hook_point, scale)
611 if new is out:
612 return None
613 return ((new,) + output[1:]) if is_tuple else new
615 return _hook
618class DelegatedAttentionBlockBridge(BlockBridge):
619 """Block whose attention is delegated wholesale to HF (no split-qkv fork).
621 For architectures with heterogeneous per-layer attention structure — e.g.
622 Gemma 4, where KV-shared layers have no ``k_proj``/``v_proj`` at all and
623 K==V layers have no ``v_proj`` — there is no uniform HookPoint that
624 represents "input that becomes Q/K/V", so the block-level ``hook_q_input``/
625 ``hook_k_input``/``hook_v_input``/``hook_attn_in`` aliases do not apply.
626 Type-level distinction means a reader of the adapter sees
627 ``DelegatedAttentionBlockBridge`` and knows those hooks are absent.
628 """
630 # Tell the component benchmark this block's attention is delegated wholesale
631 # to HF and cannot be tested in isolation (requires model-specific kwargs).
632 maintain_native_attention: bool = True
634 def __init__(
635 self,
636 name: str,
637 config: Optional[Any] = None,
638 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
639 hook_alias_overrides: Optional[Dict[str, str]] = None,
640 ):
641 super().__init__(
642 name,
643 config=config,
644 submodules=submodules,
645 hook_alias_overrides=hook_alias_overrides,
646 )
647 if self.hook_aliases is BlockBridge.hook_aliases:
648 self.hook_aliases = dict(self.hook_aliases)
649 for alias in ("hook_q_input", "hook_k_input", "hook_v_input", "hook_attn_in"):
650 self.hook_aliases.pop(alias, None)