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