nnsight.intervention.tracing

Tracing module for capturing and executing code blocks within ‘with’ statements.

This module provides the core Tracer class that enables dynamic code capture, compilation, and execution. It allows intercepting Python code blocks and executing them in controlled environments with access to the original context.

exception nnsight.intervention.tracing.base.ExitTracingException

Exception raised to exit the tracing process.

This exception is used as a control flow mechanism to cleanly exit a with block without executing the code inside it. When the tracer detects that execution has reached the traced code block, it raises this exception to prevent normal execution and instead execute the code through the tracing backend.

class nnsight.intervention.tracing.base.Tracer(*args, backend: Backend = None, _info: Info = None, **kwargs)

Captures and executes code within a tracing context.

This class provides a sophisticated mechanism for intercepting Python code blocks within ‘with’ statements, capturing their source code, and executing them in controlled environments. It’s designed to enable dynamic code manipulation and execution with full access to the original context and variables.

The tracing process follows four main steps: 1. Capture: Finds and extracts the source code from the with block 2. Parse: Uses AST to identify the exact code boundaries within the block 3. Compile: Wraps the source code in a function definition for execution 4. Execute: Runs the compiled function with appropriate context

Example

```python with Tracer() as tracer:

# This code block will be captured and executed later x = some_computation() result = process(x)

```

The tracer supports various execution environments including: - Regular Python files - IPython/Jupyter notebooks - Interactive consoles - Nested tracing contexts

args

Arguments passed to the tracer for use during execution

kwargs

Keyword arguments passed to the tracer

backend

Backend implementation that handles the actual code execution

info

Info object containing captured code metadata and context

class Info(source: List[str], frame: FrameType, start_line: int, node: With, filename: str = None, cache_key: int = None)

Container for metadata about the traced code block.

This nested class stores all necessary information about the code being traced, including the source code, execution frame context, and AST node information. It serves as a data transfer object between the capture, parse, and execution phases.

source

Source code lines from the traced block

Type:

List[str]

frame

Frame information from the call stack where tracing occurred

Type:

FrameType

start_line

Line number where the traced code block begins

Type:

int

node

AST node representing the with block

Type:

ast.With

filename

Filename for the traced code (real file or generated identifier)

Type:

str

copy()

Create a deep copy of this Info instance.

Returns:

A new Info instance with the same metadata

capture()

Capture the code block within the ‘with’ statement.

This is step 1 of the tracing process. It walks up the call stack to find the frame outside of nnsight, extracts the source code of the ‘with’ block, and prepares it for later execution. The method handles various execution environments including regular files, IPython notebooks, and interactive consoles.

The capture process: 1. Identifies the execution context (file, notebook, console, nested trace) 2. Extracts source code lines from the appropriate source 3. Handles indentation normalization for proper AST parsing 4. Calls parse() to identify the exact with block boundaries 5. Stores all metadata in a Tracer.Info object

Raises:

ValueError – If no source code can be found for the current context

compile() → Callable

Compile the captured source code into a callable function.

This is step 3 of the tracing process. Takes the captured and parsed source code and wraps it in a function definition with the necessary context parameters. The resulting function can be executed with proper variable scoping and access to the original execution environment.

The compiled function signature is: __nnsight_tracer_{id}__(__nnsight_tracer__, __nnsight_tracing_info__)

The function includes: - A call to tracer.pull() to import variables from the original scope - The original traced code block - A call to tracer.push() to export variables back to the original scope

Returns:

A callable function that executes the captured code block with proper context

execute(fn: Callable)

Execute the compiled function with proper context.

This is step 4 of the tracing process. Runs the compiled function that was created in the compile() step, passing in the tracer instance and info object as context. This allows the traced code to access the original variables and execution environment.

Parameters:

fn – The compiled function to execute (created by compile() method)

parse(source_lines: List[str], start_line: int)

Parse the source code to extract the with block contents.

This is step 2 of the tracing process. Uses the Abstract Syntax Tree (AST) to identify the exact boundaries of the with block and extract only the code that should be traced and executed later.

Parameters:

• source_lines – List of source code lines to parse

• start_line – Line number where the tracer creation statement begins

Returns:

• start_line (int): Adjusted start line of the with block body

• source_lines (List[str]): Extracted source lines from the with block

• node (ast.With): AST node representing the with block

Return type:

Tuple containing

Raises:

WithBlockNotFoundError – If no with block is found at the specified line

pull()

