tracer¶
Cache(modules: Optional[List[Union[Envoy, str]]] = None, device: Optional[device] = device('cpu'), dtype: Optional[dtype] = None, detach: Optional[bool] = True, include_output: bool = True, include_inputs: bool = False, rename: Optional[Dict[str, str]] = None, alias: Optional[Dict[str, str]] = None)
A cache for storing module activations during tracing.
Persistent cache hooks (registered via :func:hooks.cache_output_hook and
:func:hooks.cache_input_hook) fire on every forward pass and call
:meth:add to record values. Hooks are registered by
:meth:InterleavingTracer.cache when the user creates a cache, and
removed automatically when the interleaver exits.
The cache applies optional transformations (detach, device, dtype) to
values before storing them. Hook handles live on the owning
:class:Mediator (in mediator.hooks) and are drained by
:meth:Mediator.remove_hooks, not on the cache itself.
| PARAMETER | DESCRIPTION |
|---|---|
modules
|
Optional list of modules (Envoy objects or path strings)
to cache. If
TYPE:
|
device
|
Device to move cached tensors to (default: CPU).
TYPE:
|
dtype
|
Optional dtype to convert cached tensors to.
TYPE:
|
detach
|
Whether to detach tensors from the computation graph.
TYPE:
|
include_output
|
Whether to cache module outputs.
TYPE:
|
include_inputs
|
Whether to cache module inputs.
TYPE:
|
rename
|
Rename mapping for alias path resolution.
TYPE:
|
alias
|
Alias mapping for
TYPE:
|
dataclass
¶
Entry(output: Optional[Any] = None, inputs: Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]] = None)
CacheDict(data: Union[CacheDict, Dict[str, Entry]], path: str = '', alias: Dict[str, str] = dict(), rename: Dict[str, str] = dict(), alias_paths: Dict[str, str] = dict())
Bases: Dict
A dictionary subclass that provides convenient access to cached module activations.
This class extends the standard dictionary to provide both dictionary-style access and attribute-style access to cached activations. It supports hierarchical access to nested modules using dot notation and indexing for module lists.
Examples:
Access cached activations using dictionary keys:
Access using attribute notation:
Access module outputs and inputs:
>>> cache.model.transformer.h[0].output
>>> cache.model.transformer.h[0].inputs
>>> cache.model.transformer.h[0].input # First input argument
The class maintains an internal path that tracks the current location in the module hierarchy, allowing for intuitive navigation through nested modules.
Add a value to the cache with optional transformations.
Called by persistent cache hooks registered via
:func:hooks.cache_output_hook and :func:hooks.cache_input_hook.
| PARAMETER | DESCRIPTION |
|---|---|
module_path
|
The module's envoy path (e.g.
TYPE:
|
key
|
TYPE:
|
value
|
The tensor value to store.
TYPE:
|
Bases: Tracer
Tracer that manages the interleaving of model execution and interventions.
This class coordinates the execution of the model's forward pass and user-defined intervention functions through the Interleaver.
| PARAMETER | DESCRIPTION |
|---|---|
fn
|
The function to execute (typically the model's forward pass)
TYPE:
|
model
|
The model envoy to intervene on
TYPE:
|
*args
|
Additional arguments to pass to the function
DEFAULT:
|
**kwargs
|
Additional keyword arguments to pass to the function
DEFAULT:
|
Compile the captured code block into a callable function.
| RETURNS | DESCRIPTION |
|---|---|
Callable
|
A callable function that executes the captured code block |
First executes the parent Tracer's execute method to set up the context, then creates an Interleaver to manage the interventions during model execution.
Create an Invoker to capture and execute an intervention function.
| PARAMETER | DESCRIPTION |
|---|---|
*args
|
Additional arguments to pass to the intervention function
DEFAULT:
|
**kwargs
|
Additional keyword arguments to pass to the intervention function
DEFAULT:
|
| RETURNS | DESCRIPTION |
|---|---|
|
An Invoker instance |
cache(modules: Optional[List[Union[Envoy, str]]] = None, device: Optional[device] = device('cpu'), dtype: Optional[dtype] = None, detach: Optional[bool] = True, include_output: bool = True, include_inputs: bool = False) -> Union[Dict, Object]
Create a cache that records module activations during execution.
Registers persistent hooks on the target modules (all modules if
modules is None, otherwise only the specified subset). The
hooks fire after any intervention hooks (mediator_idx=inf) so
they capture post-intervention values. Hooks persist across
generation steps and are automatically removed when the interleaver
exits.
| PARAMETER | DESCRIPTION |
|---|---|
modules
|
Modules to cache — Envoy objects or path strings.
If
TYPE:
|
device
|
Device to move cached tensors to (default: CPU).
TYPE:
|
dtype
|
Optional dtype to convert cached tensors to.
TYPE:
|
detach
|
Whether to detach tensors from the computation graph.
TYPE:
|
include_output
|
Whether to cache module outputs.
TYPE:
|
include_inputs
|
Whether to cache module inputs.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
A
|
class:
TYPE:
|
nnsight barrier: A synchronization primitive for coordinating multiple concurrent invocations in nnsight.
This works similarly to a threading.Barrier, but is designed for use with nnsight's model tracing and intervention system. A barrier allows you to pause execution in multiple parallel invocations until all participants have reached the barrier, at which point all are released to continue. This is useful when you want to synchronize the execution of different model runs, for example to ensure that all have reached a certain point (such as after embedding lookup) before proceeding to the next stage (such as generation or intervention).
Example usage:
with gpt2.generate(max_new_tokens=3) as tracer:
barrier = tracer.barrier(2)
with tracer.invoke(MSG_prompt):
embeddings = gpt2.transformer.wte.output
barrier()
output1 = gpt2.generator.output.save()
with tracer.invoke("_ _ _ _ _ _ _ _ _"):
barrier()
gpt2.transformer.wte.output = embeddings
output2 = gpt2.generator.output.save()
In this example, both invocations will pause at the barrier until both have reached it, ensuring synchronization.
Bases: InterleavingTracer
A tracer that runs the model in fake tensor mode to validate operations and inspect tensor shapes.
This tracer uses PyTorch's FakeTensorMode to run the model without actual computation, allowing for shape validation and operation checking.
Execute the model in fake tensor mode.
This method: 1. Runs the model in fake tensor mode to validate operations 2. Allows intervention code inside the scan context to access fake tensor shapes
| PARAMETER | DESCRIPTION |
|---|---|
fn
|
The function to execute (typically the model's forward pass)
TYPE:
|