Coverage for transformer_lens/model_bridge/bridge_core.py: 85%
925 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Framework-agnostic bridge surface shared by TransformerBridge and RemoteBridge."""
2from __future__ import annotations
4import inspect
5import re
6import warnings
7from contextlib import contextmanager
8from functools import lru_cache, partial
9from typing import (
10 Any,
11 Callable,
12 Dict,
13 FrozenSet,
14 Iterable,
15 Iterator,
16 List,
17 Literal,
18 Mapping,
19 Optional,
20 Tuple,
21 Union,
22 cast,
23 overload,
24)
26import torch
27from torch.nn import functional as F
29from transformer_lens.ActivationCache import ActivationCache
30from transformer_lens.hook_points import HookPoint, LensHandle
31from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
32from transformer_lens.model_bridge.driver_protocol import to_torch
33from transformer_lens.model_bridge.exceptions import StopAtLayerException
34from transformer_lens.model_bridge.generalized_components.base import alias_generation
35from transformer_lens.utilities.aliases import resolve_alias
36from transformer_lens.utilities.lm_utils import lm_cross_entropy_loss
37from transformer_lens.utilities.slice import Slice, SliceInput
39# Aliased: the caching methods take a ``remove_batch_dim`` flag that would shadow it.
40from transformer_lens.utilities.tensors import remove_batch_dim as drop_batch_dim
42_BLOCK_PATTERN = re.compile("blocks\\.(\\d+)")
44# Block-list container attributes a bridge may expose.
45_BLOCK_LIST_ATTRS = ("blocks", "encoder_blocks", "decoder_blocks", "L_blocks", "H_blocks")
46# Encoder blocks name self-attention ``attn``; decoder blocks name it ``self_attn``
47# (``cross_attn`` is a separate submodule, deliberately excluded from stacking).
48_SELF_ATTENTION_NAMES = {"attn": ("attn", "self_attn")}
51def _input_batch_size(*inputs: Any) -> Optional[int]:
52 """The batch size of the first input that states it (a string, list of strings or tensor).
54 A 1-D tensor is one sequence. Returns None when no input says, so the caller can fall
55 back to inferring it from what was cached.
56 """
57 for value in inputs:
58 if isinstance(value, str):
59 return 1
60 if isinstance(value, list) and value and all(isinstance(v, str) for v in value):
61 return len(value)
62 if isinstance(value, torch.Tensor) and value.ndim > 0:
63 return 1 if value.ndim == 1 else int(value.shape[0])
64 return None
67def build_alias_to_canonical_map(hook_dict: Any, prefix: str = "") -> dict:
68 """Map alias hook names to their canonical names (where ``.name`` differs from the key)."""
69 aliases: dict = {}
70 for key, value in hook_dict.items():
71 full_key = f"{prefix}.{key}" if prefix else key
72 if isinstance(value, dict): 72 ↛ 73line 72 didn't jump to line 73 because the condition on line 72 was never true
73 aliases.update(build_alias_to_canonical_map(value, full_key))
74 elif hasattr(value, "name"): 74 ↛ 70line 74 didn't jump to line 70 because the condition on line 74 was always true
75 if key != value.name:
76 aliases[full_key] = value.name
77 return aliases
80class BridgeCore:
81 """Framework-agnostic bridge surface: hooks, cache, run_with_*, driver wiring.
83 Holds state shared by every bridge (adapter, cfg, tokenizer, driver, hook
84 registries). Subclasses add framework-specific state — ``TransformerBridge``
85 walks the wrapped ``nn.Module``; ``RemoteBridge`` builds components from
86 adapter metadata.
87 """
89 hook_aliases: Dict[str, Union[str, List[str]]] = {
90 # Prefer embed_ln.hook_out for post-LN models (Bloom, BERT)
91 "hook_embed": ["embed_ln.hook_out", "embed.hook_out"],
92 "hook_pos_embed": ["pos_embed.hook_out", "rotary_emb.hook_out"],
93 "hook_unembed": "unembed.hook_out",
94 }
96 def __init__(
97 self,
98 adapter: ArchitectureAdapter,
99 tokenizer: Any,
100 driver: Any,
101 ) -> None:
102 """Subclasses call this AFTER ``nn.Module.__init__`` (if applicable),
103 then do framework-specific setup."""
104 self.adapter = adapter
105 self.cfg = adapter.cfg
106 self._tokenizer = None
107 if tokenizer is not None:
108 self.tokenizer = tokenizer # Use the property setter
109 if self.cfg.d_vocab_out == -1:
110 self.cfg.d_vocab_out = self.cfg.d_vocab
111 self.compatibility_mode = False
112 self._weights_processed = False
113 self._hook_cache = None
114 self._hook_registry: Dict[str, HookPoint] = {}
115 self._hook_registry_initialized = False
116 self._hook_alias_registry: Dict[str, Union[str, List[str]]] = {}
117 self._block_alias_cache: Optional[Tuple[Tuple[int, int], Dict[str, str]]] = None
118 self._property_alias_registry: Dict[str, str] = {}
119 self.context_level = 0
120 self._driver = driver
121 if not hasattr(adapter, "component_mapping") or adapter.component_mapping is None: 121 ↛ 122line 121 didn't jump to line 122 because the condition on line 121 was never true
122 raise ValueError("Adapter must have a component_mapping attribute")
124 # ---- tokenizer property ----
126 @property
127 def tokenizer(self) -> Any:
128 """The tokenizer used for encoding/decoding text."""
129 return self._tokenizer
131 @tokenizer.setter
132 def tokenizer(self, value: Any) -> None:
133 """Set tokenizer and re-run wiring (d_vocab, BOS/EOS detection, padding).
135 On initial assignment (during __init__), the boot path has already called
136 configure_tokenizer, so we skip calling it again. However, we still infer
137 d_vocab if it wasn't set from the model config (d_vocab == -1).
139 On reassignment, we re-run configure_tokenizer and update d_vocab to keep
140 cfg in sync with the new tokenizer.
141 """
142 is_reassignment = getattr(self, "_tokenizer", None) is not None
143 cfg = getattr(self, "cfg", None)
144 if value is not None and cfg is not None:
145 if is_reassignment:
146 from transformer_lens.model_bridge.sources._bridge_builder import (
147 configure_tokenizer,
148 )
150 value = configure_tokenizer(value, cfg)
152 # Infer d_vocab: on initial assignment only if not set (-1),
153 # on reassignment always update to match new tokenizer.
154 # Use getattr for cfg attributes since tests may use SimpleNamespace.
155 d_vocab = getattr(cfg, "d_vocab", None)
156 if d_vocab == -1 or is_reassignment:
157 if hasattr(value, "get_vocab"): 157 ↛ 160line 157 didn't jump to line 160 because the condition on line 157 was always true
158 vocab = value.get_vocab()
159 cfg.d_vocab = max(vocab.values()) + 1
160 elif hasattr(value, "vocab"):
161 cfg.d_vocab = max(value.vocab.values()) + 1
162 else:
163 cfg.d_vocab = getattr(value, "vocab_size", 50257)
164 d_vocab_out = getattr(cfg, "d_vocab_out", None)
165 if d_vocab_out == -1 or is_reassignment:
166 cfg.d_vocab_out = getattr(cfg, "d_vocab", d_vocab_out)
167 self._tokenizer = value
169 # ---- hook registry ----
171 def _initialize_hook_registry(self) -> None:
172 """Initialize the hook registry by scanning existing components."""
173 if self._hook_registry_initialized: 173 ↛ 174line 173 didn't jump to line 174 because the condition on line 173 was never true
174 return
175 self._scan_existing_hooks(self, "")
176 self._hook_registry_initialized = True
178 def _scan_existing_hooks(self, module: Any, prefix: str = "") -> None:
179 """Walk components for HookPoint instances. Framework-specific."""
180 raise NotImplementedError(
181 f"{type(self).__name__} must implement _scan_existing_hooks "
182 "(walks the component tree to find HookPoint instances)."
183 )
185 def clear_hook_registry(self) -> None:
186 """Clear the hook registry and force re-initialization."""
187 self._hook_registry.clear()
188 self._hook_registry_initialized = False
190 @property
191 def hook_dict(self) -> dict[str, HookPoint]:
192 """All HookPoint objects, including aliases — TransformerLens-compatible."""
193 hooks = self._hook_registry.copy()
194 self._add_aliases_to_hooks(hooks)
195 return hooks
197 @property
198 def mod_dict(self) -> Dict[str, Any]:
199 """Module/hook name -> object, HookedRootModule-compatible.
201 Union of the named-module tree and the aliased hook view, so both
202 canonical (``blocks.0.mlp.hook_out``) and HT-style
203 (``blocks.0.hook_mlp_out``) names resolve to the same HookPoint.
204 """
205 # BridgeCore is a mixin; concrete bridges (TransformerBridge/RemoteBridge)
206 # are nn.Modules, so named_modules() is always present at runtime.
207 mods: Dict[str, Any] = {
208 name: module
209 for name, module in cast(torch.nn.Module, self).named_modules()
210 if name != ""
211 }
212 mods.update(self.hook_dict)
213 return mods
215 # ---- alias registry ----
217 def _register_aliases(self) -> None:
218 """Register bridge-level aliases (hook_embed, hook_pos_embed, etc.) by
219 resolving each alias target path and installing the target HookPoint
220 as a direct attribute under the alias name."""
221 if self.hook_aliases: 221 ↛ exitline 221 didn't return from function '_register_aliases' because the condition on line 221 was always true
222 self._hook_alias_registry.update(self.hook_aliases)
223 for alias_name, target_path in self.hook_aliases.items():
224 try:
225 if isinstance(target_path, list):
226 for single_target in target_path:
227 try:
228 target_obj = self
229 for part in single_target.split("."):
230 target_obj = getattr(target_obj, part)
231 object.__setattr__(self, alias_name, target_obj)
232 break
233 except AttributeError:
234 continue
235 else:
236 target_obj = self
237 for part in target_path.split("."):
238 target_obj = getattr(target_obj, part)
239 object.__setattr__(self, alias_name, target_obj)
240 except AttributeError:
241 pass
243 def _collect_component_aliases(
244 self, component_mapping: Any, prefix: str = "", _ancestors: FrozenSet[int] = frozenset()
245 ) -> dict:
246 """Recursively collect aliases from the architecture's component templates.
248 ``_ancestors`` is path-scoped, not globally-visited: a cycle is cut
249 (else RecursionError at boot) while a diamond-shared component still
250 contributes aliases under both names.
251 """
252 aliases: dict = {}
253 if id(component_mapping) in _ancestors:
254 return aliases
255 _ancestors = _ancestors | {id(component_mapping)}
256 if isinstance(component_mapping, dict):
257 for name, component in component_mapping.items():
258 sub_prefix = f"{prefix}.{name}" if prefix else name
259 aliases.update(self._collect_component_aliases(component, sub_prefix, _ancestors))
260 else:
261 if hasattr(component_mapping, "hook_aliases") and component_mapping.hook_aliases:
262 for alias_name, target in component_mapping.hook_aliases.items():
263 full_alias = f"{prefix}.{alias_name}" if prefix else alias_name
264 if isinstance(target, list):
265 # Fallback chain: candidates in declared priority order.
266 full_target: Any = tuple(f"{prefix}.{t}" if prefix else t for t in target)
267 else:
268 full_target = f"{prefix}.{target}" if prefix else target
269 aliases[full_alias] = full_target
270 if hasattr(component_mapping, "submodules") and component_mapping.submodules:
271 for sub_name, sub_component in component_mapping.submodules.items():
272 sub_prefix = f"{prefix}.{sub_name}" if prefix else sub_name
273 aliases.update(
274 self._collect_component_aliases(sub_component, sub_prefix, _ancestors)
275 )
276 return aliases
278 @staticmethod
279 @lru_cache(maxsize=128)
280 def _compute_hook_aliases_cached(
281 hook_names_tuple: Tuple[str, ...],
282 component_aliases_tuple: Tuple[Tuple[str, Union[str, Tuple[str, ...]]], ...],
283 ) -> Tuple[Tuple[str, str], ...]:
284 """Cached computation of hook aliases."""
285 aliases: dict = {}
286 # For list-valued (fallback-chain) targets, remember which candidate
287 # resolved each alias so an earlier (higher-priority) candidate wins.
288 alias_priority: dict = {}
289 component_aliases = dict(component_aliases_tuple)
290 for hook_name in hook_names_tuple:
291 for alias_pattern, target_patterns in component_aliases.items():
292 candidates = (
293 target_patterns if isinstance(target_patterns, tuple) else (target_patterns,)
294 )
295 for priority, target_pattern in enumerate(candidates):
296 if "blocks." in target_pattern and "blocks." in hook_name:
297 block_match = _BLOCK_PATTERN.search(hook_name)
298 if not block_match: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 continue
300 block_num = block_match.group(1)
301 dynamic_alias_pattern = alias_pattern.replace(
302 "blocks.", f"blocks.{block_num}."
303 )
304 dynamic_target_pattern = target_pattern.replace(
305 "blocks.", f"blocks.{block_num}."
306 )
307 if hook_name.endswith(dynamic_target_pattern):
308 target_len = len(dynamic_target_pattern)
309 alias_name = hook_name[:-target_len] + dynamic_alias_pattern
310 if alias_priority.get(alias_name, len(candidates)) > priority:
311 aliases[alias_name] = hook_name
312 alias_priority[alias_name] = priority
313 elif hook_name.endswith(target_pattern):
314 target_len = len(target_pattern)
315 alias_name = hook_name[:-target_len] + alias_pattern
316 if alias_priority.get(alias_name, len(candidates)) > priority: 316 ↛ 295line 316 didn't jump to line 295 because the condition on line 316 was always true
317 aliases[alias_name] = hook_name
318 alias_priority[alias_name] = priority
319 return tuple(aliases.items())
321 def _collect_hook_aliases_from_registry(self) -> dict:
322 """Collect aliases based on existing hooks in the registry."""
323 if hasattr(self.adapter, "component_mapping"):
324 component_aliases = self._collect_component_aliases(self.adapter.component_mapping)
325 hook_names_tuple = tuple(sorted(self._hook_registry.keys()))
326 component_aliases_tuple = tuple(sorted(component_aliases.items()))
327 aliases_tuple = self._compute_hook_aliases_cached(
328 hook_names_tuple, component_aliases_tuple
329 )
330 aliases = dict(aliases_tuple)
331 aliases.update(self._collect_block_instance_aliases())
332 return aliases
333 return {}
335 def _collect_block_instance_aliases(self) -> Dict[str, str]:
336 """Collect per-block-instance aliases, overriding template-derived ones.
338 Templates cannot see per-layer rebinds (OlmoHybrid, MoE dense/sparse).
339 Memoized on (registry size, alias generation): size alone misses a
340 dense<->sparse rebind, which changes targets without changing size.
341 """
342 cache_key = (len(self._hook_registry), alias_generation())
343 cached = self._block_alias_cache
344 if cached is not None and cached[0] == cache_key:
345 return cached[1]
346 aliases: Dict[str, str] = {}
347 unresolved: List[str] = []
348 for bl_name in _BLOCK_LIST_ATTRS:
349 block_list = getattr(self, bl_name, None)
350 if block_list is None:
351 continue
352 for i, block in enumerate(block_list):
353 # A block with no registered hooks means the registry hasn't
354 # scanned it yet — unresolved aliases there are timing, not drops.
355 block_prefix = f"{bl_name}.{i}."
356 if f"{block_prefix}hook_in" not in self._hook_registry:
357 continue
358 # Walk the block and its submodule tree: components rebind
359 # aliases per layer at bind time either at block level
360 # (OlmoHybrid) or one level down (MoEBridge's dense/sparse
361 # dispatch). id()-seen guards against shared/cyclic submodule
362 # references, which would otherwise hang boot.
363 # (prefix, component, ids-on-this-path): path-scoped rather than
364 # globally-visited so a cycle is cut while a component shared
365 # under two names still contributes aliases at both.
366 stack: List[Tuple[str, Any, FrozenSet[int]]] = [("", block, frozenset())]
367 while stack:
368 sub_prefix, component, ancestors = stack.pop()
369 if id(component) in ancestors:
370 continue
371 ancestors = ancestors | {id(component)}
372 component_aliases = getattr(component, "hook_aliases", None)
373 if component_aliases:
374 for alias_name, target in component_aliases.items():
375 targets = target if isinstance(target, list) else [target]
376 for single_target in targets:
377 full_target = f"{block_prefix}{sub_prefix}{single_target}"
378 if full_target in self._hook_registry:
379 aliases[f"{block_prefix}{sub_prefix}{alias_name}"] = full_target
380 break
381 else:
382 unresolved.append(f"{block_prefix}{sub_prefix}{alias_name}")
383 for nested_name, nested in (
384 getattr(component, "submodules", None) or {}
385 ).items():
386 stack.append((f"{sub_prefix}{nested_name}.", nested, ancestors))
387 if unresolved:
388 # Surface drops instead of silently swallowing, mirroring
389 # GeneralizedComponent._register_aliases.
390 warnings.warn(
391 f"{len(unresolved)} block hook alias(es) did not resolve to a "
392 f"registered hook (e.g. '{unresolved[0]}'). Any such alias falls "
393 "back to the architecture template's mapping, which for a "
394 "per-layer rebind is the wrong tensor for this layer.",
395 stacklevel=2,
396 )
397 self._block_alias_cache = (cache_key, aliases)
398 return aliases
400 def _add_aliases_to_hooks(self, hooks: Dict[str, HookPoint]) -> None:
401 """Add aliases to hooks in place. Registry-first so RemoteBridge works."""
402 component_aliases = self._collect_hook_aliases_from_registry()
403 all_aliases = {**self.hook_aliases, **component_aliases}
404 if not all_aliases: 404 ↛ 405line 404 didn't jump to line 405 because the condition on line 404 was never true
405 return
406 for alias_name, target in all_aliases.items():
407 targets = target if isinstance(target, list) else [target]
408 for t in targets:
409 hp = self._hook_registry.get(t)
410 if hp is not None:
411 hooks[alias_name] = hp
412 break
413 # Fall back to attribute walk for TransformerBridge nested paths
414 # not directly keyed in the registry.
415 try:
416 target_hook = resolve_alias(self, alias_name, {alias_name: t})
417 if target_hook is not None: 417 ↛ 408line 417 didn't jump to line 408 because the condition on line 417 was always true
418 hooks[alias_name] = target_hook
419 break
420 except AttributeError:
421 continue
423 # ---- captures from driver ----
425 def _replay_captures(self, captured: Mapping[str, Any]) -> None:
426 """Fire driver-delivered captures through the registry. Unknown names dropped silently."""
427 for hook_name, activation in captured.items():
428 hp = self._hook_registry.get(hook_name)
429 if hp is None:
430 continue
431 hp(to_torch(activation))
433 # ---- driver dispatch (subclasses concrete-override) ----
435 def forward(self, *args: Any, **kwargs: Any) -> Any:
436 """Subclasses implement how the driver gets called."""
437 raise NotImplementedError(f"{type(self).__name__} must implement forward()")
439 def to_tokens(self, *args: Any, **kwargs: Any) -> Any:
440 """Subclasses implement against their tokenizer surface."""
441 raise NotImplementedError(f"{type(self).__name__} must implement to_tokens()")
443 def close(self) -> None:
444 """Release driver-managed resources. Idempotent — safe to call multiple times."""
445 self._driver.close()
447 def _input_device(self) -> Any:
448 """Driver's expected device for inputs; None for remote / meta / dispatched drivers."""
449 if not self._driver.supports("parameters"):
450 return None
451 params = getattr(self._driver, "parameters", None)
452 if not callable(params): 452 ↛ 453line 452 didn't jump to line 453 because the condition on line 452 was never true
453 return None
454 try:
455 device = next(params()).device
456 except (StopIteration, NotImplementedError, RuntimeError):
457 return None
458 # Meta device → inputs would silently become meta tensors with no data.
459 if device.type == "meta":
460 return None
461 return device
463 def loss_fn(
464 self,
465 logits: torch.Tensor,
466 tokens: torch.Tensor,
467 attention_mask: Optional[torch.Tensor] = None,
468 per_token: bool = False,
469 ) -> torch.Tensor:
470 """Cross-entropy loss matching HookedTransformer's formula (log_softmax + gather)."""
471 if tokens.device != logits.device: 471 ↛ 472line 471 didn't jump to line 472 because the condition on line 471 was never true
472 tokens = tokens.to(logits.device)
473 if attention_mask is not None:
474 if attention_mask.device != logits.device: 474 ↛ 475line 474 didn't jump to line 475 because the condition on line 474 was never true
475 attention_mask = attention_mask.to(logits.device)
476 attention_mask = self._prepare_loss_attention_mask(attention_mask, tokens)
477 return lm_cross_entropy_loss(logits, tokens, attention_mask, per_token)
479 def _causal_labels_loss(
480 self,
481 logits: torch.Tensor,
482 labels: torch.Tensor,
483 attention_mask: Optional[torch.Tensor] = None,
484 per_token: bool = False,
485 ) -> torch.Tensor:
486 """Compute shifted causal loss against explicit labels, ignoring ``-100``."""
487 if labels.device != logits.device: 487 ↛ 488line 487 didn't jump to line 488 because the condition on line 487 was never true
488 labels = labels.to(logits.device)
489 if labels.shape != logits.shape[:-1]: 489 ↛ 490line 489 didn't jump to line 490 because the condition on line 489 was never true
490 raise ValueError(
491 "causal labels must match the logits batch and position dimensions, "
492 f"got labels {tuple(labels.shape)} and logits {tuple(logits.shape)}"
493 )
495 losses = F.cross_entropy(
496 logits[:, :-1].flatten(0, 1),
497 labels[:, 1:].flatten(),
498 reduction="none",
499 ignore_index=-100,
500 ).view_as(labels[:, 1:])
501 valid_targets = labels[:, 1:] != -100
502 if attention_mask is not None:
503 if attention_mask.device != logits.device: 503 ↛ 504line 503 didn't jump to line 504 because the condition on line 503 was never true
504 attention_mask = attention_mask.to(logits.device)
505 token_mask = self._prepare_loss_attention_mask(attention_mask, labels)
506 valid_targets &= token_mask[:, :-1] & token_mask[:, 1:]
507 losses = losses.masked_fill(~valid_targets, 0.0)
508 return losses if per_token else losses.sum() / valid_targets.sum()
510 def _prepare_loss_attention_mask(
511 self, attention_mask: torch.Tensor, tokens: torch.Tensor
512 ) -> torch.Tensor:
513 """Reduce a forward mask using the adapter's 4D boolean polarity."""
514 batch, pos = tokens.shape
515 if attention_mask.ndim not in (2, 4): 515 ↛ 516line 515 didn't jump to line 516 because the condition on line 515 was never true
516 raise ValueError(
517 "attention_mask must be 2D [batch, key_pos] or 4D "
518 f"[batch, *, query_pos, key_pos], got shape {tuple(attention_mask.shape)}"
519 )
520 if attention_mask.shape[0] != batch: 520 ↛ 521line 520 didn't jump to line 521 because the condition on line 520 was never true
521 raise ValueError(
522 "attention_mask batch dimension must match tokens, "
523 f"got {attention_mask.shape[0]} and {batch}"
524 )
526 if attention_mask.ndim == 2:
527 if attention_mask.shape[1] < pos: 527 ↛ 528line 527 didn't jump to line 528 because the condition on line 527 was never true
528 raise ValueError(
529 "attention_mask must cover every scored token, "
530 f"got length {attention_mask.shape[1]} for {pos} tokens"
531 )
532 return attention_mask[:, -pos:].bool()
534 query_pos, key_pos = attention_mask.shape[-2:]
535 if key_pos < pos: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 raise ValueError(
537 "attention_mask must cover every scored token, "
538 f"got key length {key_pos} for {pos} tokens"
539 )
541 if attention_mask.dtype is torch.bool:
542 blocked = ~attention_mask if self.adapter.bool_4d_mask_is_keep else attention_mask
543 else:
544 blocked = attention_mask < -1.0
545 if query_pos == 1:
546 # Broadcast key-only masks use one query row for the full sequence.
547 keep = ~blocked[..., 0, -pos:]
548 else:
549 if query_pos < pos: 549 ↛ 550line 549 didn't jump to line 550 because the condition on line 549 was never true
550 raise ValueError(
551 "attention_mask must contain a query row for every scored token, "
552 f"got {query_pos} rows for {pos} tokens"
553 )
554 # The aligned diagonal excludes causal masking while retaining padding.
555 diagonal = torch.diagonal(
556 blocked,
557 offset=key_pos - query_pos,
558 dim1=-2,
559 dim2=-1,
560 )
561 if diagonal.shape[-1] < pos: 561 ↛ 562line 561 didn't jump to line 562 because the condition on line 561 was never true
562 raise ValueError(
563 "attention_mask diagonal must cover every scored token, "
564 f"got length {diagonal.shape[-1]} for {pos} tokens"
565 )
566 keep = ~diagonal[..., -pos:]
568 # A token is padding only when every broadcast/head mask blocks its key.
569 return keep.reshape(batch, -1, pos).any(dim=1)
571 def _finalize_return(
572 self,
573 return_type: Optional[str],
574 logits: Optional[torch.Tensor],
575 input_ids: Optional[torch.Tensor],
576 *,
577 attention_mask: Optional[torch.Tensor] = None,
578 labels: Optional[torch.Tensor] = None,
579 is_audio_model: bool = False,
580 is_visual_model: bool = False,
581 inputs_embeds_was_used: bool = False,
582 loss_per_token: bool = False,
583 ) -> Any:
584 """Post-process driver output into the user's requested return_type."""
585 if return_type == "logits":
586 return logits
587 if return_type is None:
588 return None
589 self._check_loss_supported(return_type)
590 if return_type == "loss":
591 if is_audio_model: 591 ↛ 592line 591 didn't jump to line 592 because the condition on line 591 was never true
592 raise ValueError(
593 "Audio models do not support return_type='loss'. "
594 "CTC loss requires aligned frame-level labels."
595 )
596 if is_visual_model: 596 ↛ 597line 596 didn't jump to line 597 because the condition on line 596 was never true
597 raise ValueError(
598 "Vision classification models do not support return_type='loss' "
599 "via this path (no next-token LM target exists for image "
600 "classification). Compute cross-entropy against `labels` "
601 "yourself from the returned logits, or use hf_generate()-style "
602 "direct access to self.original_model for HF's own loss."
603 )
604 if inputs_embeds_was_used and labels is None: 604 ↛ 605line 604 didn't jump to line 605 because the condition on line 604 was never true
605 raise ValueError(
606 "Cannot compute loss with inputs_embeds — token IDs required for labels."
607 )
608 assert isinstance(logits, torch.Tensor), f"Expected logits tensor, got {type(logits)}"
609 if labels is not None:
610 return self._causal_labels_loss(
611 logits,
612 labels,
613 attention_mask=attention_mask,
614 per_token=loss_per_token,
615 )
616 assert input_ids is not None, "input_ids required for return_type='loss'"
617 return self.loss_fn(
618 logits,
619 input_ids,
620 attention_mask=attention_mask,
621 per_token=loss_per_token,
622 )
623 if return_type == "both":
624 if is_audio_model: 624 ↛ 625line 624 didn't jump to line 625 because the condition on line 624 was never true
625 raise ValueError(
626 "Audio models do not support return_type='both'. "
627 "CTC loss requires aligned frame-level labels."
628 )
629 if is_visual_model: 629 ↛ 630line 629 didn't jump to line 630 because the condition on line 629 was never true
630 raise ValueError(
631 "Vision classification models do not support return_type='both' "
632 "via this path (no next-token LM target exists for image "
633 "classification). Compute cross-entropy against `labels` "
634 "yourself from the returned logits."
635 )
636 if inputs_embeds_was_used and labels is None: 636 ↛ 637line 636 didn't jump to line 637 because the condition on line 636 was never true
637 raise ValueError(
638 "Cannot compute loss with inputs_embeds — token IDs required for labels."
639 )
640 assert isinstance(logits, torch.Tensor), f"Expected logits tensor, got {type(logits)}"
641 if labels is not None:
642 loss = self._causal_labels_loss(
643 logits,
644 labels,
645 attention_mask=attention_mask,
646 per_token=loss_per_token,
647 )
648 else:
649 assert input_ids is not None, "input_ids required for return_type='both'"
650 loss = self.loss_fn(
651 logits,
652 input_ids,
653 attention_mask=attention_mask,
654 per_token=loss_per_token,
655 )
656 return (logits, loss)
657 if return_type == "predictions":
658 assert self.tokenizer is not None, "Tokenizer required for return_type='predictions'"
659 assert isinstance(logits, torch.Tensor), f"Expected logits tensor, got {type(logits)}"
660 if logits.shape[-1] == 2: 660 ↛ 670line 660 didn't jump to line 670 because the condition on line 660 was always true
661 # Next Sentence Prediction — 2-class output
662 logprobs = logits.log_softmax(dim=-1)
663 predictions = [
664 "The sentences are sequential",
665 "The sentences are NOT sequential",
666 ]
667 return predictions[int(logprobs.argmax(dim=-1).item())]
668 else:
669 # Masked Language Modeling — decode [MASK] tokens
670 assert input_ids is not None, "input_ids required for MLM predictions"
671 logprobs = logits[input_ids == self.tokenizer.mask_token_id].log_softmax(dim=-1)
672 preds = self.tokenizer.decode(logprobs.argmax(dim=-1))
673 if " " in preds:
674 parts = preds.split(" ")
675 return [f"Prediction {i}: {p}" for i, p in enumerate(parts)]
676 return preds
677 raise ValueError(f"Invalid return_type: {return_type}")
679 def _check_loss_supported(self, return_type: Optional[str]) -> None:
680 """Loss needs full-sequence logits; final-position-only drivers would NaN."""
681 if return_type in ("loss", "both") and not getattr(
682 self._driver, "provides_sequence_logits", True
683 ):
684 raise NotImplementedError(
685 f"return_type={return_type!r} is unsupported on this driver: it "
686 "provides next-token logits for the final position only, so loss "
687 "over earlier positions is undefined. Use return_type='logits' "
688 "and read logits[..., -1, :]."
689 )
691 # ---- hook lookup / mutation ----
693 @staticmethod
694 def _is_embedding_stage_hook(name: str) -> bool:
695 """Hooks belonging to the pre-block token/positional embedding stage.
697 Excluded from ``start_at_layer`` output: the caller's residual already
698 carries the embedding, so the embedding stage is logically skipped even
699 though HF still computes (and discards) it. ``unembed``/``hook_unembed``
700 are the output stage and deliberately not matched.
701 """
702 return name in ("hook_embed", "hook_pos_embed", "hook_tokens") or name.startswith(
703 ("embed.", "pos_embed.")
704 )
706 def _check_hook_fireable(self, *names: str) -> None:
707 """Fail loud when the driver declares it can't fire a requested hook —
708 attaching anyway would yield a silently-unhooked forward / empty cache."""
709 non_fireable: frozenset = getattr(self._driver, "non_fireable_hook_points", frozenset())
710 for name in names:
711 if name in non_fireable:
712 raise NotImplementedError(
713 f"this backend cannot fire {name!r}; use boot_transformers() "
714 "for full hook coverage."
715 )
717 def _resolve_hook_point(
718 self, name: str, aliases: Dict[str, str], hook_dict: Dict[str, HookPoint]
719 ) -> Tuple[str, HookPoint]:
720 """Resolve an (aliased) string hook name to its HookPoint, enforcing
721 fireability. A name that resolves to nothing raises — a typo'd name must
722 not run unhooked."""
723 canonical = aliases.get(name, name)
724 self._check_hook_fireable(name, canonical)
725 hook_point = hook_dict.get(canonical)
726 if hook_point is None:
727 raise KeyError(f"Hook name {name!r} does not exist on this model.")
728 return canonical, hook_point
730 def get_hook_point(self, hook_name: str) -> Optional[HookPoint]:
731 """Get a hook point by name from the bridge's hook system."""
732 if hook_name in self._hook_registry: 732 ↛ 734line 732 didn't jump to line 734 because the condition on line 732 was always true
733 return self._hook_registry[hook_name]
734 try:
735 parts = hook_name.split(".")
736 current: Any = self
737 for part in parts:
738 current = getattr(current, part)
739 if isinstance(current, HookPoint):
740 return current
741 except AttributeError:
742 pass
743 return None
745 def check_hooks_to_add(
746 self,
747 hook_point: HookPoint,
748 hook_point_name: str,
749 hook: Callable,
750 dir: Literal["fwd", "bwd"] = "fwd",
751 is_permanent: bool = False,
752 prepend: bool = False,
753 ) -> None:
754 """Validate a hook before it is added; override to add checks.
756 Raises for a gated-off hook point — a targeted attach there would
757 silently never fire. Every explicit-name attach path routes through
758 here; filter sweeps pre-skip gated matches (with a warning) before
759 reaching it, since a filter was not necessarily targeting them.
761 Gating keys on the POINT's own canonical name, not the requested
762 spelling: on adapters with ``hook_alias_overrides`` a gated HT name
763 (e.g. ``blocks.0.hook_mlp_in`` on BERT) resolves to an always-firing
764 point, and refusing it would reject a hook that works.
765 Driver-fireability is enforced separately in ``_check_hook_fireable``.
766 """
767 reason = self._gated_hook_reason(hook_point.name or hook_point_name)
768 if reason is not None:
769 raise ValueError(
770 f"Cannot add hook {hook_point_name} because {reason} is False. "
771 f"Call set_{reason}(True) first."
772 )
774 def _add_fn_to_hook_point(
775 self,
776 hook_point: HookPoint,
777 name: str,
778 hook_fn: Callable,
779 dir: Literal["fwd", "bwd"],
780 is_permanent: bool,
781 ) -> None:
782 """Run the extension-point check, then attach the hook function."""
783 self.check_hooks_to_add(hook_point, name, hook_fn, dir=dir, is_permanent=is_permanent)
784 hook_point.add_hook(hook_fn, dir=dir, is_permanent=is_permanent)
786 def _gated_hook_reason(self, hook_point_name: str) -> Optional[str]:
787 """Return the disabled setter name if hook_point_name is gated off, else None."""
788 if hook_point_name.endswith("attn.hook_result") and not self.cfg.use_attn_result:
789 return "use_attn_result"
790 if (
791 hook_point_name.endswith(("hook_q_input", "hook_k_input", "hook_v_input"))
792 and not self.cfg.use_split_qkv_input
793 ):
794 return "use_split_qkv_input"
795 if hook_point_name.endswith("mlp_in") and not self.cfg.use_hook_mlp_in:
796 return "use_hook_mlp_in"
797 if hook_point_name.endswith("attn_in") and not self.cfg.use_attn_in:
798 return "use_attn_in"
799 return None
801 def _warn_gated_skipped(self, api: str, names: List[str], stacklevel: int = 3) -> None:
802 """One warning per sweep for filter-matched gated-off names."""
803 if names:
804 warnings.warn(
805 f"{api}: skipped {len(names)} gated-off hook name(s) "
806 f"that would never fire: {names}. Call the relevant "
807 "set_use_*(True) setter first to enable them.",
808 stacklevel=stacklevel,
809 )
811 def add_hook(
812 self,
813 name: Union[str, Callable[[str], bool]],
814 hook_fn: Any,
815 dir: Literal["fwd", "bwd"] = "fwd",
816 is_permanent: bool = False,
817 ) -> None:
818 """Add a hook to a specific component or to all components matching a filter.
820 Args:
821 name: Either a string hook point name (e.g. "blocks.0.attn.hook_q")
822 or a callable filter ``(str) -> bool`` that is applied to every
823 hook point name; the hook is added to each point where the filter
824 returns True.
825 hook_fn: The hook function ``(activation, hook) -> activation | None``.
826 dir: Hook direction, ``"fwd"`` or ``"bwd"``.
827 is_permanent: If True the hook survives ``reset_hooks()`` calls.
828 """
829 if callable(name) and not isinstance(name, str):
830 hook_dict = self.hook_dict
831 seen_hooks: set = set()
832 gated_names_skipped: List[str] = []
833 for hook_name, hook_point in hook_dict.items():
834 if name(hook_name):
835 hook_id = id(hook_point)
836 if hook_id in seen_hooks: 836 ↛ 837line 836 didn't jump to line 837 because the condition on line 836 was never true
837 continue
838 seen_hooks.add(hook_id)
839 # A filter is a sweep, not a targeted request, so a gated-off
840 # match is skipped rather than raised on — but silently
841 # attaching here would leave a dead hook that never fires,
842 # which is the failure this warns about.
843 if self._gated_hook_reason(hook_point.name or hook_name) is not None:
844 gated_names_skipped.append(hook_name)
845 continue
846 self._add_fn_to_hook_point(hook_point, hook_name, hook_fn, dir, is_permanent)
847 self._warn_gated_skipped("add_hook", gated_names_skipped)
848 return
850 # Explicit-name attaches raise on gated-off points inside
851 # check_hooks_to_add (via _add_fn_to_hook_point) — every path below
852 # funnels through it.
853 # Fast path: canonical registry names skip the alias-map build (hook_dict +
854 # map construction cost ~ms on large models; add_hook is often called per layer).
855 registry_hp = self._hook_registry.get(name)
856 if registry_hp is not None:
857 self._check_hook_fireable(name)
858 self._add_fn_to_hook_point(registry_hp, name, hook_fn, dir, is_permanent)
859 return
860 # Same alias resolution run_with_hooks uses, so HT-style names work here too.
861 canonical = build_alias_to_canonical_map(self.hook_dict).get(name, name)
862 if canonical != name:
863 self._check_hook_fireable(name, canonical)
864 else:
865 self._check_hook_fireable(name)
866 registry_hp = self._hook_registry.get(canonical)
867 if registry_hp is not None: 867 ↛ 871line 867 didn't jump to line 871 because the condition on line 867 was always true
868 self._add_fn_to_hook_point(registry_hp, canonical, hook_fn, dir, is_permanent)
869 return
871 component: Any = self
872 parts = name.split(".")
873 for part in parts[:-1]:
874 if hasattr(component, part):
875 component = getattr(component, part)
876 else:
877 raise AttributeError(f"Component path '{'.'.join(parts[:-1])}' not found")
878 hook_name = parts[-1]
879 if hasattr(component, hook_name):
880 hook_point = getattr(component, hook_name)
881 if isinstance(hook_point, HookPoint):
882 self._add_fn_to_hook_point(hook_point, name, hook_fn, dir, is_permanent)
883 else:
884 raise AttributeError(
885 f"'{hook_name}' is not a hook point. Found object of type: {type(hook_point)} with value: {hook_point}"
886 )
887 else:
888 raise AttributeError(f"Hook point '{hook_name}' not found on component")
890 def add_perma_hook(
891 self,
892 name: Union[str, Callable[[str], bool]],
893 hook_fn: Callable,
894 dir: Literal["fwd", "bwd"] = "fwd",
895 ) -> None:
896 """Add a permanent hook that survives ``reset_hooks()`` calls.
898 Convenience wrapper for ``add_hook(..., is_permanent=True)``. To remove,
899 call ``reset_hooks(including_permanent=True)`` or remove from the
900 underlying ``HookPoint`` directly.
901 """
902 self.add_hook(name, hook_fn, dir=dir, is_permanent=True)
904 def hook_points(self) -> Iterable[HookPoint]:
905 """All :class:`HookPoint` instances (registry is canonical and complete)."""
906 return self._hook_registry.values()
908 def clear_contexts(self) -> None:
909 """Clear the stored ``ctx`` on every hook point."""
910 for hp in self._hook_registry.values():
911 hp.clear_context()
913 def remove_all_hook_fns(
914 self,
915 direction: Literal["fwd", "bwd", "both"] = "both",
916 including_permanent: bool = False,
917 level: Optional[int] = None,
918 ) -> None:
919 """Remove hook functions from every hook point."""
920 for hp in self._hook_registry.values():
921 hp.remove_hooks(direction, including_permanent=including_permanent, level=level)
923 def reset_hooks(
924 self,
925 clear_contexts: bool = True,
926 direction: Literal["fwd", "bwd", "both"] = "both",
927 including_permanent: bool = False,
928 level: Optional[int] = None,
929 ) -> None:
930 """Remove hooks from every hook point; mirrors ``HookedRootModule.reset_hooks``.
932 The hook registry is canonical and complete (every component's HookPoint
933 is registered), so a single pass covers the whole model.
935 Args:
936 clear_contexts: Also clear each hook point's stored ``ctx``.
937 direction: Which direction(s) to remove — ``"fwd"``, ``"bwd"``, or ``"both"``.
938 including_permanent: If True, also remove hooks added via ``add_perma_hook``.
939 level: If set, only remove hooks registered at this context level.
940 """
941 if clear_contexts:
942 self.clear_contexts()
943 self.remove_all_hook_fns(direction, including_permanent=including_permanent, level=level)
945 @staticmethod
946 def _normalize_names_filter(
947 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]],
948 ) -> Callable[[str], bool]:
949 """Turn None / str / list / callable into a name predicate."""
950 if names_filter is None:
951 return lambda name: True
952 if isinstance(names_filter, str):
953 return lambda name: name == names_filter
954 if isinstance(names_filter, list):
955 return lambda name: name in names_filter
956 if callable(names_filter): 956 ↛ 958line 956 didn't jump to line 958 because the condition on line 956 was always true
957 return names_filter
958 raise ValueError("names_filter must be None, a string, a list of strings, or a callable")
960 @staticmethod
961 def _pos_slice_dim(name: str, hook_point: Any = None) -> int:
962 """Position axis of a hook's activation in its batched layout.
964 Position is dim 1 for most bridge activations (resid ``[b, p, d]``, per-head
965 ``[b, p, h, d]``, token ids ``[b, p]``) and -2 for attention patterns/scores
966 (``[b, h, q_pos, k_pos]``). Hooks fired on another layout, such as a
967 post-reshape qk-norm (``[b, h, p, d]``) or a Mamba conv input
968 (``[b, channels, p]``), declare their axis via ``HookPoint.pos_dim``.
969 """
970 pos_dim = getattr(hook_point, "pos_dim", None)
971 if pos_dim is not None:
972 return pos_dim
973 return -2 if name.endswith(("hook_pattern", "hook_attn_scores")) else 1
975 def get_caching_hooks(
976 self,
977 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
978 incl_bwd: bool = False,
979 device: Any = None,
980 remove_batch_dim: bool = False,
981 cache: Optional[dict] = None,
982 pos_slice: Optional[Union[Slice, SliceInput]] = None,
983 ) -> Tuple[dict, list, list]:
984 """Build caching hooks without adding them. Mirrors ``HookedRootModule.get_caching_hooks``.
986 ``remove_batch_dim`` drops a leading dimension of size 1 and leaves any other tensor as
987 it is, so a batch larger than 1 keeps its batch dimension instead of losing examples.
988 The check is per tensor, so a broadcast activation with a size-1 leading dimension
989 (e.g. rotary ``hook_cos``/``hook_sin``) is squeezed even when the batch is larger.
990 Unlike ``run_with_cache``, which raises for a batch larger than 1, a hook sees one tensor
991 at a time and cannot tell a batch dimension from a flattened ``[batch * pos, ...]`` one.
992 Hooks that run under ``run_with_hooks(remove_batch_dim=True)`` already receive tensors
993 without a batch dimension, so set the flag on only one of the two.
995 Returns ``(cache, fwd_hooks, bwd_hooks)`` where each hook is a
996 ``(name, hook_fn)`` pair suitable for ``hooks()`` / ``run_with_hooks``.
997 Activations are keyed by the HookPoint's canonical name, and also by any
998 alias spelling that ``names_filter`` matched. This is broader than
999 ``run_with_cache`` for callable filters: one that matches only an alias
1000 yields both the canonical and alias keys here, while ``run_with_cache``
1001 stores only the alias. Backward hooks append ``"_grad"``. ``bwd_hooks``
1002 is empty unless ``incl_bwd``.
1003 """
1004 if cache is None:
1005 cache = {}
1006 pos_slice_obj = Slice.unwrap(pos_slice)
1007 filter_fn = self._normalize_names_filter(names_filter)
1009 def save_hook(
1010 tensor: Any, hook: Any, is_backward: bool = False, alias_names: Tuple[str, ...] = ()
1011 ) -> None:
1012 assert hook.name is not None
1013 suffix = "_grad" if is_backward else ""
1014 stored = tensor.detach().to(device)
1015 # Slice before dropping the batch dim: _pos_slice_dim indexes the
1016 # batched layout, so the order must match run_with_cache.
1017 if pos_slice_obj is not None and stored.dim() >= 2: 1017 ↛ 1019line 1017 didn't jump to line 1019 because the condition on line 1017 was always true
1018 stored = pos_slice_obj.apply(stored, dim=self._pos_slice_dim(hook.name, hook))
1019 if remove_batch_dim:
1020 stored = drop_batch_dim(stored)
1021 for key in (hook.name, *alias_names):
1022 cache[key + suffix] = stored
1024 # One hook point can be reachable under several names (canonical + aliases).
1025 # Gather every name the filter matched so the cache has an entry under the
1026 # spelling the caller asked for, not just the canonical one.
1027 matched: Dict[int, Tuple[str, HookPoint, List[str]]] = {}
1028 for name, hook_point in self.hook_dict.items():
1029 if not filter_fn(name):
1030 continue
1031 entry = matched.setdefault(id(hook_point), (name, hook_point, []))
1032 if name != hook_point.name:
1033 entry[2].append(name)
1035 fwd_hooks: list = []
1036 bwd_hooks: list = []
1037 for name, hook_point, alias_names in matched.values():
1038 # The default sweep must not emit gated-off points: they never
1039 # fire (empty cache entries at best), and feeding them to
1040 # hooks()/run_with_hooks — the advertised composition — would
1041 # raise. An explicit filter keeps them so downstream can warn.
1042 if names_filter is None and self._gated_hook_reason(hook_point.name or name):
1043 continue
1044 # The unfiltered sweep matches every alias too; keep it canonical-only.
1045 extra_names = () if names_filter is None else tuple(alias_names)
1046 fwd_hooks.append((name, partial(save_hook, is_backward=False, alias_names=extra_names)))
1047 if incl_bwd:
1048 bwd_hooks.append(
1049 (name, partial(save_hook, is_backward=True, alias_names=extra_names))
1050 )
1051 return cache, fwd_hooks, bwd_hooks
1053 def add_caching_hooks(
1054 self,
1055 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
1056 incl_bwd: bool = False,
1057 device: Any = None,
1058 remove_batch_dim: bool = False,
1059 cache: Optional[dict] = None,
1060 ) -> dict:
1061 """Attach caching hooks to the model (does not run it). Returns the cache dict.
1063 Mirrors ``HookedRootModule.add_caching_hooks``. The hooks persist until
1064 ``reset_hooks()``. ``remove_batch_dim`` behaves as in :meth:`get_caching_hooks`: a leading
1065 dimension of size 1 is dropped and a larger batch keeps its batch dimension.
1066 """
1067 cache, fwd_hooks, bwd_hooks = self.get_caching_hooks(
1068 names_filter,
1069 incl_bwd=incl_bwd,
1070 device=device,
1071 remove_batch_dim=remove_batch_dim,
1072 cache=cache,
1073 )
1075 # Gated-off hook points never fire, so caching them is a no-op; skip
1076 # them rather than trip add_hook's explicit-name guard. Warn only when
1077 # the caller EXPLICITLY asked for them (run_with_cache's rule: the
1078 # default filter matches everything and must not read as a request).
1079 def _point_reason(name: str):
1080 hp = self.hook_dict.get(name)
1081 return self._gated_hook_reason(hp.name or name if hp is not None else name)
1083 gated_skipped = []
1084 for name, hook_fn in fwd_hooks:
1085 if _point_reason(name) is None:
1086 self.add_hook(name, hook_fn, dir="fwd")
1087 else:
1088 gated_skipped.append(name)
1089 for name, hook_fn in bwd_hooks: 1089 ↛ 1090line 1089 didn't jump to line 1090 because the loop on line 1089 never started
1090 if _point_reason(name) is None:
1091 self.add_hook(name, hook_fn, dir="bwd")
1092 elif name not in gated_skipped:
1093 gated_skipped.append(name)
1094 if names_filter is not None:
1095 self._warn_gated_skipped("add_caching_hooks", gated_skipped)
1096 return cache
1098 def cache_all(
1099 self,
1100 cache: Optional[dict],
1101 incl_bwd: bool = False,
1102 device: Any = None,
1103 remove_batch_dim: bool = False,
1104 ) -> None:
1105 """Deprecated: cache every activation. Use ``run_with_cache`` / ``add_caching_hooks``."""
1106 warnings.warn(
1107 "cache_all is deprecated; use run_with_cache or add_caching_hooks.",
1108 DeprecationWarning,
1109 stacklevel=2,
1110 )
1111 self.add_caching_hooks(
1112 names_filter=None,
1113 cache=cache,
1114 incl_bwd=incl_bwd,
1115 device=device,
1116 remove_batch_dim=remove_batch_dim,
1117 )
1119 def cache_some(
1120 self,
1121 cache: Optional[dict],
1122 names: Callable[[str], bool],
1123 incl_bwd: bool = False,
1124 device: Any = None,
1125 remove_batch_dim: bool = False,
1126 ) -> None:
1127 """Deprecated: cache activations matching ``names``. Use ``run_with_cache``."""
1128 warnings.warn(
1129 "cache_some is deprecated; use run_with_cache or add_caching_hooks.",
1130 DeprecationWarning,
1131 stacklevel=2,
1132 )
1133 self.add_caching_hooks(
1134 names_filter=names,
1135 cache=cache,
1136 incl_bwd=incl_bwd,
1137 device=device,
1138 remove_batch_dim=remove_batch_dim,
1139 )
1141 @staticmethod
1142 def _remove_temporary_hook(
1143 hook_point: HookPoint, direction: Literal["fwd", "bwd"], handle: LensHandle
1144 ) -> None:
1145 """Remove an invocation's handle, not retained hooks at a reused lexical level."""
1146 if handle.is_permanent: 1146 ↛ 1147line 1146 didn't jump to line 1147 because the condition on line 1146 was never true
1147 return
1148 handle.hook.remove()
1149 handles = hook_point.fwd_hooks if direction == "fwd" else hook_point.bwd_hooks
1150 handles[:] = [existing for existing in handles if existing is not handle]
1152 def hooks(
1153 self,
1154 fwd_hooks: List = [],
1155 bwd_hooks: List = [],
1156 reset_hooks_end: bool = True,
1157 clear_contexts: bool = False,
1158 ) -> Any:
1159 """Context manager for temporarily adding hooks.
1161 ``reset_hooks_end`` removes the hooks this context added when it exits —
1162 hooks the caller attached beforehand are left alone either way;
1163 ``clear_contexts`` also wipes the touched hook points' ``ctx``.
1165 Example:
1166 with model.hooks(fwd_hooks=[("hook_embed", my_hook)]):
1167 output = model("Hello world")
1168 """
1170 @contextmanager
1171 def _hooks_context() -> Iterator["BridgeCore"]:
1172 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = []
1173 context_level = getattr(self, "context_level", 0) + 1
1174 self.context_level = context_level
1176 def add_hook_to_point(
1177 hook_point: HookPoint,
1178 hook_fn: Callable,
1179 name: str,
1180 dir: Literal["fwd", "bwd"] = "fwd",
1181 ) -> None:
1182 self.check_hooks_to_add(hook_point, name, hook_fn, dir=dir)
1183 if self.compatibility_mode and name != hook_point.name: 1183 ↛ 1184line 1183 didn't jump to line 1184 because the condition on line 1183 was never true
1184 alias_names_list: list = []
1185 if hook_point.name is not None:
1186 alias_names_list.append(hook_point.name)
1187 alias_names_list.append(name)
1188 handle = hook_point.add_hook(
1189 hook_fn,
1190 dir=dir,
1191 level=context_level,
1192 alias_names=alias_names_list,
1193 )
1194 else:
1195 handle = hook_point.add_hook(hook_fn, dir=dir, level=context_level)
1196 added_hooks.append((hook_point, dir, handle))
1198 def apply_hooks(hook_list: List[Tuple[Any, Callable]], is_fwd: bool) -> None:
1199 direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
1200 hook_dict = self.hook_dict
1201 aliases = build_alias_to_canonical_map(hook_dict)
1202 for hook_name_or_filter, hook_fn in hook_list:
1203 if isinstance(hook_name_or_filter, str):
1204 actual_hook_name, hook_point = self._resolve_hook_point(
1205 hook_name_or_filter, aliases, hook_dict
1206 )
1207 add_hook_to_point(hook_point, hook_fn, actual_hook_name, direction)
1208 else:
1209 seen_hooks = set()
1210 gated_skipped: List[str] = []
1211 for n, hook_point in hook_dict.items():
1212 if hook_name_or_filter(n):
1213 hook_id = id(hook_point)
1214 if hook_id in seen_hooks:
1215 continue
1216 seen_hooks.add(hook_id)
1217 if self._gated_hook_reason(hook_point.name or n) is not None:
1218 gated_skipped.append(n)
1219 continue
1220 hook_name_to_use = hook_point.name if hook_point.name else n
1221 add_hook_to_point(hook_point, hook_fn, hook_name_to_use, direction)
1222 self._warn_gated_skipped("hooks", gated_skipped, stacklevel=4)
1224 try:
1225 apply_hooks(fwd_hooks, True)
1226 apply_hooks(bwd_hooks, False)
1227 yield self
1228 finally:
1229 try:
1230 for hook_point, direction, handle in added_hooks:
1231 if reset_hooks_end:
1232 self._remove_temporary_hook(hook_point, direction, handle)
1233 if clear_contexts: 1233 ↛ 1234line 1233 didn't jump to line 1234 because the condition on line 1233 was never true
1234 hook_point.clear_context()
1235 finally:
1236 self.context_level -= 1
1238 return _hooks_context()
1240 # ---- high-level execution: run_with_hooks ----
1242 def run_with_hooks(
1243 self,
1244 input: Any,
1245 fwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
1246 bwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
1247 reset_hooks_end: bool = True,
1248 clear_contexts: bool = False,
1249 return_type: Optional[str] = "logits",
1250 stop_at_layer: Optional[int] = None,
1251 start_at_layer: Optional[int] = None,
1252 remove_batch_dim: bool = False,
1253 **kwargs: Any,
1254 ) -> Any:
1255 """Run the model with specified forward and backward hooks.
1257 ``stop_at_layer`` raises :class:`StopAtLayerException` to stop early
1258 (KV cache cleaned up on stop). ``start_at_layer`` treats ``input`` as the
1259 residual entering block ``k`` (see :meth:`forward`); hooks on blocks below
1260 ``k`` are skipped to match HookedTransformer. ``remove_batch_dim``
1261 squeezes/unsqueezes the leading dim around each hook callback whenever
1262 that tensor's leading dim is 1 — a per-tensor rule, so broadcast
1263 activations such as rotary ``hook_cos``/``hook_sin`` lose their size-1
1264 leading dim even at batch > 1.
1265 ``reset_hooks_end`` removes the hooks this call added when it finishes —
1266 hooks the caller attached beforehand are left alone either way;
1267 ``clear_contexts`` also wipes the touched hook points' ``ctx``.
1268 """
1269 if "names_filter" in kwargs:
1270 # **kwargs would silently absorb it; fail loud.
1271 raise TypeError(
1272 "run_with_hooks() got an unexpected keyword argument 'names_filter'; "
1273 "use run_with_cache(names_filter=...) to scope caching."
1274 )
1275 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = []
1276 effective_stop_layer = None
1277 if stop_at_layer is not None and hasattr(self, "blocks"):
1278 if stop_at_layer < 0: 1278 ↛ 1279line 1278 didn't jump to line 1279 because the condition on line 1278 was never true
1279 effective_stop_layer = len(self.blocks) + stop_at_layer
1280 else:
1281 effective_stop_layer = stop_at_layer
1282 effective_start_layer = None
1283 if start_at_layer is not None and hasattr(self, "blocks"): 1283 ↛ 1288line 1283 didn't jump to line 1288 because the condition on line 1283 was always true
1284 effective_start_layer = (
1285 len(self.blocks) + start_at_layer if start_at_layer < 0 else start_at_layer
1286 )
1288 def add_hook_to_point(
1289 hook_point: HookPoint,
1290 hook_fn: Callable,
1291 name: str,
1292 dir: Literal["fwd", "bwd"] = "fwd",
1293 ) -> None:
1294 if effective_start_layer is not None and self._is_embedding_stage_hook(name): 1294 ↛ 1295line 1294 didn't jump to line 1295 because the condition on line 1294 was never true
1295 return
1296 if name.startswith("blocks."):
1297 try:
1298 layer_num: Optional[int] = int(name.split(".")[1])
1299 except (IndexError, ValueError):
1300 layer_num = None
1301 if layer_num is not None: 1301 ↛ 1306line 1301 didn't jump to line 1306 because the condition on line 1301 was always true
1302 if effective_stop_layer is not None and layer_num >= effective_stop_layer:
1303 return
1304 if effective_start_layer is not None and layer_num < effective_start_layer:
1305 return
1306 self.check_hooks_to_add(hook_point, name, hook_fn, dir=dir)
1307 if self.compatibility_mode and name != hook_point.name: 1307 ↛ 1308line 1307 didn't jump to line 1308 because the condition on line 1307 was never true
1308 alias_names_list: list = []
1309 if hook_point.name is not None:
1310 alias_names_list.append(hook_point.name)
1311 alias_names_list.append(name)
1312 handle = hook_point.add_hook(
1313 hook_fn,
1314 dir=dir,
1315 level=context_level,
1316 alias_names=alias_names_list,
1317 )
1318 else:
1319 handle = hook_point.add_hook(hook_fn, dir=dir, level=context_level)
1320 added_hooks.append((hook_point, dir, handle))
1322 def apply_hooks(
1323 hook_list: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool
1324 ) -> None:
1325 direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
1326 hook_dict = self.hook_dict
1327 aliases = build_alias_to_canonical_map(hook_dict)
1328 for hook_name_or_filter, hook_fn in hook_list:
1329 if remove_batch_dim:
1330 original_hook_fn = hook_fn
1332 # Default arg captures hook_fn by value (avoids closure issue)
1333 def wrapped_hook_fn(tensor, hook, _orig_fn=original_hook_fn):
1334 if tensor.shape[0] == 1:
1335 tensor_no_batch = tensor.squeeze(0)
1336 result = _orig_fn(tensor_no_batch, hook)
1337 if result is not None and result.dim() == tensor_no_batch.dim():
1338 result = result.unsqueeze(0)
1339 return result
1340 else:
1341 return _orig_fn(tensor, hook)
1343 hook_fn = wrapped_hook_fn
1344 if isinstance(hook_name_or_filter, str):
1345 # Resolve BEFORE gating: an aliased HT name may point at an
1346 # always-firing override, which must not be refused. The
1347 # message still names the caller's spelling.
1348 actual_hook_name, hook_point = self._resolve_hook_point(
1349 hook_name_or_filter, aliases, hook_dict
1350 )
1351 reason = self._gated_hook_reason(hook_point.name or actual_hook_name)
1352 if reason is not None:
1353 raise ValueError(
1354 f"Cannot add hook {hook_name_or_filter} because {reason} is False. "
1355 f"Call set_{reason}(True) first."
1356 )
1357 add_hook_to_point(hook_point, hook_fn, actual_hook_name, direction)
1358 else:
1359 seen_hooks: set = set()
1360 gated_skipped: List[str] = []
1361 for n, hook_point in hook_dict.items():
1362 if hook_name_or_filter(n):
1363 hook_id = id(hook_point)
1364 if hook_id in seen_hooks: 1364 ↛ 1365line 1364 didn't jump to line 1365 because the condition on line 1364 was never true
1365 continue
1366 seen_hooks.add(hook_id)
1367 if self._gated_hook_reason(hook_point.name or n) is not None:
1368 gated_skipped.append(n)
1369 continue
1370 hook_name_to_use = hook_point.name if hook_point.name else n
1371 add_hook_to_point(hook_point, hook_fn, hook_name_to_use, direction)
1372 self._warn_gated_skipped("run_with_hooks", gated_skipped, stacklevel=4)
1374 context_level = getattr(self, "context_level", 0) + 1
1375 self.context_level = context_level
1376 try:
1377 if effective_stop_layer is not None:
1378 stop_at_layer = effective_stop_layer
1379 apply_hooks(fwd_hooks, True)
1380 apply_hooks(bwd_hooks, False)
1381 if start_at_layer is not None:
1382 kwargs["start_at_layer"] = start_at_layer
1383 try:
1384 output = self.forward(
1385 input, return_type=return_type, stop_at_layer=stop_at_layer, **kwargs
1386 )
1387 except StopAtLayerException as e:
1388 output = e.layer_output
1389 return output
1390 finally:
1391 try:
1392 for hook_point, direction, handle in added_hooks:
1393 if reset_hooks_end:
1394 self._remove_temporary_hook(hook_point, direction, handle)
1395 if clear_contexts:
1396 hook_point.clear_context()
1397 finally:
1398 self.context_level -= 1
1400 # ---- high-level execution: run_with_cache ----
1402 @overload
1403 def run_with_cache(
1404 self,
1405 input: Union[str, List[str], torch.Tensor],
1406 return_cache_object: Literal[True] = True,
1407 remove_batch_dim: bool = False,
1408 **kwargs,
1409 ) -> Tuple[Any, ActivationCache]:
1410 """Run with cache - placeholder implementation."""
1411 pass
1413 @overload
1414 def run_with_cache(
1415 self,
1416 input: Union[str, List[str], torch.Tensor],
1417 return_cache_object: Literal[False],
1418 remove_batch_dim: bool = False,
1419 **kwargs,
1420 ) -> Tuple[Any, Dict[str, torch.Tensor]]:
1421 """Run with cache - placeholder implementation."""
1422 pass
1424 def run_with_cache(
1425 self,
1426 input: Union[str, List[str], torch.Tensor],
1427 return_cache_object: bool = True,
1428 remove_batch_dim: bool = False,
1429 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
1430 stop_at_layer: Optional[int] = None,
1431 start_at_layer: Optional[int] = None,
1432 pos_slice: Optional[Union[Slice, SliceInput]] = None,
1433 incl_bwd: bool = False,
1434 reset_hooks_end: bool = True,
1435 clear_contexts: bool = False,
1436 **kwargs,
1437 ) -> Tuple[Any, Union[ActivationCache, Dict[str, torch.Tensor]]]:
1438 """Run the model and cache activations. Returns ``(output, cache)``.
1440 ``stop_at_layer`` raises :class:`StopAtLayerException` to stop early.
1441 ``start_at_layer`` treats ``input`` as the residual entering block ``k``
1442 (see :meth:`forward`); blocks below ``k`` are excluded from the cache to
1443 match HookedTransformer. ``pos_slice`` slices each cached activation along
1444 its position dimension (dim 1 for resid/per-head/token-id activations; the
1445 query position ``-2`` for attention patterns/scores; ``HookPoint.pos_dim``
1446 for hooks fired on another layout). ``incl_bwd`` also caches gradients
1447 under ``"<name>_grad"`` by running ``output.backward()``;
1448 the caller must request a scalar output (``return_type="loss"``) and the
1449 model must be on the gradients-capable transformers driver.
1450 ``reset_hooks_end`` removes the hooks this call added when it finishes —
1451 hooks the caller attached beforehand are left alone either way;
1452 ``clear_contexts`` also wipes the touched hook points' ``ctx``. ``device``
1453 offloads cached activations (matches ``ActivationCache.to``); the model and
1454 inputs stay where the caller put them.
1455 """
1456 if incl_bwd and stop_at_layer is not None:
1457 raise ValueError(
1458 "incl_bwd=True cannot be combined with stop_at_layer: the run returns the "
1459 "intermediate activation at that layer, not a scalar to call backward() on."
1460 )
1461 if incl_bwd and not torch.is_grad_enabled():
1462 raise ValueError(
1463 "incl_bwd=True needs autograd, but gradient tracking is off "
1464 "(torch.no_grad(), set_grad_enabled(False) or inference mode). "
1465 "Run the call with gradients enabled."
1466 )
1467 pos_slice_obj = Slice.unwrap(pos_slice)
1468 aliases = build_alias_to_canonical_map(self.hook_dict)
1470 def create_names_filter_fn(filter_input):
1471 if filter_input is None:
1472 return lambda name: True
1473 elif isinstance(filter_input, str):
1474 mapped_name = aliases.get(filter_input, None)
1475 if mapped_name:
1476 return lambda name: name == mapped_name or name == filter_input
1477 else:
1478 return lambda name: name == filter_input
1479 elif isinstance(filter_input, list):
1480 mapped_list = []
1481 for item in filter_input:
1482 mapped_list.append(item)
1483 mapped_name = aliases.get(item, None)
1484 if mapped_name:
1485 mapped_list.append(mapped_name)
1486 return lambda name: name in mapped_list
1487 elif callable(filter_input): 1487 ↛ 1490line 1487 didn't jump to line 1490 because the condition on line 1487 was always true
1488 return filter_input
1489 else:
1490 raise ValueError("names_filter must be a string, list of strings, or callable")
1492 names_filter_fn = create_names_filter_fn(names_filter)
1493 if isinstance(names_filter, (str, list)):
1494 requested = [names_filter] if isinstance(names_filter, str) else names_filter
1495 for name in requested:
1496 self._check_hook_fireable(name, aliases.get(name, name))
1497 cache: Dict[str, torch.Tensor] = {}
1498 hooks: List[Tuple[HookPoint, str]] = []
1499 visited: set[int] = set()
1500 stop_hook_point: Optional[HookPoint] = None
1501 stop_hook_fn: Optional[Callable] = None
1503 # None → no-op .to(None), tensors stay on their current device.
1504 cache_device = kwargs.pop("device", None)
1506 def _store(name: str, value: torch.Tensor, hook: Any, suffix: str = "") -> None:
1507 stored = value.detach().to(cache_device)
1508 if pos_slice_obj is not None and stored.dim() >= 2: 1508 ↛ 1512line 1508 didn't jump to line 1512 because the condition on line 1508 was always true
1509 # A gradient has its activation's layout, so the axis comes from
1510 # the unsuffixed name — see _pos_slice_dim.
1511 stored = pos_slice_obj.apply(stored, dim=self._pos_slice_dim(name, hook))
1512 cache[name + suffix] = stored
1514 def make_cache_hook(name: str):
1515 def cache_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
1516 if tensor is None: 1516 ↛ 1517line 1516 didn't jump to line 1517 because the condition on line 1516 was never true
1517 cache[name] = None
1518 elif isinstance(tensor, torch.Tensor): 1518 ↛ 1520line 1518 didn't jump to line 1520 because the condition on line 1518 was always true
1519 _store(name, tensor, hook)
1520 elif isinstance(tensor, tuple):
1521 if len(tensor) > 0 and isinstance(tensor[0], torch.Tensor):
1522 _store(name, tensor[0], hook)
1523 else:
1524 pass
1525 else:
1526 try:
1527 if hasattr(tensor, "detach"):
1528 _store(name, tensor, hook)
1529 except:
1530 pass
1531 return tensor
1533 return cache_hook
1535 hook_dict = self.hook_dict
1536 effective_stop_layer = None
1537 if stop_at_layer is not None and hasattr(self, "blocks"):
1538 if stop_at_layer < 0:
1539 effective_stop_layer = len(self.blocks) + stop_at_layer
1540 else:
1541 effective_stop_layer = stop_at_layer
1542 effective_start_layer = None
1543 if start_at_layer is not None and hasattr(self, "blocks"):
1544 effective_start_layer = (
1545 len(self.blocks) + start_at_layer if start_at_layer < 0 else start_at_layer
1546 )
1547 matched_any = False
1548 for hook_name, hook in hook_dict.items():
1549 if names_filter_fn(hook_name):
1550 matched_any = True
1551 if effective_start_layer is not None and self._is_embedding_stage_hook(hook_name):
1552 continue
1553 if hook_name.startswith("blocks."):
1554 try:
1555 layer_num = int(hook_name.split(".")[1])
1556 except (IndexError, ValueError):
1557 layer_num = None
1558 if layer_num is not None: 1558 ↛ 1566line 1558 didn't jump to line 1566 because the condition on line 1558 was always true
1559 # stop/start bound the executed range; blocks outside it
1560 # either don't run (stop) or run on discarded input (start),
1561 # so their activations must not enter the cache.
1562 if effective_stop_layer is not None and layer_num >= effective_stop_layer:
1563 continue
1564 if effective_start_layer is not None and layer_num < effective_start_layer:
1565 continue
1566 hooks.append((hook, hook_name))
1567 # Explicit string/list filters matching nothing must not return (logits, {}) silently.
1568 if not matched_any and names_filter and isinstance(names_filter, (str, list)): 1568 ↛ 1574line 1568 didn't jump to line 1574 because the condition on line 1568 was always true
1569 raise KeyError(
1570 f"names_filter {names_filter!r} matched no hook points on this model; "
1571 "check the name against model.hook_dict (this backend may not serve it)."
1572 )
1574 def make_grad_cache_hook(name: str):
1575 def grad_hook(tensor: torch.Tensor, *, hook: Any) -> None:
1576 if isinstance(tensor, torch.Tensor): 1576 ↛ 1579line 1576 didn't jump to line 1579 because the condition on line 1576 was always true
1577 _store(name, tensor, hook, suffix="_grad")
1578 # A non-None return from a backward hook replaces grad_input; stay read-only.
1579 return None
1581 return grad_hook
1583 processed_args = [input]
1584 # Driver-aware input placement: torch drivers move input_ids to the model's
1585 # device; remote drivers (no local parameters) leave them as-is.
1586 target_device = self._input_device()
1587 if processed_args and isinstance(processed_args[0], str):
1588 assert self.tokenizer is not None, "Tokenizer must be set to pass string input."
1589 prepend_bos = kwargs.pop("prepend_bos", None)
1590 input_ids = self.to_tokens(processed_args[0], prepend_bos=prepend_bos)
1591 if target_device is not None: 1591 ↛ 1593line 1591 didn't jump to line 1593 because the condition on line 1591 was always true
1592 input_ids = input_ids.to(target_device)
1593 kwargs["input_ids"] = input_ids
1594 processed_args = processed_args[1:]
1595 elif "input" in kwargs and isinstance(kwargs["input"], str): 1595 ↛ 1596line 1595 didn't jump to line 1596 because the condition on line 1595 was never true
1596 assert self.tokenizer is not None, "Tokenizer must be set to pass string input."
1597 prepend_bos = kwargs.pop("prepend_bos", None)
1598 input_ids = self.to_tokens(kwargs["input"], prepend_bos=prepend_bos)
1599 if target_device is not None:
1600 input_ids = input_ids.to(target_device)
1601 kwargs["input_ids"] = input_ids
1602 del kwargs["input"]
1603 if stop_at_layer is not None and hasattr(self, "blocks"):
1604 if stop_at_layer < 0:
1605 stop_at_layer = len(self.blocks) + stop_at_layer
1606 last_layer_to_process = stop_at_layer - 1
1608 def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
1609 raise StopAtLayerException(tensor)
1611 if stop_at_layer >= 0 and stop_at_layer < len(self.blocks): 1611 ↛ 1618line 1611 didn't jump to line 1618 because the condition on line 1611 was always true
1612 # Stop at the beginning of the specified block, not at the end of the previous block
1613 block_hook_name = f"blocks.{stop_at_layer}.hook_in"
1614 hook_dict = self.hook_dict
1615 if block_hook_name in hook_dict: 1615 ↛ 1618line 1615 didn't jump to line 1618 because the condition on line 1615 was always true
1616 stop_hook_point = hook_dict[block_hook_name]
1617 stop_hook_fn = stop_hook
1618 filtered_kwargs = kwargs.copy()
1619 # ``cache_device`` is honored by ``make_cache_hook`` above (``tensor.detach().to(cache_device)``);
1620 # the model and inputs stay where the caller put them, matching ``ActivationCache.to``.
1621 if cache_device is not None and getattr(self.cfg, "n_devices", 1) > 1:
1622 # Moving a dispatched model to a single device collapses accelerate's
1623 # split and breaks its routing hooks. The cache will stay spread across
1624 # the per-layer devices; callers can .to(cache_device) on cache entries
1625 # after the fact if they need a single-device cache.
1626 warnings.warn(
1627 f"run_with_cache(device={cache_device!r}) ignored: model is dispatched "
1628 f"across {self.cfg.n_devices} devices via device_map. Cached activations "
1629 "will remain on their per-layer devices.",
1630 stacklevel=2,
1631 )
1632 if start_at_layer is not None:
1633 filtered_kwargs["start_at_layer"] = start_at_layer
1634 # Only validate gated hooks when the caller explicitly supplied a
1635 # names_filter. The default filter matches every hook and must not
1636 # cause gated hooks to be treated as explicitly requested.
1637 gated_names_skipped: List[str] = []
1638 if names_filter is not None:
1639 kept = []
1640 for hp, name in hooks:
1641 if self._gated_hook_reason(hp.name or name) is not None:
1642 gated_names_skipped.append(name)
1643 else:
1644 kept.append((hp, name))
1645 hooks = kept
1646 if gated_names_skipped:
1647 warnings.warn(
1648 f"run_with_cache: skipped {len(gated_names_skipped)} gated-off hook name(s) "
1649 f"that will never be cached: {gated_names_skipped}. Call the relevant "
1650 "set_use_*(True) setter first to enable them.",
1651 stacklevel=2,
1652 )
1653 context_level = getattr(self, "context_level", 0) + 1
1654 self.context_level = context_level
1655 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = []
1656 stop_handle: Optional[LensHandle] = None
1657 try:
1658 for hp, name in hooks:
1659 fwd_handle = hp.add_hook(make_cache_hook(name), level=context_level)
1660 added_hooks.append((hp, "fwd", fwd_handle))
1661 if incl_bwd:
1662 bwd_handle = hp.add_hook(
1663 make_grad_cache_hook(name), dir="bwd", level=context_level
1664 )
1665 added_hooks.append((hp, "bwd", bwd_handle))
1666 if stop_hook_point is not None and stop_hook_fn is not None:
1667 stop_handle = stop_hook_point.add_hook(stop_hook_fn, level=context_level)
1668 added_hooks.append((stop_hook_point, "fwd", stop_handle))
1669 except Exception:
1670 try:
1671 for hp, direction, handle in added_hooks:
1672 self._remove_temporary_hook(hp, direction, handle)
1673 finally:
1674 self.context_level -= 1
1675 raise
1676 try:
1677 if (
1678 "output_attentions" not in filtered_kwargs
1679 and self.adapter.supports_hf_output_attentions
1680 ):
1681 # Attention-free remote-code models (e.g. HyenaDNA) reject the
1682 # kwarg outright; only pass it when the forward actually accepts
1683 # it. Non-torch drivers expose no local module to introspect —
1684 # keep the kwarg for them, as before.
1685 underlying = getattr(self._driver, "underlying_model", None)
1686 if underlying is None:
1687 filtered_kwargs["output_attentions"] = True
1688 else:
1689 fwd_params = inspect.signature(underlying.forward).parameters
1690 if "output_attentions" in fwd_params or any(
1691 p.kind is inspect.Parameter.VAR_KEYWORD for p in fwd_params.values()
1692 ):
1693 filtered_kwargs["output_attentions"] = True
1694 if processed_args:
1695 output = self.forward(processed_args[0], **filtered_kwargs)
1696 elif "input_ids" in filtered_kwargs: 1696 ↛ 1702line 1696 didn't jump to line 1702 because the condition on line 1696 was always true
1697 output = self.forward(
1698 filtered_kwargs["input_ids"],
1699 **{k: v for k, v in filtered_kwargs.items() if k != "input_ids"},
1700 )
1701 else:
1702 output = self.forward(**filtered_kwargs)
1703 if hasattr(output, "logits"): 1703 ↛ 1704line 1703 didn't jump to line 1704 because the condition on line 1703 was never true
1704 output = output.logits
1705 if incl_bwd:
1706 # Gradients land in the cache via the bwd hooks, which the finally below
1707 # removes, so the backward pass has to happen inside this try.
1708 if not isinstance(output, torch.Tensor) or output.numel() != 1:
1709 shape = tuple(output.shape) if isinstance(output, torch.Tensor) else None
1710 raise ValueError(
1711 "incl_bwd=True needs a scalar output to call backward() on, got "
1712 f"{type(output).__name__}{f' of shape {shape}' if shape else ''}. "
1713 'Pass return_type="loss".'
1714 )
1715 if not output.requires_grad:
1716 raise ValueError(
1717 "incl_bwd=True got an output with no grad_fn — the model's parameters "
1718 "have requires_grad=False, so there is nothing to differentiate."
1719 )
1720 output.backward()
1721 except StopAtLayerException as e:
1722 output = e.layer_output
1723 except Exception as e:
1724 raise e
1725 finally:
1726 try:
1727 for hp, direction, handle in added_hooks:
1728 # Retain requested cache hooks, but never internal early-stop hooks.
1729 if reset_hooks_end or handle is stop_handle:
1730 self._remove_temporary_hook(hp, direction, handle)
1731 if clear_contexts:
1732 for hp, _ in hooks:
1733 hp.clear_context()
1734 if stop_hook_point is not None: 1734 ↛ 1735line 1734 didn't jump to line 1735 because the condition on line 1734 was never true
1735 stop_hook_point.clear_context()
1736 finally:
1737 self.context_level -= 1
1738 if self.compatibility_mode == True:
1739 reverse_aliases = {}
1740 for old_name, new_name in aliases.items():
1741 if isinstance(new_name, list): 1741 ↛ 1742line 1741 didn't jump to line 1742 because the condition on line 1741 was never true
1742 for single_new_name in new_name:
1743 reverse_aliases[single_new_name] = old_name
1744 else:
1745 reverse_aliases[new_name] = old_name
1746 # Gradient entries are keyed "<hook_name>_grad", so alias lookups run on the
1747 # base name and the suffix is re-attached to the aliased key.
1748 suffixes = ("", "_grad") if incl_bwd else ("",)
1749 cache_items_to_add = {}
1750 for cache_name, cached_value in cache.items():
1751 base_name, suffix = (
1752 (cache_name[: -len("_grad")], "_grad")
1753 if cache_name.endswith("_grad")
1754 else (cache_name, "")
1755 )
1756 old_name = reverse_aliases.get(base_name)
1757 if old_name is not None:
1758 cache_items_to_add[old_name + suffix] = cached_value
1759 cache.update(cache_items_to_add)
1760 for alias_name, target_name in aliases.items():
1761 targets = target_name if isinstance(target_name, list) else [target_name]
1762 for suffix in suffixes:
1763 if alias_name + suffix in cache:
1764 continue
1765 for single_target in targets:
1766 if single_target + suffix in cache: 1766 ↛ 1767line 1766 didn't jump to line 1767 because the condition on line 1766 was never true
1767 cache[alias_name + suffix] = cache[single_target + suffix]
1768 break
1769 # The input states the batch size; inferring it from the cached shapes goes wrong
1770 # when the filter selects only flattened or position-indexed hooks.
1771 batch_size = _input_batch_size(
1772 processed_args[0] if processed_args else None,
1773 filtered_kwargs.get("input_ids"),
1774 filtered_kwargs.get("inputs_embeds"),
1775 )
1776 if return_cache_object:
1777 activation_cache = ActivationCache(cache, self, has_batch_dim=True)
1778 if remove_batch_dim:
1779 activation_cache.remove_batch_dim(batch_size=batch_size)
1780 return (output, activation_cache)
1781 else:
1782 if remove_batch_dim:
1783 # Same batch-size check and squeeze rule as the ActivationCache path.
1784 tensors = {k: v for k, v in cache.items() if isinstance(v, torch.Tensor)}
1785 cache.update(
1786 ActivationCache(tensors, self)
1787 .remove_batch_dim(batch_size=batch_size)
1788 .cache_dict
1789 )
1790 return (output, cache)