Coverage for transformer_lens/HookedRootModule.py: 63%
179 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""HookedRootModule.
3Base class extending :class:`torch.nn.Module` with hook-based introspection
4utilities: wrap any module's tensors in :class:`HookPoint` wrappers, call ``setup()``,
5and ``run_with_hooks`` / ``run_with_cache`` work on it. This is permanent
6infrastructure — it survives the 4.0 removal of the legacy model classes and is
7the supported way to add TransformerLens-style hooks to arbitrary ``nn.Module`` objects.
8Lives in its own module so that downstream code (e.g. :class:`ActivationCache`)
9can type-hint against it without the broader ``hook_points`` import surface.
10"""
12from __future__ import annotations
14import logging
15from collections.abc import Callable, Iterable
16from contextlib import contextmanager
17from functools import partial
18from typing import Any, Literal, Optional, Union, cast
20import torch
21import torch.nn as nn
22from torch import Tensor
24from transformer_lens.hook_points import (
25 DeviceType,
26 HookFunction,
27 HookIntrospectionMixin,
28 HookPoint,
29 NamesFilter,
30)
31from transformer_lens.utilities import Slice, SliceInput
32from transformer_lens.utilities import remove_batch_dim as drop_batch_dim
33from transformer_lens.utilities import warn_if_mps
36class HookedRootModule(HookIntrospectionMixin, nn.Module):
37 """A class building on nn.Module to interface nicely with HookPoints.
39 Adds various nice utilities, most notably run_with_hooks to run the model with temporary hooks,
40 and run_with_cache to run the model on some input and return a cache of all activations.
42 Notes:
44 The main footgun with PyTorch hooking is that hooks are GLOBAL state. If you add a hook to the
45 module, and then run it a bunch of times, the hooks persist. If you debug a broken hook and add
46 the fixed version, the broken one is still there. To solve this, run_with_hooks will remove
47 hooks at the end by default, and I recommend using the API of this and run_with_cache. If you
48 want to add hooks into global state, I recommend being intentional about this, and I recommend
49 using reset_hooks liberally in your code to remove any accidentally remaining global state.
51 The main time this goes wrong is when you want to use backward hooks (to cache or intervene on
52 gradients). In this case, you need to keep the hooks around as global state until you've run
53 loss.backward() (and so need to disable the reset_hooks_end flag on run_with_hooks)
54 """
56 name: Optional[str]
57 mod_dict: dict[str, nn.Module]
58 hook_dict: dict[str, HookPoint]
60 def __init__(self, *args: Any):
61 super().__init__()
62 self.is_caching = False
63 self.context_level = 0
65 def setup(self):
66 """
67 Sets up model.
69 This function must be called in the model's `__init__` method AFTER defining all layers. It
70 adds a parameter to each module containing its name, and builds a dictionary mapping module
71 names to the module instances. It also initializes a hook dictionary for modules of type
72 "HookPoint".
73 """
74 self.mod_dict = {}
75 self.hook_dict = {}
76 for name, module in self.named_modules():
77 if name == "":
78 continue
79 module.name = name
80 self.mod_dict[name] = module
81 # TODO: is the bottom line the same as "if "HookPoint" in str(type(module)):"
82 if isinstance(module, HookPoint): 82 ↛ 76line 82 didn't jump to line 76 because the condition on line 82 was always true
83 self.hook_dict[name] = module
85 def hook_points(self):
86 return self.hook_dict.values()
88 def remove_all_hook_fns(
89 self,
90 direction: Literal["fwd", "bwd", "both"] = "both",
91 including_permanent: bool = False,
92 level: Optional[int] = None,
93 ):
94 for hp in self.hook_points():
95 hp.remove_hooks(direction, including_permanent=including_permanent, level=level)
97 def clear_contexts(self):
98 for hp in self.hook_points():
99 hp.clear_context()
101 def reset_hooks(
102 self,
103 clear_contexts: bool = True,
104 direction: Literal["fwd", "bwd", "both"] = "both",
105 including_permanent: bool = False,
106 level: Optional[int] = None,
107 ):
108 if clear_contexts:
109 self.clear_contexts()
110 self.remove_all_hook_fns(direction, including_permanent, level=level)
111 self.is_caching = False
113 def check_and_add_hook(
114 self,
115 hook_point: HookPoint,
116 hook_point_name: str,
117 hook: HookFunction,
118 dir: Literal["fwd", "bwd"] = "fwd",
119 is_permanent: bool = False,
120 level: Optional[int] = None,
121 prepend: bool = False,
122 ) -> None:
123 """Runs checks on the hook, and then adds it to the hook point"""
125 self.check_hooks_to_add(
126 hook_point,
127 hook_point_name,
128 hook,
129 dir=dir,
130 is_permanent=is_permanent,
131 prepend=prepend,
132 )
133 hook_point.add_hook(hook, dir=dir, is_permanent=is_permanent, level=level, prepend=prepend)
135 def check_hooks_to_add(
136 self,
137 hook_point: HookPoint,
138 hook_point_name: str,
139 hook: HookFunction,
140 dir: Literal["fwd", "bwd"] = "fwd",
141 is_permanent: bool = False,
142 prepend: bool = False,
143 ) -> None:
144 """Override this function to add checks on which hooks should be added"""
145 pass
147 def add_hook(
148 self,
149 name: Union[str, Callable[[str], bool]],
150 hook: HookFunction,
151 dir: Literal["fwd", "bwd"] = "fwd",
152 is_permanent: bool = False,
153 level: Optional[int] = None,
154 prepend: bool = False,
155 ) -> None:
156 if isinstance(name, str):
157 hook_point = self.mod_dict[name]
158 assert isinstance(
159 hook_point, HookPoint
160 ) # TODO does adding assert meaningfully slow down performance? I've added them for type checking purposes.
161 self.check_and_add_hook(
162 hook_point,
163 name,
164 hook,
165 dir=dir,
166 is_permanent=is_permanent,
167 level=level,
168 prepend=prepend,
169 )
170 else:
171 # Otherwise, name is a Boolean function on names
172 for hook_point_name, hp in self.hook_dict.items():
173 if name(hook_point_name):
174 self.check_and_add_hook(
175 hp,
176 hook_point_name,
177 hook,
178 dir=dir,
179 is_permanent=is_permanent,
180 level=level,
181 prepend=prepend,
182 )
184 def add_perma_hook(
185 self,
186 name: Union[str, Callable[[str], bool]],
187 hook: HookFunction,
188 dir: Literal["fwd", "bwd"] = "fwd",
189 ) -> None:
190 self.add_hook(name, hook, dir=dir, is_permanent=True)
192 def _enable_hook_with_name(self, name: str, hook: Callable, dir: Literal["fwd", "bwd"]):
193 """This function takes a key for the mod_dict and enables the related hook for that module
195 Args:
196 name (str): The module name
197 hook (Callable): The hook to add
198 dir (Literal["fwd", "bwd"]): The direction for the hook
199 """
200 hook_point_module = self.mod_dict[name]
201 if not hasattr(hook_point_module, "add_hook"): 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 raise TypeError(f"Expected a module with add_hook, got {type(hook_point_module)}")
203 if isinstance(hook_point_module, torch.Tensor): 203 ↛ 204line 203 didn't jump to line 204 because the condition on line 203 was never true
204 raise TypeError(
205 "Module set as Tensor for some reason!"
206 ) # mypy seems to think these could be tensors after a torch update no idea why, or if this is possible
207 module_with_hook = cast(HookPoint, hook_point_module)
208 module_with_hook.add_hook(hook, dir=dir, level=self.context_level)
210 def _enable_hooks_for_points(
211 self,
212 hook_points: Iterable[tuple[str, HookPoint]],
213 enabled: Callable,
214 hook: Callable,
215 dir: Literal["fwd", "bwd"],
216 ):
217 """Enables hooks for a list of points
219 Args:
220 hook_points (Dict[str, HookPoint]): The hook points
221 enabled (Callable): _description_
222 hook (Callable): _description_
223 dir (Literal["fwd", "bwd"]): _description_
224 """
225 for hook_name, hook_point in hook_points:
226 if enabled(hook_name):
227 hook_point.add_hook(hook, dir=dir, level=self.context_level)
229 def _enable_hook(self, name: Union[str, Callable], hook: Callable, dir: Literal["fwd", "bwd"]):
230 """Enables an individual hook on a hook point
232 Args:
233 name (str): The name of the hook
234 hook (Callable): The actual hook
235 dir (Literal["fwd", "bwd"], optional): The direction of the hook. Defaults to "fwd".
236 """
237 if isinstance(name, str): 237 ↛ 240line 237 didn't jump to line 240 because the condition on line 237 was always true
238 self._enable_hook_with_name(name=name, hook=hook, dir=dir)
239 else:
240 self._enable_hooks_for_points(
241 hook_points=self.hook_dict.items(), enabled=name, hook=hook, dir=dir
242 )
244 @contextmanager
245 def hooks(
246 self,
247 fwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
248 bwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
249 reset_hooks_end: bool = True,
250 clear_contexts: bool = False,
251 ):
252 """
253 A context manager for adding temporary hooks to the model.
255 Args:
256 fwd_hooks: List[Tuple[name, hook]], where name is either the name of a hook point or a
257 Boolean function on hook names and hook is the function to add to that hook point.
258 bwd_hooks: Same as fwd_hooks, but for the backward pass.
259 reset_hooks_end (bool): If True, removes all hooks added by this context manager when the context manager exits.
260 clear_contexts (bool): If True, clears hook contexts whenever hooks are reset.
262 Example:
264 .. code-block:: python
266 with model.hooks(fwd_hooks=my_hooks):
267 hooked_loss = model(text, return_type="loss")
268 """
269 try:
270 self.context_level += 1
272 for name, hook in fwd_hooks:
273 self._enable_hook(name=name, hook=hook, dir="fwd")
274 for name, hook in bwd_hooks: 274 ↛ 275line 274 didn't jump to line 275 because the loop on line 274 never started
275 self._enable_hook(name=name, hook=hook, dir="bwd")
276 yield self
277 finally:
278 if reset_hooks_end: 278 ↛ 282line 278 didn't jump to line 282 because the condition on line 278 was always true
279 self.reset_hooks(
280 clear_contexts, including_permanent=False, level=self.context_level
281 )
282 self.context_level -= 1
284 def run_with_hooks(
285 self,
286 *model_args: Any, # TODO: unsure about whether or not this Any typing is correct or not; may need to be replaced with something more specific?
287 fwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
288 bwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
289 reset_hooks_end: bool = True,
290 clear_contexts: bool = False,
291 **model_kwargs: Any,
292 ):
293 """
294 Runs the model with specified forward and backward hooks.
296 Args:
297 fwd_hooks (List[Tuple[Union[str, Callable], Callable]]): A list of (name, hook), where name is
298 either the name of a hook point or a boolean function on hook names, and hook is the
299 function to add to that hook point. Hooks with names that evaluate to True are added
300 respectively.
301 bwd_hooks (List[Tuple[Union[str, Callable], Callable]]): Same as fwd_hooks, but for the
302 backward pass.
303 reset_hooks_end (bool): If True, all hooks are removed at the end, including those added
304 during this run. Default is True.
305 clear_contexts (bool): If True, clears hook contexts whenever hooks are reset. Default is
306 False.
307 *model_args: Positional arguments for the model.
308 **model_kwargs: Keyword arguments for the model's forward function. See your related
309 models forward pass for details as to what sort of arguments you can pass through.
311 Note:
312 If you want to use backward hooks, set `reset_hooks_end` to False, so the backward hooks
313 remain active. This function only runs a forward pass.
314 """
315 if len(bwd_hooks) > 0 and reset_hooks_end:
316 logging.warning(
317 "WARNING: Hooks will be reset at the end of run_with_hooks. This removes the backward hooks before a backward pass can occur."
318 )
320 with self.hooks(fwd_hooks, bwd_hooks, reset_hooks_end, clear_contexts) as hooked_model:
321 return hooked_model.forward(*model_args, **model_kwargs)
323 def add_caching_hooks(
324 self,
325 names_filter: NamesFilter = None,
326 incl_bwd: bool = False,
327 device: DeviceType = None, # TODO: unsure about whether or not this device typing is correct or not?
328 remove_batch_dim: bool = False,
329 cache: Optional[dict] = None,
330 ) -> dict:
331 """Adds hooks to the model to cache activations. Note: It does NOT actually run the model to get activations, that must be done separately.
333 Args:
334 names_filter (NamesFilter, optional): Which activations to cache. Can be a list of strings (hook names) or a filter function mapping hook names to booleans. Defaults to lambda name: True.
335 incl_bwd (bool, optional): Whether to also do backwards hooks. Defaults to False.
336 device (_type_, optional): The device to store on. Defaults to same device as model.
337 remove_batch_dim (bool, optional): Whether to remove the batch dimension (only works for batch_size==1). A leading dimension that isn't size 1 is left in place, so a larger batch keeps its batch dimension. Defaults to False.
338 cache (Optional[dict], optional): The cache to store activations in, a new dict is created by default. Defaults to None.
340 Returns:
341 cache (dict): The cache where activations will be stored.
342 """
343 if device is not None: 343 ↛ 344line 343 didn't jump to line 344 because the condition on line 343 was never true
344 warn_if_mps(device)
345 if cache is None: 345 ↛ 348line 345 didn't jump to line 348 because the condition on line 345 was always true
346 cache = {}
348 if names_filter is None: 348 ↛ 350line 348 didn't jump to line 350 because the condition on line 348 was always true
349 names_filter = lambda name: True
350 elif isinstance(names_filter, str):
351 filter_str = names_filter
352 names_filter = lambda name: name == filter_str
353 elif isinstance(names_filter, list):
354 filter_list = names_filter
355 names_filter = lambda name: name in filter_list
357 assert callable(names_filter), "names_filter must be a callable"
359 self.is_caching = True
361 def save_hook(tensor: Tensor, hook: HookPoint, is_backward: bool):
362 assert hook.name is not None
363 hook_name = hook.name
364 if is_backward: 364 ↛ 365line 364 didn't jump to line 365 because the condition on line 364 was never true
365 hook_name += "_grad"
366 stored = tensor.detach().to(device)
367 cache[hook_name] = drop_batch_dim(stored) if remove_batch_dim else stored
369 for name, hp in self.hook_dict.items():
370 if names_filter(name): 370 ↛ 369line 370 didn't jump to line 369 because the condition on line 370 was always true
371 hp.add_hook(partial(save_hook, is_backward=False), "fwd")
372 if incl_bwd: 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true
373 hp.add_hook(partial(save_hook, is_backward=True), "bwd")
374 return cache
376 def run_with_cache(
377 self,
378 *model_args: Any,
379 names_filter: NamesFilter = None,
380 device: DeviceType = None,
381 remove_batch_dim: bool = False,
382 incl_bwd: bool = False,
383 reset_hooks_end: bool = True,
384 clear_contexts: bool = False,
385 pos_slice: Optional[Union[Slice, SliceInput]] = None,
386 **model_kwargs: Any,
387 ):
388 """
389 Runs the model and returns the model output and a Cache object.
391 Args:
392 *model_args: Positional arguments for the model.
393 names_filter (NamesFilter, optional): A filter for which activations to cache. Accepts None, str,
394 list of str, or a function that takes a string and returns a bool. Defaults to None, which
395 means cache everything.
396 device (str or torch.Device, optional): The device to cache activations on. Defaults to the
397 model device. WARNING: Setting a different device than the one used by the model leads to
398 significant performance degradation.
399 remove_batch_dim (bool, optional): If True, removes the batch dimension when caching. Only
400 makes sense with batch_size=1 inputs: an AssertionError is raised when the first
401 positional input has a larger batch. Defaults to False.
402 incl_bwd (bool, optional): If True, calls backward on the model output and caches gradients
403 as well. Assumes that the model outputs a scalar (e.g., return_type="loss"). Custom loss
404 functions are not supported. Defaults to False.
405 reset_hooks_end (bool, optional): If True, removes all hooks added by this function at the
406 end of the run. Defaults to True.
407 clear_contexts (bool, optional): If True, clears hook contexts whenever hooks are reset.
408 Defaults to False.
409 pos_slice:
410 The slice to apply to the cache output. Defaults to None, do nothing.
411 **model_kwargs: Keyword arguments for the model's forward function. See your related
412 models forward pass for details as to what sort of arguments you can pass through.
414 Returns:
415 tuple: A tuple containing the model output and a Cache object.
417 """
419 if remove_batch_dim and model_args and isinstance(model_args[0], Tensor): 419 ↛ 425line 419 didn't jump to line 425 because the condition on line 419 was always true
420 batch_size = model_args[0].shape[0] if model_args[0].ndim > 0 else 1
421 assert (
422 batch_size == 1
423 ), f"Cannot remove batch dimension from cache with batch size {batch_size}"
425 pos_slice = Slice.unwrap(pos_slice)
427 cache_dict, fwd, bwd = self.get_caching_hooks(
428 names_filter,
429 incl_bwd,
430 device,
431 remove_batch_dim=remove_batch_dim,
432 pos_slice=pos_slice,
433 )
435 with self.hooks(
436 fwd_hooks=fwd,
437 bwd_hooks=bwd,
438 reset_hooks_end=reset_hooks_end,
439 clear_contexts=clear_contexts,
440 ):
441 model_out = self(*model_args, **model_kwargs)
442 if incl_bwd: 442 ↛ 443line 442 didn't jump to line 443 because the condition on line 442 was never true
443 model_out.backward()
445 return model_out, cache_dict
447 def get_caching_hooks(
448 self,
449 names_filter: NamesFilter = None,
450 incl_bwd: bool = False,
451 device: DeviceType = None,
452 remove_batch_dim: bool = False,
453 cache: Optional[dict] = None,
454 pos_slice: Optional[Union[Slice, SliceInput]] = None,
455 ) -> tuple[dict, list, list]:
456 """Creates hooks to cache activations. Note: It does not add the hooks to the model.
458 Args:
459 names_filter (NamesFilter, optional): Which activations to cache. Can be a list of strings (hook names) or a filter function mapping hook names to booleans. Defaults to lambda name: True.
460 incl_bwd (bool, optional): Whether to also do backwards hooks. Defaults to False.
461 device (_type_, optional): The device to store on. Keeps on the same device as the layer if None.
462 remove_batch_dim (bool, optional): Whether to remove the batch dimension (only works for batch_size==1). A leading dimension that isn't size 1 is left in place, so a larger batch keeps its batch dimension. Defaults to False.
463 cache (Optional[dict], optional): The cache to store activations in, a new dict is created by default. Defaults to None.
465 Returns:
466 cache (dict): The cache where activations will be stored.
467 fwd_hooks (list): The forward hooks.
468 bwd_hooks (list): The backward hooks. Empty if incl_bwd is False.
469 """
470 if device is not None: 470 ↛ 471line 470 didn't jump to line 471 because the condition on line 470 was never true
471 warn_if_mps(device)
472 if cache is None: 472 ↛ 475line 472 didn't jump to line 475 because the condition on line 472 was always true
473 cache = {}
475 pos_slice = Slice.unwrap(pos_slice)
477 if names_filter is None: 477 ↛ 479line 477 didn't jump to line 479 because the condition on line 477 was always true
478 names_filter = lambda name: True
479 elif isinstance(names_filter, str):
480 filter_str = names_filter
481 names_filter = lambda name: name == filter_str
482 elif isinstance(names_filter, list):
483 filter_list = names_filter
484 names_filter = lambda name: name in filter_list
485 elif callable(names_filter):
486 names_filter = names_filter
487 else:
488 raise ValueError("names_filter must be a string, list of strings, or function")
489 assert callable(names_filter) # Callable[[str], bool]
491 self.is_caching = True
493 def save_hook(tensor: Tensor, hook: HookPoint, is_backward: bool = False):
494 # for attention heads the pos dimension is the third from last
495 if hook.name is None: 495 ↛ 496line 495 didn't jump to line 496 because the condition on line 495 was never true
496 raise RuntimeError("Hook should have been provided a name")
498 hook_name = hook.name
499 if is_backward: 499 ↛ 500line 499 didn't jump to line 500 because the condition on line 499 was never true
500 hook_name += "_grad"
501 resid_stream = tensor.detach().to(device)
502 if remove_batch_dim: 502 ↛ 505line 502 didn't jump to line 505 because the condition on line 502 was always true
503 resid_stream = drop_batch_dim(resid_stream)
505 if ( 505 ↛ 512line 505 didn't jump to line 512 because the condition on line 505 was never true
506 hook.name.endswith("hook_q")
507 or hook.name.endswith("hook_k")
508 or hook.name.endswith("hook_v")
509 or hook.name.endswith("hook_z")
510 or hook.name.endswith("hook_result")
511 ):
512 pos_dim = -3
513 else:
514 # for all other components the pos dimension is the second from last
515 # including the attn scores where the dest token is the second from last
516 pos_dim = -2
518 if ( 518 ↛ 521line 518 didn't jump to line 521 because the condition on line 518 was never true
519 resid_stream.dim() >= -pos_dim
520 ): # check if the residual stream has a pos dimension before trying to slice
521 resid_stream = pos_slice.apply(resid_stream, dim=pos_dim)
522 cache[hook_name] = resid_stream
524 fwd_hooks = []
525 bwd_hooks = []
526 for name, _ in self.hook_dict.items():
527 if names_filter(name): 527 ↛ 526line 527 didn't jump to line 526 because the condition on line 527 was always true
528 fwd_hooks.append((name, partial(save_hook, is_backward=False)))
529 if incl_bwd: 529 ↛ 530line 529 didn't jump to line 530 because the condition on line 529 was never true
530 bwd_hooks.append((name, partial(save_hook, is_backward=True)))
532 return cache, fwd_hooks, bwd_hooks
534 def cache_all(
535 self,
536 cache: Optional[dict],
537 incl_bwd: bool = False,
538 device: DeviceType = None,
539 remove_batch_dim: bool = False,
540 ):
541 logging.warning(
542 "cache_all is deprecated and will eventually be removed, use add_caching_hooks or run_with_cache"
543 )
544 self.add_caching_hooks(
545 names_filter=lambda name: True,
546 cache=cache,
547 incl_bwd=incl_bwd,
548 device=device,
549 remove_batch_dim=remove_batch_dim,
550 )
552 def cache_some(
553 self,
554 cache: Optional[dict],
555 names: Callable[[str], bool],
556 incl_bwd: bool = False,
557 device: DeviceType = None,
558 remove_batch_dim: bool = False,
559 ):
560 """Cache a list of hook provided by names, Boolean function on names"""
561 logging.warning(
562 "cache_some is deprecated and will eventually be removed, use add_caching_hooks or run_with_cache"
563 )
564 self.add_caching_hooks(
565 names_filter=names,
566 cache=cache,
567 incl_bwd=incl_bwd,
568 device=device,
569 remove_batch_dim=remove_batch_dim,
570 )