Pull variables from the original execution frame into the current context.

This method imports variables from the original scope where the tracer was created into the current traced code execution context. This ensures that the traced code has access to all variables that were available when the tracer was instantiated.

This is the opposite operation of push() and is called at the beginning of traced code execution.

push(state: Dict = None)

Push local variables back to the original execution frame.

This method exports variables from the traced code execution back to the original scope where the tracer was created. This allows changes made during tracing to persist and affect the original execution environment.

The method handles variable filtering to only push non-nnsight variables, and includes special logic for nested tracing contexts using Globals.stack.

Parameters:

state – Dictionary of variable names and values to push to the frame. If None, automatically collects variables from the current execution frame.

exception nnsight.intervention.tracing.base.WithBlockNotFoundError

Exception raised when a with block is not found in the source code.

This exception indicates that the AST parser could not locate a with block at the expected line number during the code capture process. This typically occurs when there are issues with source code parsing or line number mapping.

class nnsight.intervention.tracing.tracer.Cache(modules: List[Any | str] | None = None, device: device | None = device(type='cpu'), dtype: dtype | None = None, detach: bool | None = True, include_output: bool = True, include_inputs: bool = False, rename: Dict[str, str] | None = None, alias: Dict[str, str] | None = None)

A cache for storing and transforming tensor values during tracing.

This class provides functionality to store tensor values with optional transformations such as detaching from computation graph, moving to a specific device, or converting to a specific dtype.

class CacheDict(data: CacheDict | Dict[str, Entry], path: str = '', alias: Dict[str, str] = {}, rename: Dict[str, str] = {}, alias_paths: Dict[str, str] = {})

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: >>> cache[‘model.transformer.h.0.attn’]

Access using attribute notation: >>> cache.model.transformer.h[0].attn

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.

property input

Returns the input property from the Cache.Entry at the current path.

property inputs

Returns the inputs attribute from the Cache.Entry at the current path.

keys() → a set-like object providing a view on D's keys

property output

Returns the output attribute from the Cache.Entry at the current path.

class Entry(output: Any | None = None, inputs: Tuple[Tuple[Any, ...], Dict[str, Any]] | None = None)

property input

Gets the first positional argument of the inputs value to the cached module. Returns None if no inputs were cached.

add(provider: str, value: Any)

Add a value to the cache with optional transformations.

Parameters:

• provider – The key to store the value under

• value – The tensor value to store

class nnsight.intervention.tracing.tracer.InterleavingTracer(fn: Callable, model: Any, *args, backend: Backend = None, **kwargs)

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.

barrier(n_participants: int)

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.

cache(modules: List[Any | str] | None = None, device: device | None = device(type='cpu'), dtype: dtype | None = None, detach: bool | None = True, include_output: bool = True, include_inputs: bool = False) → Dict | Object

Get or create a cache for storing intermediate values during tracing.

Parameters:

• modules – Optional list of modules to cache, defaults to all modules

• device – Optional device to move tensors to, defaults to cpu

• dtype – Optional dtype to convert tensors to, defaults to None

• detach – Whether to detach tensors from computation graph, defaults to True

• include_output – Whether to include output in the cached activations

• include_inputs – Whether to include inputs in the cached activations

Returns:

A dictionary containing the cached values

capture()

Capture the code block within the ‘with’ statement.

compile() → Callable

Compile the captured code block into a callable function.

Returns:

A callable function that executes the captured code block

execute(fn: Callable)

First executes the parent Tracer’s execute method to set up the context, then creates an Interleaver to manage the interventions during model execution.

get_frame()

Get the frame of the tracer.

invoke(*args, **kwargs)

Create an Invoker to capture and execute an intervention function.

Parameters:

• *args – Additional arguments to pass to the intervention function

• **kwargs – Additional keyword arguments to pass to the intervention function

Returns:

An Invoker instance

property result: Object

Get the result of the method being traced.

This property allows access to the return values produced by the method being traced.

Example

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)
>>> with model.generate("Hello World") as tracer:
>>>     result = tracer.result.save()
>>> print(result)

Returns:

The result of the method being traced

stop()

Raise an EarlyStopException to stop the execution of the model.

class nnsight.intervention.tracing.tracer.ScanningTracer(fn: Callable, model: Any, *args, backend: Backend = None, **kwargs)

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. It populates the _fake_inputs and _fake_output attributes on each Envoy to store the shapes and types of tensors that would flow through the model during a real forward pass.

