Coverage for transformer_lens/HookedRootModule.py: 15%
175 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +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, warn_if_mps
34class HookedRootModule(HookIntrospectionMixin, nn.Module):
35 """A class building on nn.Module to interface nicely with HookPoints.
37 Adds various nice utilities, most notably run_with_hooks to run the model with temporary hooks,
38 and run_with_cache to run the model on some input and return a cache of all activations.
40 Notes:
42 The main footgun with PyTorch hooking is that hooks are GLOBAL state. If you add a hook to the
43 module, and then run it a bunch of times, the hooks persist. If you debug a broken hook and add
44 the fixed version, the broken one is still there. To solve this, run_with_hooks will remove
45 hooks at the end by default, and I recommend using the API of this and run_with_cache. If you
46 want to add hooks into global state, I recommend being intentional about this, and I recommend
47 using reset_hooks liberally in your code to remove any accidentally remaining global state.
49 The main time this goes wrong is when you want to use backward hooks (to cache or intervene on
50 gradients). In this case, you need to keep the hooks around as global state until you've run
51 loss.backward() (and so need to disable the reset_hooks_end flag on run_with_hooks)
52 """
54 name: Optional[str]
55 mod_dict: dict[str, nn.Module]
56 hook_dict: dict[str, HookPoint]
58 def __init__(self, *args: Any):
59 super().__init__()
60 self.is_caching = False
61 self.context_level = 0
63 def setup(self):
64 """
65 Sets up model.
67 This function must be called in the model's `__init__` method AFTER defining all layers. It
68 adds a parameter to each module containing its name, and builds a dictionary mapping module
69 names to the module instances. It also initializes a hook dictionary for modules of type
70 "HookPoint".
71 """
72 self.mod_dict = {}
73 self.hook_dict = {}
74 for name, module in self.named_modules():
75 if name == "":
76 continue
77 module.name = name
78 self.mod_dict[name] = module
79 # TODO: is the bottom line the same as "if "HookPoint" in str(type(module)):"
80 if isinstance(module, HookPoint):
81 self.hook_dict[name] = module
83 def hook_points(self):
84 return self.hook_dict.values()
86 def remove_all_hook_fns(
87 self,
88 direction: Literal["fwd", "bwd", "both"] = "both",
89 including_permanent: bool = False,
90 level: Optional[int] = None,
91 ):
92 for hp in self.hook_points():
93 hp.remove_hooks(direction, including_permanent=including_permanent, level=level)
95 def clear_contexts(self):
96 for hp in self.hook_points():
97 hp.clear_context()
99 def reset_hooks(
100 self,
101 clear_contexts: bool = True,
102 direction: Literal["fwd", "bwd", "both"] = "both",
103 including_permanent: bool = False,
104 level: Optional[int] = None,
105 ):
106 if clear_contexts:
107 self.clear_contexts()
108 self.remove_all_hook_fns(direction, including_permanent, level=level)
109 self.is_caching = False
111 def check_and_add_hook(
112 self,
113 hook_point: HookPoint,
114 hook_point_name: str,
115 hook: HookFunction,
116 dir: Literal["fwd", "bwd"] = "fwd",
117 is_permanent: bool = False,
118 level: Optional[int] = None,
119 prepend: bool = False,
120 ) -> None:
121 """Runs checks on the hook, and then adds it to the hook point"""
123 self.check_hooks_to_add(
124 hook_point,
125 hook_point_name,
126 hook,
127 dir=dir,
128 is_permanent=is_permanent,
129 prepend=prepend,
130 )
131 hook_point.add_hook(hook, dir=dir, is_permanent=is_permanent, level=level, prepend=prepend)
133 def check_hooks_to_add(
134 self,
135 hook_point: HookPoint,
136 hook_point_name: str,
137 hook: HookFunction,
138 dir: Literal["fwd", "bwd"] = "fwd",
139 is_permanent: bool = False,
140 prepend: bool = False,
141 ) -> None:
142 """Override this function to add checks on which hooks should be added"""
143 pass
145 def add_hook(
146 self,
147 name: Union[str, Callable[[str], bool]],
148 hook: HookFunction,
149 dir: Literal["fwd", "bwd"] = "fwd",
150 is_permanent: bool = False,
151 level: Optional[int] = None,
152 prepend: bool = False,
153 ) -> None:
154 if isinstance(name, str):
155 hook_point = self.mod_dict[name]
156 assert isinstance(
157 hook_point, HookPoint
158 ) # TODO does adding assert meaningfully slow down performance? I've added them for type checking purposes.
159 self.check_and_add_hook(
160 hook_point,
161 name,
162 hook,
163 dir=dir,
164 is_permanent=is_permanent,
165 level=level,
166 prepend=prepend,
167 )
168 else:
169 # Otherwise, name is a Boolean function on names
170 for hook_point_name, hp in self.hook_dict.items():
171 if name(hook_point_name):
172 self.check_and_add_hook(
173 hp,
174 hook_point_name,
175 hook,
176 dir=dir,
177 is_permanent=is_permanent,
178 level=level,
179 prepend=prepend,
180 )
182 def add_perma_hook(
183 self,
184 name: Union[str, Callable[[str], bool]],
185 hook: HookFunction,
186 dir: Literal["fwd", "bwd"] = "fwd",
187 ) -> None:
188 self.add_hook(name, hook, dir=dir, is_permanent=True)
190 def _enable_hook_with_name(self, name: str, hook: Callable, dir: Literal["fwd", "bwd"]):
191 """This function takes a key for the mod_dict and enables the related hook for that module
193 Args:
194 name (str): The module name
195 hook (Callable): The hook to add
196 dir (Literal["fwd", "bwd"]): The direction for the hook
197 """
198 hook_point_module = self.mod_dict[name]
199 if not hasattr(hook_point_module, "add_hook"):
200 raise TypeError(f"Expected a module with add_hook, got {type(hook_point_module)}")
201 if isinstance(hook_point_module, torch.Tensor):
202 raise TypeError(
203 "Module set as Tensor for some reason!"
204 ) # mypy seems to think these could be tensors after a torch update no idea why, or if this is possible
205 module_with_hook = cast(HookPoint, hook_point_module)
206 module_with_hook.add_hook(hook, dir=dir, level=self.context_level)
208 def _enable_hooks_for_points(
209 self,
210 hook_points: Iterable[tuple[str, HookPoint]],
211 enabled: Callable,
212 hook: Callable,
213 dir: Literal["fwd", "bwd"],
214 ):
215 """Enables hooks for a list of points
217 Args:
218 hook_points (Dict[str, HookPoint]): The hook points
219 enabled (Callable): _description_
220 hook (Callable): _description_
221 dir (Literal["fwd", "bwd"]): _description_
222 """
223 for hook_name, hook_point in hook_points:
224 if enabled(hook_name):
225 hook_point.add_hook(hook, dir=dir, level=self.context_level)
227 def _enable_hook(self, name: Union[str, Callable], hook: Callable, dir: Literal["fwd", "bwd"]):
228 """Enables an individual hook on a hook point
230 Args:
231 name (str): The name of the hook
232 hook (Callable): The actual hook
233 dir (Literal["fwd", "bwd"], optional): The direction of the hook. Defaults to "fwd".
234 """
235 if isinstance(name, str):
236 self._enable_hook_with_name(name=name, hook=hook, dir=dir)
237 else:
238 self._enable_hooks_for_points(
239 hook_points=self.hook_dict.items(), enabled=name, hook=hook, dir=dir
240 )
242 @contextmanager
243 def hooks(
244 self,
245 fwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
246 bwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
247 reset_hooks_end: bool = True,
248 clear_contexts: bool = False,
249 ):
250 """
251 A context manager for adding temporary hooks to the model.
253 Args:
254 fwd_hooks: List[Tuple[name, hook]], where name is either the name of a hook point or a
255 Boolean function on hook names and hook is the function to add to that hook point.
256 bwd_hooks: Same as fwd_hooks, but for the backward pass.
257 reset_hooks_end (bool): If True, removes all hooks added by this context manager when the context manager exits.
258 clear_contexts (bool): If True, clears hook contexts whenever hooks are reset.
260 Example:
262 .. code-block:: python
264 with model.hooks(fwd_hooks=my_hooks):
265 hooked_loss = model(text, return_type="loss")
266 """
267 try:
268 self.context_level += 1
270 for name, hook in fwd_hooks:
271 self._enable_hook(name=name, hook=hook, dir="fwd")
272 for name, hook in bwd_hooks:
273 self._enable_hook(name=name, hook=hook, dir="bwd")
274 yield self
275 finally:
276 if reset_hooks_end:
277 self.reset_hooks(
278 clear_contexts, including_permanent=False, level=self.context_level
279 )
280 self.context_level -= 1
282 def run_with_hooks(
283 self,
284 *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?
285 fwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
286 bwd_hooks: list[tuple[Union[str, Callable], Callable]] = [],
287 reset_hooks_end: bool = True,
288 clear_contexts: bool = False,
289 **model_kwargs: Any,
290 ):
291 """
292 Runs the model with specified forward and backward hooks.
294 Args:
295 fwd_hooks (List[Tuple[Union[str, Callable], Callable]]): A list of (name, hook), where name is
296 either the name of a hook point or a boolean function on hook names, and hook is the
297 function to add to that hook point. Hooks with names that evaluate to True are added
298 respectively.
299 bwd_hooks (List[Tuple[Union[str, Callable], Callable]]): Same as fwd_hooks, but for the
300 backward pass.
301 reset_hooks_end (bool): If True, all hooks are removed at the end, including those added
302 during this run. Default is True.
303 clear_contexts (bool): If True, clears hook contexts whenever hooks are reset. Default is
304 False.
305 *model_args: Positional arguments for the model.
306 **model_kwargs: Keyword arguments for the model's forward function. See your related
307 models forward pass for details as to what sort of arguments you can pass through.
309 Note:
310 If you want to use backward hooks, set `reset_hooks_end` to False, so the backward hooks
311 remain active. This function only runs a forward pass.
312 """
313 if len(bwd_hooks) > 0 and reset_hooks_end:
314 logging.warning(
315 "WARNING: Hooks will be reset at the end of run_with_hooks. This removes the backward hooks before a backward pass can occur."
316 )
318 with self.hooks(fwd_hooks, bwd_hooks, reset_hooks_end, clear_contexts) as hooked_model:
319 return hooked_model.forward(*model_args, **model_kwargs)
321 def add_caching_hooks(
322 self,
323 names_filter: NamesFilter = None,
324 incl_bwd: bool = False,
325 device: DeviceType = None, # TODO: unsure about whether or not this device typing is correct or not?
326 remove_batch_dim: bool = False,
327 cache: Optional[dict] = None,
328 ) -> dict:
329 """Adds hooks to the model to cache activations. Note: It does NOT actually run the model to get activations, that must be done separately.
331 Args:
332 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.
333 incl_bwd (bool, optional): Whether to also do backwards hooks. Defaults to False.
334 device (_type_, optional): The device to store on. Defaults to same device as model.
335 remove_batch_dim (bool, optional): Whether to remove the batch dimension (only works for batch_size==1). Defaults to False.
336 cache (Optional[dict], optional): The cache to store activations in, a new dict is created by default. Defaults to None.
338 Returns:
339 cache (dict): The cache where activations will be stored.
340 """
341 if device is not None:
342 warn_if_mps(device)
343 if cache is None:
344 cache = {}
346 if names_filter is None:
347 names_filter = lambda name: True
348 elif isinstance(names_filter, str):
349 filter_str = names_filter
350 names_filter = lambda name: name == filter_str
351 elif isinstance(names_filter, list):
352 filter_list = names_filter
353 names_filter = lambda name: name in filter_list
355 assert callable(names_filter), "names_filter must be a callable"
357 self.is_caching = True
359 def save_hook(tensor: Tensor, hook: HookPoint, is_backward: bool):
360 assert hook.name is not None
361 hook_name = hook.name
362 if is_backward:
363 hook_name += "_grad"
364 if remove_batch_dim:
365 cache[hook_name] = tensor.detach().to(device)[0]
366 else:
367 cache[hook_name] = tensor.detach().to(device)
369 for name, hp in self.hook_dict.items():
370 if names_filter(name):
371 hp.add_hook(partial(save_hook, is_backward=False), "fwd")
372 if incl_bwd:
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. Defaults to False.
401 incl_bwd (bool, optional): If True, calls backward on the model output and caches gradients
402 as well. Assumes that the model outputs a scalar (e.g., return_type="loss"). Custom loss
403 functions are not supported. Defaults to False.
404 reset_hooks_end (bool, optional): If True, removes all hooks added by this function at the
405 end of the run. Defaults to True.
406 clear_contexts (bool, optional): If True, clears hook contexts whenever hooks are reset.
407 Defaults to False.
408 pos_slice:
409 The slice to apply to the cache output. Defaults to None, do nothing.
410 **model_kwargs: Keyword arguments for the model's forward function. See your related
411 models forward pass for details as to what sort of arguments you can pass through.
413 Returns:
414 tuple: A tuple containing the model output and a Cache object.
416 """
418 pos_slice = Slice.unwrap(pos_slice)
420 cache_dict, fwd, bwd = self.get_caching_hooks(
421 names_filter,
422 incl_bwd,
423 device,
424 remove_batch_dim=remove_batch_dim,
425 pos_slice=pos_slice,
426 )
428 with self.hooks(
429 fwd_hooks=fwd,
430 bwd_hooks=bwd,
431 reset_hooks_end=reset_hooks_end,
432 clear_contexts=clear_contexts,
433 ):
434 model_out = self(*model_args, **model_kwargs)
435 if incl_bwd:
436 model_out.backward()
438 return model_out, cache_dict
440 def get_caching_hooks(
441 self,
442 names_filter: NamesFilter = None,
443 incl_bwd: bool = False,
444 device: DeviceType = None,
445 remove_batch_dim: bool = False,
446 cache: Optional[dict] = None,
447 pos_slice: Optional[Union[Slice, SliceInput]] = None,
448 ) -> tuple[dict, list, list]:
449 """Creates hooks to cache activations. Note: It does not add the hooks to the model.
451 Args:
452 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.
453 incl_bwd (bool, optional): Whether to also do backwards hooks. Defaults to False.
454 device (_type_, optional): The device to store on. Keeps on the same device as the layer if None.
455 remove_batch_dim (bool, optional): Whether to remove the batch dimension (only works for batch_size==1). Defaults to False.
456 cache (Optional[dict], optional): The cache to store activations in, a new dict is created by default. Defaults to None.
458 Returns:
459 cache (dict): The cache where activations will be stored.
460 fwd_hooks (list): The forward hooks.
461 bwd_hooks (list): The backward hooks. Empty if incl_bwd is False.
462 """
463 if device is not None:
464 warn_if_mps(device)
465 if cache is None:
466 cache = {}
468 pos_slice = Slice.unwrap(pos_slice)
470 if names_filter is None:
471 names_filter = lambda name: True
472 elif isinstance(names_filter, str):
473 filter_str = names_filter
474 names_filter = lambda name: name == filter_str
475 elif isinstance(names_filter, list):
476 filter_list = names_filter
477 names_filter = lambda name: name in filter_list
478 elif callable(names_filter):
479 names_filter = names_filter
480 else:
481 raise ValueError("names_filter must be a string, list of strings, or function")
482 assert callable(names_filter) # Callable[[str], bool]
484 self.is_caching = True
486 def save_hook(tensor: Tensor, hook: HookPoint, is_backward: bool = False):
487 # for attention heads the pos dimension is the third from last
488 if hook.name is None:
489 raise RuntimeError("Hook should have been provided a name")
491 hook_name = hook.name
492 if is_backward:
493 hook_name += "_grad"
494 resid_stream = tensor.detach().to(device)
495 if remove_batch_dim:
496 resid_stream = resid_stream[0]
498 if (
499 hook.name.endswith("hook_q")
500 or hook.name.endswith("hook_k")
501 or hook.name.endswith("hook_v")
502 or hook.name.endswith("hook_z")
503 or hook.name.endswith("hook_result")
504 ):
505 pos_dim = -3
506 else:
507 # for all other components the pos dimension is the second from last
508 # including the attn scores where the dest token is the second from last
509 pos_dim = -2
511 if (
512 tensor.dim() >= -pos_dim
513 ): # check if the residual stream has a pos dimension before trying to slice
514 resid_stream = pos_slice.apply(resid_stream, dim=pos_dim)
515 cache[hook_name] = resid_stream
517 fwd_hooks = []
518 bwd_hooks = []
519 for name, _ in self.hook_dict.items():
520 if names_filter(name):
521 fwd_hooks.append((name, partial(save_hook, is_backward=False)))
522 if incl_bwd:
523 bwd_hooks.append((name, partial(save_hook, is_backward=True)))
525 return cache, fwd_hooks, bwd_hooks
527 def cache_all(
528 self,
529 cache: Optional[dict],
530 incl_bwd: bool = False,
531 device: DeviceType = None,
532 remove_batch_dim: bool = False,
533 ):
534 logging.warning(
535 "cache_all is deprecated and will eventually be removed, use add_caching_hooks or run_with_cache"
536 )
537 self.add_caching_hooks(
538 names_filter=lambda name: True,
539 cache=cache,
540 incl_bwd=incl_bwd,
541 device=device,
542 remove_batch_dim=remove_batch_dim,
543 )
545 def cache_some(
546 self,
547 cache: Optional[dict],
548 names: Callable[[str], bool],
549 incl_bwd: bool = False,
550 device: DeviceType = None,
551 remove_batch_dim: bool = False,
552 ):
553 """Cache a list of hook provided by names, Boolean function on names"""
554 logging.warning(
555 "cache_some is deprecated and will eventually be removed, use add_caching_hooks or run_with_cache"
556 )
557 self.add_caching_hooks(
558 names_filter=names,
559 cache=cache,
560 incl_bwd=incl_bwd,
561 device=device,
562 remove_batch_dim=remove_batch_dim,
563 )