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