execute(fn: Callable)

Execute the model in fake tensor mode.

This method: 1. Registers forward hooks on all modules to capture fake input/output 2. Runs the model in fake tensor mode to validate operations 3. Stores the fake inputs/outputs on each Envoy for later inspection

Parameters:

fn – The function to execute (typically the model’s forward pass)

class nnsight.intervention.tracing.invoker.Invoker(tracer: Any, *args, **kwargs)

Extends the Tracer class to invoke intervention functions.

This class captures code blocks and compiles them into intervention functions that can be executed by the Interleaver.

compile()

Compile the captured code block into an intervention function.

The function is wrapped with try-catch logic to handle exceptions and signal completion to the mediator.

Returns:

A callable intervention function

execute(fn: Callable)

Execute the compiled intervention function.

Creates a new Mediator for the intervention function and adds it to the parent tracer’s mediators list.

Parameters:

fn – The compiled intervention function

exception nnsight.intervention.tracing.util.ExceptionWrapper(info: Tracer.Info, original: Exception, *args, **kwargs)

Wrapper for exceptions that provides additional details for tracer created code.

This class helps provide better error messages by including source code context and proper line numbers from the original code being traced.

set_info(info: Tracer.Info)

Updates the tracer information and recalculates line offsets.

Parameters:

info – New tracer information to use

nnsight.intervention.tracing.util.get_dependencies(fn: Callable)

Extracts global dependencies used by a function.

Parameters:

fn – The function to analyze for dependencies

Returns:

Dictionary mapping names to their corresponding global objects used by the function

nnsight.intervention.tracing.util.indent(source: List[str], indent: int = 1)

Indents each line in the source list by a specified number of indentation levels.

Parameters:

• source – List of strings to indent

• indent – Number of indentation levels to apply (default: 1)

Returns:

List of indented strings

nnsight.intervention.tracing.util.try_catch(source: List[str], exception_source: List[str] = ['raise\n'], else_source: List[str] = ['pass\n'], finally_source: List[str] = ['pass\n'])

Wraps source code in a try-except-else-finally block.

Args:

source: The code to be wrapped in the try block exception_source: Code for the except block (default: [“raise

“])

else_source: Code for the else block (default: [“pass

“])

finally_source: Code for the finally block (default: [“pass

“])

Returns:

List of strings representing the complete try-catch block, properly indented

nnsight.intervention.tracing.util.wrap_exception(exception: Exception, info: Tracer.Info)

Wraps an exception with additional context from the tracer.

This function either updates an existing ExceptionWrapper or creates a new dynamically-typed exception class that inherits from both the original exception type and ExceptionWrapper.

Parameters:

• exception – The exception to wrap

• info – Tracer information containing context about where the exception occurred

Returns:

A wrapped exception with enhanced traceback information

class nnsight.intervention.tracing.backwards.BackwardsMediator(intervention: Callable, info: Tracer.Info, name: str | None = None, batch_group: List[int] | None = None, stop: int | None = None)

request(requester: Any)

Request a value from a specific provider.

Parameters:

requester (str) – The identifier of the provider to request a value from

Returns:

The requested value

Return type:

Any

class nnsight.intervention.tracing.backwards.BackwardsTracer(tensor: Tensor, fn: Callable, *args, **kwargs)

execute(fn: Callable)

Execute the compiled intervention function.

Creates a new Mediator for the intervention function and adds it to the parent tracer’s mediators list.

Parameters:

fn – The compiled intervention function

nnsight.intervention.tracing.backwards.wrap_grad(interleaver: Interleaver)

Create a hook for gradient intervention.

Returns:

A function that can be used to intercept gradients

class nnsight.intervention.tracing.editing.EditingTracer(*args, backend: ~nnsight.intervention.backends.base.Backend = <nnsight.intervention.backends.editing.EditingBackend object>, inplace: bool = False, **kwargs)

class nnsight.intervention.tracing.iterator.IteratorTracer(iteration: int | slice, interleaver: Interleaver)

compile()

Compile the captured source code as a callable function.

Wraps the captured code in a function definition that accepts the necessary context parameters for execution.

Returns:

A callable function that executes the captured code block

execute(fn: Callable)

Execute the compiled function with proper context.

This is step 4 of the tracing process. Runs the compiled function that was created in the compile() step, passing in the tracer instance and info object as context. This allows the traced code to access the original variables and execution environment.

Parameters:

fn – The compiled function to execute (created by compile() method)