Coverage for transformer_lens/model_bridge/bridge_core.py: 85%
919 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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
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 def hooks(
1142 self,
1143 fwd_hooks: List = [],
1144 bwd_hooks: List = [],
1145 reset_hooks_end: bool = True,
1146 clear_contexts: bool = False,
1147 ) -> Any:
1148 """Context manager for temporarily adding hooks.
1150 ``reset_hooks_end`` removes the hooks this context added when it exits —
1151 hooks the caller attached beforehand are left alone either way;
1152 ``clear_contexts`` also wipes the touched hook points' ``ctx``.
1154 Example:
1155 with model.hooks(fwd_hooks=[("hook_embed", my_hook)]):
1156 output = model("Hello world")
1157 """
1159 @contextmanager
1160 def _hooks_context() -> Iterator["BridgeCore"]:
1161 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = []
1162 context_level = getattr(self, "context_level", 0) + 1
1163 self.context_level = context_level
1165 def add_hook_to_point(
1166 hook_point: HookPoint,
1167 hook_fn: Callable,
1168 name: str,
1169 dir: Literal["fwd", "bwd"] = "fwd",
1170 ) -> None:
1171 self.check_hooks_to_add(hook_point, name, hook_fn, dir=dir)
1172 if self.compatibility_mode and name != hook_point.name: 1172 ↛ 1173line 1172 didn't jump to line 1173 because the condition on line 1172 was never true
1173 alias_names_list: list = []
1174 if hook_point.name is not None:
1175 alias_names_list.append(hook_point.name)
1176 alias_names_list.append(name)
1177 hook_point.add_hook(
1178 hook_fn,
1179 dir=dir,
1180 level=context_level,
1181 alias_names=alias_names_list,
1182 )
1183 else:
1184 hook_point.add_hook(hook_fn, dir=dir, level=context_level)
1185 added_hooks.append((hook_point, dir))
1187 def apply_hooks(hook_list: List[Tuple[Any, Callable]], is_fwd: bool) -> None:
1188 direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
1189 hook_dict = self.hook_dict
1190 aliases = build_alias_to_canonical_map(hook_dict)
1191 for hook_name_or_filter, hook_fn in hook_list:
1192 if isinstance(hook_name_or_filter, str):
1193 actual_hook_name, hook_point = self._resolve_hook_point(
1194 hook_name_or_filter, aliases, hook_dict
1195 )
1196 add_hook_to_point(hook_point, hook_fn, actual_hook_name, direction)
1197 else:
1198 seen_hooks = set()
1199 gated_skipped: List[str] = []
1200 for n, hook_point in hook_dict.items():
1201 if hook_name_or_filter(n):
1202 hook_id = id(hook_point)
1203 if hook_id in seen_hooks:
1204 continue
1205 seen_hooks.add(hook_id)
1206 if self._gated_hook_reason(hook_point.name or n) is not None:
1207 gated_skipped.append(n)
1208 continue
1209 hook_name_to_use = hook_point.name if hook_point.name else n
1210 add_hook_to_point(hook_point, hook_fn, hook_name_to_use, direction)
1211 self._warn_gated_skipped("hooks", gated_skipped, stacklevel=4)
1213 try:
1214 apply_hooks(fwd_hooks, True)
1215 apply_hooks(bwd_hooks, False)
1216 yield self
1217 finally:
1218 try:
1219 for hook_point, direction in added_hooks:
1220 if reset_hooks_end:
1221 # `level` keeps this to the hooks added above — hooks the
1222 # caller attached before the context survive.
1223 hook_point.remove_hooks(dir=direction, level=context_level)
1224 if clear_contexts: 1224 ↛ 1225line 1224 didn't jump to line 1225 because the condition on line 1224 was never true
1225 hook_point.clear_context()
1226 finally:
1227 self.context_level -= 1
1229 return _hooks_context()
1231 # ---- high-level execution: run_with_hooks ----
1233 def run_with_hooks(
1234 self,
1235 input: Any,
1236 fwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
1237 bwd_hooks: List[Tuple[Union[str, Callable], Callable]] = [],
1238 reset_hooks_end: bool = True,
1239 clear_contexts: bool = False,
1240 return_type: Optional[str] = "logits",
1241 stop_at_layer: Optional[int] = None,
1242 start_at_layer: Optional[int] = None,
1243 remove_batch_dim: bool = False,
1244 **kwargs: Any,
1245 ) -> Any:
1246 """Run the model with specified forward and backward hooks.
1248 ``stop_at_layer`` raises :class:`StopAtLayerException` to stop early
1249 (KV cache cleaned up on stop). ``start_at_layer`` treats ``input`` as the
1250 residual entering block ``k`` (see :meth:`forward`); hooks on blocks below
1251 ``k`` are skipped to match HookedTransformer. ``remove_batch_dim``
1252 squeezes/unsqueezes the leading dim around each hook callback whenever
1253 that tensor's leading dim is 1 — a per-tensor rule, so broadcast
1254 activations such as rotary ``hook_cos``/``hook_sin`` lose their size-1
1255 leading dim even at batch > 1.
1256 ``reset_hooks_end`` removes the hooks this call added when it finishes —
1257 hooks the caller attached beforehand are left alone either way;
1258 ``clear_contexts`` also wipes the touched hook points' ``ctx``.
1259 """
1260 if "names_filter" in kwargs:
1261 # **kwargs would silently absorb it; fail loud.
1262 raise TypeError(
1263 "run_with_hooks() got an unexpected keyword argument 'names_filter'; "
1264 "use run_with_cache(names_filter=...) to scope caching."
1265 )
1266 added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"]]] = []
1267 effective_stop_layer = None
1268 if stop_at_layer is not None and hasattr(self, "blocks"):
1269 if stop_at_layer < 0: 1269 ↛ 1270line 1269 didn't jump to line 1270 because the condition on line 1269 was never true
1270 effective_stop_layer = len(self.blocks) + stop_at_layer
1271 else:
1272 effective_stop_layer = stop_at_layer
1273 effective_start_layer = None
1274 if start_at_layer is not None and hasattr(self, "blocks"): 1274 ↛ 1279line 1274 didn't jump to line 1279 because the condition on line 1274 was always true
1275 effective_start_layer = (
1276 len(self.blocks) + start_at_layer if start_at_layer < 0 else start_at_layer
1277 )
1279 def add_hook_to_point(
1280 hook_point: HookPoint,
1281 hook_fn: Callable,
1282 name: str,
1283 dir: Literal["fwd", "bwd"] = "fwd",
1284 ) -> None:
1285 if effective_start_layer is not None and self._is_embedding_stage_hook(name): 1285 ↛ 1286line 1285 didn't jump to line 1286 because the condition on line 1285 was never true
1286 return
1287 if name.startswith("blocks."):
1288 try:
1289 layer_num: Optional[int] = int(name.split(".")[1])
1290 except (IndexError, ValueError):
1291 layer_num = None
1292 if layer_num is not None: 1292 ↛ 1297line 1292 didn't jump to line 1297 because the condition on line 1292 was always true
1293 if effective_stop_layer is not None and layer_num >= effective_stop_layer:
1294 return
1295 if effective_start_layer is not None and layer_num < effective_start_layer:
1296 return
1297 self.check_hooks_to_add(hook_point, name, hook_fn, dir=dir)
1298 if self.compatibility_mode and name != hook_point.name: 1298 ↛ 1299line 1298 didn't jump to line 1299 because the condition on line 1298 was never true
1299 alias_names_list: list = []
1300 if hook_point.name is not None:
1301 alias_names_list.append(hook_point.name)
1302 alias_names_list.append(name)
1303 hook_point.add_hook(
1304 hook_fn,
1305 dir=dir,
1306 level=context_level,
1307 alias_names=alias_names_list,
1308 )
1309 else:
1310 hook_point.add_hook(hook_fn, dir=dir, level=context_level)
1311 added_hooks.append((hook_point, dir))
1313 def apply_hooks(
1314 hook_list: List[Tuple[Union[str, Callable], Callable]], is_fwd: bool
1315 ) -> None:
1316 direction: Literal["fwd", "bwd"] = "fwd" if is_fwd else "bwd"
1317 hook_dict = self.hook_dict
1318 aliases = build_alias_to_canonical_map(hook_dict)
1319 for hook_name_or_filter, hook_fn in hook_list:
1320 if remove_batch_dim:
1321 original_hook_fn = hook_fn
1323 # Default arg captures hook_fn by value (avoids closure issue)
1324 def wrapped_hook_fn(tensor, hook, _orig_fn=original_hook_fn):
1325 if tensor.shape[0] == 1:
1326 tensor_no_batch = tensor.squeeze(0)
1327 result = _orig_fn(tensor_no_batch, hook)
1328 if result is not None and result.dim() == tensor_no_batch.dim():
1329 result = result.unsqueeze(0)
1330 return result
1331 else:
1332 return _orig_fn(tensor, hook)
1334 hook_fn = wrapped_hook_fn
1335 if isinstance(hook_name_or_filter, str):
1336 # Resolve BEFORE gating: an aliased HT name may point at an
1337 # always-firing override, which must not be refused. The
1338 # message still names the caller's spelling.
1339 actual_hook_name, hook_point = self._resolve_hook_point(
1340 hook_name_or_filter, aliases, hook_dict
1341 )
1342 reason = self._gated_hook_reason(hook_point.name or actual_hook_name)
1343 if reason is not None:
1344 raise ValueError(
1345 f"Cannot add hook {hook_name_or_filter} because {reason} is False. "
1346 f"Call set_{reason}(True) first."
1347 )
1348 add_hook_to_point(hook_point, hook_fn, actual_hook_name, direction)
1349 else:
1350 seen_hooks: set = set()
1351 gated_skipped: List[str] = []
1352 for n, hook_point in hook_dict.items():
1353 if hook_name_or_filter(n):
1354 hook_id = id(hook_point)
1355 if hook_id in seen_hooks: 1355 ↛ 1356line 1355 didn't jump to line 1356 because the condition on line 1355 was never true
1356 continue
1357 seen_hooks.add(hook_id)
1358 if self._gated_hook_reason(hook_point.name or n) is not None:
1359 gated_skipped.append(n)
1360 continue
1361 hook_name_to_use = hook_point.name if hook_point.name else n
1362 add_hook_to_point(hook_point, hook_fn, hook_name_to_use, direction)
1363 self._warn_gated_skipped("run_with_hooks", gated_skipped, stacklevel=4)
1365 context_level = getattr(self, "context_level", 0) + 1
1366 self.context_level = context_level
1367 try:
1368 if stop_at_layer is not None and hasattr(self, "blocks"):
1369 if stop_at_layer < 0: 1369 ↛ 1370line 1369 didn't jump to line 1370 because the condition on line 1369 was never true
1370 stop_at_layer = len(self.blocks) + stop_at_layer
1371 if stop_at_layer >= 0 and stop_at_layer < len(self.blocks): 1371 ↛ anywhereline 1371 didn't jump anywhere: it always raised an exception.
1373 def stop_hook(tensor: Any, *, hook: Any) -> Any:
1374 raise StopAtLayerException(tensor)
1376 # Stop at the beginning of the specified block, not at the end of the previous block
1377 block_hook_name = f"blocks.{stop_at_layer}.hook_in"
1378 hook_dict = self.hook_dict
1379 if block_hook_name in hook_dict: 1379 ↛ 1384line 1379 didn't jump to line 1384 because the condition on line 1379 was always true
1380 add_hook_to_point(
1381 hook_dict[block_hook_name], stop_hook, block_hook_name, "fwd"
1382 )
1384 apply_hooks(fwd_hooks, True)
1385 apply_hooks(bwd_hooks, False)
1386 if start_at_layer is not None:
1387 kwargs["start_at_layer"] = start_at_layer
1388 try:
1389 output = self.forward(
1390 input, return_type=return_type, stop_at_layer=stop_at_layer, **kwargs
1391 )
1392 except StopAtLayerException as e:
1393 output = e.layer_output
1394 return output
1395 finally:
1396 try:
1397 for hook_point, direction in added_hooks:
1398 if reset_hooks_end:
1399 # `level` keeps this to the hooks added above — hooks the caller
1400 # attached before the call survive.
1401 hook_point.remove_hooks(dir=direction, level=context_level)
1402 if clear_contexts:
1403 hook_point.clear_context()
1404 finally:
1405 self.context_level -= 1
1407 # ---- high-level execution: run_with_cache ----
1409 @overload
1410 def run_with_cache(
1411 self,
1412 input: Union[str, List[str], torch.Tensor],
1413 return_cache_object: Literal[True] = True,
1414 remove_batch_dim: bool = False,
1415 **kwargs,
1416 ) -> Tuple[Any, ActivationCache]:
1417 """Run with cache - placeholder implementation."""
1418 pass
1420 @overload
1421 def run_with_cache(
1422 self,
1423 input: Union[str, List[str], torch.Tensor],
1424 return_cache_object: Literal[False],
1425 remove_batch_dim: bool = False,
1426 **kwargs,
1427 ) -> Tuple[Any, Dict[str, torch.Tensor]]:
1428 """Run with cache - placeholder implementation."""
1429 pass
1431 def run_with_cache(
1432 self,
1433 input: Union[str, List[str], torch.Tensor],
1434 return_cache_object: bool = True,
1435 remove_batch_dim: bool = False,
1436 names_filter: Optional[Union[str, List[str], Callable[[str], bool]]] = None,
1437 stop_at_layer: Optional[int] = None,
1438 start_at_layer: Optional[int] = None,
1439 pos_slice: Optional[Union[Slice, SliceInput]] = None,
1440 incl_bwd: bool = False,
1441 reset_hooks_end: bool = True,
1442 clear_contexts: bool = False,
1443 **kwargs,
1444 ) -> Tuple[Any, Union[ActivationCache, Dict[str, torch.Tensor]]]:
1445 """Run the model and cache activations. Returns ``(output, cache)``.
1447 ``stop_at_layer`` raises :class:`StopAtLayerException` to stop early.
1448 ``start_at_layer`` treats ``input`` as the residual entering block ``k``
1449 (see :meth:`forward`); blocks below ``k`` are excluded from the cache to
1450 match HookedTransformer. ``pos_slice`` slices each cached activation along
1451 its position dimension (dim 1 for resid/per-head/token-id activations; the
1452 query position ``-2`` for attention patterns/scores; ``HookPoint.pos_dim``
1453 for hooks fired on another layout). ``incl_bwd`` also caches gradients
1454 under ``"<name>_grad"`` by running ``output.backward()``;
1455 the caller must request a scalar output (``return_type="loss"``) and the
1456 model must be on the gradients-capable transformers driver.
1457 ``reset_hooks_end`` removes the hooks this call added when it finishes —
1458 hooks the caller attached beforehand are left alone either way;
1459 ``clear_contexts`` also wipes the touched hook points' ``ctx``. ``device``
1460 offloads cached activations (matches ``ActivationCache.to``); the model and
1461 inputs stay where the caller put them.
1462 """
1463 if incl_bwd and stop_at_layer is not None:
1464 raise ValueError(
1465 "incl_bwd=True cannot be combined with stop_at_layer: the run returns the "
1466 "intermediate activation at that layer, not a scalar to call backward() on."
1467 )
1468 if incl_bwd and not torch.is_grad_enabled():
1469 raise ValueError(
1470 "incl_bwd=True needs autograd, but gradient tracking is off "
1471 "(torch.no_grad(), set_grad_enabled(False) or inference mode). "
1472 "Run the call with gradients enabled."
1473 )
1474 pos_slice_obj = Slice.unwrap(pos_slice)
1475 aliases = build_alias_to_canonical_map(self.hook_dict)
1477 def create_names_filter_fn(filter_input):
1478 if filter_input is None:
1479 return lambda name: True
1480 elif isinstance(filter_input, str):
1481 mapped_name = aliases.get(filter_input, None)
1482 if mapped_name:
1483 return lambda name: name == mapped_name or name == filter_input
1484 else:
1485 return lambda name: name == filter_input
1486 elif isinstance(filter_input, list):
1487 mapped_list = []
1488 for item in filter_input:
1489 mapped_list.append(item)
1490 mapped_name = aliases.get(item, None)
1491 if mapped_name:
1492 mapped_list.append(mapped_name)
1493 return lambda name: name in mapped_list
1494 elif callable(filter_input): 1494 ↛ 1497line 1494 didn't jump to line 1497 because the condition on line 1494 was always true
1495 return filter_input
1496 else:
1497 raise ValueError("names_filter must be a string, list of strings, or callable")
1499 names_filter_fn = create_names_filter_fn(names_filter)
1500 if isinstance(names_filter, (str, list)):
1501 requested = [names_filter] if isinstance(names_filter, str) else names_filter
1502 for name in requested:
1503 self._check_hook_fireable(name, aliases.get(name, name))
1504 cache: Dict[str, torch.Tensor] = {}
1505 hooks: List[Tuple[HookPoint, str]] = []
1506 visited: set[int] = set()
1507 stop_hook_point: Optional[HookPoint] = None
1508 stop_hook_fn: Optional[Callable] = None
1510 # None → no-op .to(None), tensors stay on their current device.
1511 cache_device = kwargs.pop("device", None)
1513 def _store(name: str, value: torch.Tensor, hook: Any, suffix: str = "") -> None:
1514 stored = value.detach().to(cache_device)
1515 if pos_slice_obj is not None and stored.dim() >= 2: 1515 ↛ 1519line 1515 didn't jump to line 1519 because the condition on line 1515 was always true
1516 # A gradient has its activation's layout, so the axis comes from
1517 # the unsuffixed name — see _pos_slice_dim.
1518 stored = pos_slice_obj.apply(stored, dim=self._pos_slice_dim(name, hook))
1519 cache[name + suffix] = stored
1521 def make_cache_hook(name: str):
1522 def cache_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
1523 if tensor is None: 1523 ↛ 1524line 1523 didn't jump to line 1524 because the condition on line 1523 was never true
1524 cache[name] = None
1525 elif isinstance(tensor, torch.Tensor): 1525 ↛ 1527line 1525 didn't jump to line 1527 because the condition on line 1525 was always true
1526 _store(name, tensor, hook)
1527 elif isinstance(tensor, tuple):
1528 if len(tensor) > 0 and isinstance(tensor[0], torch.Tensor):
1529 _store(name, tensor[0], hook)
1530 else:
1531 pass
1532 else:
1533 try:
1534 if hasattr(tensor, "detach"):
1535 _store(name, tensor, hook)
1536 except:
1537 pass
1538 return tensor
1540 return cache_hook
1542 hook_dict = self.hook_dict
1543 effective_stop_layer = None
1544 if stop_at_layer is not None and hasattr(self, "blocks"):
1545 if stop_at_layer < 0:
1546 effective_stop_layer = len(self.blocks) + stop_at_layer
1547 else:
1548 effective_stop_layer = stop_at_layer
1549 effective_start_layer = None
1550 if start_at_layer is not None and hasattr(self, "blocks"):
1551 effective_start_layer = (
1552 len(self.blocks) + start_at_layer if start_at_layer < 0 else start_at_layer
1553 )
1554 matched_any = False
1555 for hook_name, hook in hook_dict.items():
1556 if names_filter_fn(hook_name):
1557 matched_any = True
1558 if effective_start_layer is not None and self._is_embedding_stage_hook(hook_name):
1559 continue
1560 if hook_name.startswith("blocks."):
1561 try:
1562 layer_num = int(hook_name.split(".")[1])
1563 except (IndexError, ValueError):
1564 layer_num = None
1565 if layer_num is not None: 1565 ↛ 1573line 1565 didn't jump to line 1573 because the condition on line 1565 was always true
1566 # stop/start bound the executed range; blocks outside it
1567 # either don't run (stop) or run on discarded input (start),
1568 # so their activations must not enter the cache.
1569 if effective_stop_layer is not None and layer_num >= effective_stop_layer:
1570 continue
1571 if effective_start_layer is not None and layer_num < effective_start_layer:
1572 continue
1573 hooks.append((hook, hook_name))
1574 # Explicit string/list filters matching nothing must not return (logits, {}) silently.
1575 if not matched_any and names_filter and isinstance(names_filter, (str, list)): 1575 ↛ 1581line 1575 didn't jump to line 1581 because the condition on line 1575 was always true
1576 raise KeyError(
1577 f"names_filter {names_filter!r} matched no hook points on this model; "
1578 "check the name against model.hook_dict (this backend may not serve it)."
1579 )
1581 def make_grad_cache_hook(name: str):
1582 def grad_hook(tensor: torch.Tensor, *, hook: Any) -> None:
1583 if isinstance(tensor, torch.Tensor): 1583 ↛ 1586line 1583 didn't jump to line 1586 because the condition on line 1583 was always true
1584 _store(name, tensor, hook, suffix="_grad")
1585 # A non-None return from a backward hook replaces grad_input; stay read-only.
1586 return None
1588 return grad_hook
1590 processed_args = [input]
1591 # Driver-aware input placement: torch drivers move input_ids to the model's
1592 # device; remote drivers (no local parameters) leave them as-is.
1593 target_device = self._input_device()
1594 if processed_args and isinstance(processed_args[0], str):
1595 assert self.tokenizer is not None, "Tokenizer must be set to pass string input."
1596 prepend_bos = kwargs.pop("prepend_bos", None)
1597 input_ids = self.to_tokens(processed_args[0], prepend_bos=prepend_bos)
1598 if target_device is not None: 1598 ↛ 1600line 1598 didn't jump to line 1600 because the condition on line 1598 was always true
1599 input_ids = input_ids.to(target_device)
1600 kwargs["input_ids"] = input_ids
1601 processed_args = processed_args[1:]
1602 elif "input" in kwargs and isinstance(kwargs["input"], str): 1602 ↛ 1603line 1602 didn't jump to line 1603 because the condition on line 1602 was never true
1603 assert self.tokenizer is not None, "Tokenizer must be set to pass string input."
1604 prepend_bos = kwargs.pop("prepend_bos", None)
1605 input_ids = self.to_tokens(kwargs["input"], prepend_bos=prepend_bos)
1606 if target_device is not None:
1607 input_ids = input_ids.to(target_device)
1608 kwargs["input_ids"] = input_ids
1609 del kwargs["input"]
1610 if stop_at_layer is not None and hasattr(self, "blocks"):
1611 if stop_at_layer < 0:
1612 stop_at_layer = len(self.blocks) + stop_at_layer
1613 last_layer_to_process = stop_at_layer - 1
1615 def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
1616 raise StopAtLayerException(tensor)
1618 if stop_at_layer >= 0 and stop_at_layer < len(self.blocks): 1618 ↛ 1625line 1618 didn't jump to line 1625 because the condition on line 1618 was always true
1619 # Stop at the beginning of the specified block, not at the end of the previous block
1620 block_hook_name = f"blocks.{stop_at_layer}.hook_in"
1621 hook_dict = self.hook_dict
1622 if block_hook_name in hook_dict: 1622 ↛ 1625line 1622 didn't jump to line 1625 because the condition on line 1622 was always true
1623 stop_hook_point = hook_dict[block_hook_name]
1624 stop_hook_fn = stop_hook
1625 filtered_kwargs = kwargs.copy()
1626 # ``cache_device`` is honored by ``make_cache_hook`` above (``tensor.detach().to(cache_device)``);
1627 # the model and inputs stay where the caller put them, matching ``ActivationCache.to``.
1628 if cache_device is not None and getattr(self.cfg, "n_devices", 1) > 1:
1629 # Moving a dispatched model to a single device collapses accelerate's
1630 # split and breaks its routing hooks. The cache will stay spread across
1631 # the per-layer devices; callers can .to(cache_device) on cache entries
1632 # after the fact if they need a single-device cache.
1633 warnings.warn(
1634 f"run_with_cache(device={cache_device!r}) ignored: model is dispatched "
1635 f"across {self.cfg.n_devices} devices via device_map. Cached activations "
1636 "will remain on their per-layer devices.",
1637 stacklevel=2,
1638 )
1639 if start_at_layer is not None:
1640 filtered_kwargs["start_at_layer"] = start_at_layer
1641 # Only validate gated hooks when the caller explicitly supplied a
1642 # names_filter. The default filter matches every hook and must not
1643 # cause gated hooks to be treated as explicitly requested.
1644 gated_names_skipped: List[str] = []
1645 if names_filter is not None:
1646 kept = []
1647 for hp, name in hooks:
1648 if self._gated_hook_reason(hp.name or name) is not None:
1649 gated_names_skipped.append(name)
1650 else:
1651 kept.append((hp, name))
1652 hooks = kept
1653 if gated_names_skipped:
1654 warnings.warn(
1655 f"run_with_cache: skipped {len(gated_names_skipped)} gated-off hook name(s) "
1656 f"that will never be cached: {gated_names_skipped}. Call the relevant "
1657 "set_use_*(True) setter first to enable them.",
1658 stacklevel=2,
1659 )
1660 context_level = getattr(self, "context_level", 0) + 1
1661 self.context_level = context_level
1662 try:
1663 for hp, name in hooks:
1664 hp.add_hook(make_cache_hook(name), level=context_level)
1665 if incl_bwd:
1666 hp.add_hook(make_grad_cache_hook(name), dir="bwd", level=context_level)
1667 if stop_hook_point is not None and stop_hook_fn is not None:
1668 stop_hook_point.add_hook(stop_hook_fn, level=context_level)
1669 except Exception:
1670 try:
1671 self.remove_all_hook_fns(level=context_level)
1672 finally:
1673 self.context_level -= 1
1674 raise
1675 try:
1676 if (
1677 "output_attentions" not in filtered_kwargs
1678 and self.adapter.supports_hf_output_attentions
1679 ):
1680 # Attention-free remote-code models (e.g. HyenaDNA) reject the
1681 # kwarg outright; only pass it when the forward actually accepts
1682 # it. Non-torch drivers expose no local module to introspect —
1683 # keep the kwarg for them, as before.
1684 underlying = getattr(self._driver, "underlying_model", None)
1685 if underlying is None:
1686 filtered_kwargs["output_attentions"] = True
1687 else:
1688 fwd_params = inspect.signature(underlying.forward).parameters
1689 if "output_attentions" in fwd_params or any(
1690 p.kind is inspect.Parameter.VAR_KEYWORD for p in fwd_params.values()
1691 ):
1692 filtered_kwargs["output_attentions"] = True
1693 if processed_args:
1694 output = self.forward(processed_args[0], **filtered_kwargs)
1695 elif "input_ids" in filtered_kwargs: 1695 ↛ 1701line 1695 didn't jump to line 1701 because the condition on line 1695 was always true
1696 output = self.forward(
1697 filtered_kwargs["input_ids"],
1698 **{k: v for k, v in filtered_kwargs.items() if k != "input_ids"},
1699 )
1700 else:
1701 output = self.forward(**filtered_kwargs)
1702 if hasattr(output, "logits"): 1702 ↛ 1703line 1702 didn't jump to line 1703 because the condition on line 1702 was never true
1703 output = output.logits
1704 if incl_bwd:
1705 # Gradients land in the cache via the bwd hooks, which the finally below
1706 # removes, so the backward pass has to happen inside this try.
1707 if not isinstance(output, torch.Tensor) or output.numel() != 1:
1708 shape = tuple(output.shape) if isinstance(output, torch.Tensor) else None
1709 raise ValueError(
1710 "incl_bwd=True needs a scalar output to call backward() on, got "
1711 f"{type(output).__name__}{f' of shape {shape}' if shape else ''}. "
1712 'Pass return_type="loss".'
1713 )
1714 if not output.requires_grad:
1715 raise ValueError(
1716 "incl_bwd=True got an output with no grad_fn — the model's parameters "
1717 "have requires_grad=False, so there is nothing to differentiate."
1718 )
1719 output.backward()
1720 except StopAtLayerException as e:
1721 output = e.layer_output
1722 except Exception as e:
1723 raise e
1724 finally:
1725 try:
1726 if reset_hooks_end:
1727 # `level` keeps this to the hooks added above — hooks the caller
1728 # attached before the call survive.
1729 self.remove_all_hook_fns(level=context_level)
1730 if clear_contexts:
1731 for hp, _ in hooks:
1732 hp.clear_context()
1733 if stop_hook_point is not None: 1733 ↛ 1734line 1733 didn't jump to line 1734 because the condition on line 1733 was never true
1734 stop_hook_point.clear_context()
1735 finally:
1736 self.context_level -= 1
1737 if self.compatibility_mode == True:
1738 reverse_aliases = {}
1739 for old_name, new_name in aliases.items():
1740 if isinstance(new_name, list): 1740 ↛ 1741line 1740 didn't jump to line 1741 because the condition on line 1740 was never true
1741 for single_new_name in new_name:
1742 reverse_aliases[single_new_name] = old_name
1743 else:
1744 reverse_aliases[new_name] = old_name
1745 # Gradient entries are keyed "<hook_name>_grad", so alias lookups run on the
1746 # base name and the suffix is re-attached to the aliased key.
1747 suffixes = ("", "_grad") if incl_bwd else ("",)
1748 cache_items_to_add = {}
1749 for cache_name, cached_value in cache.items():
1750 base_name, suffix = (
1751 (cache_name[: -len("_grad")], "_grad")
1752 if cache_name.endswith("_grad")
1753 else (cache_name, "")
1754 )
1755 old_name = reverse_aliases.get(base_name)
1756 if old_name is not None:
1757 cache_items_to_add[old_name + suffix] = cached_value
1758 cache.update(cache_items_to_add)
1759 for alias_name, target_name in aliases.items():
1760 targets = target_name if isinstance(target_name, list) else [target_name]
1761 for suffix in suffixes:
1762 if alias_name + suffix in cache:
1763 continue
1764 for single_target in targets:
1765 if single_target + suffix in cache: 1765 ↛ 1766line 1765 didn't jump to line 1766 because the condition on line 1765 was never true
1766 cache[alias_name + suffix] = cache[single_target + suffix]
1767 break
1768 # The input states the batch size; inferring it from the cached shapes goes wrong
1769 # when the filter selects only flattened or position-indexed hooks.
1770 batch_size = _input_batch_size(
1771 processed_args[0] if processed_args else None,
1772 filtered_kwargs.get("input_ids"),
1773 filtered_kwargs.get("inputs_embeds"),
1774 )
1775 if return_cache_object:
1776 activation_cache = ActivationCache(cache, self, has_batch_dim=True)
1777 if remove_batch_dim:
1778 activation_cache.remove_batch_dim(batch_size=batch_size)
1779 return (output, activation_cache)
1780 else:
1781 if remove_batch_dim:
1782 # Same batch-size check and squeeze rule as the ActivationCache path.
1783 tensors = {k: v for k, v in cache.items() if isinstance(v, torch.Tensor)}
1784 cache.update(
1785 ActivationCache(tensors, self)
1786 .remove_batch_dim(batch_size=batch_size)
1787 .cache_dict
1788 )
1789 return (output, cache)