nnsight.contexts

The contexts module contains logic for managing the tracing and running of models with nnsight.tracing and nnsight.envoy

The primary two classes involved here are Tracer and Invoker.

The Tracer class creates a Graph around the underlying model of an NNsight. The graph tracks and manages the operations performed on the inputs and outputs of said model. Module’s envoys in the model expose their .output and .input attributes which when accessed, add to the computation graph of the tracer. To do this, they need to know about the current Tracer object, so each module’s envoy’s .tracer object is set to be the current Tracer.

The Tracer object also keeps track of the batch_size of the most recent input, the generation index for multi iteration generations, and all of the inputs made during its context in the .batched_input attribute. Inputs added to this attribute should be in a format where each index is a batch index and allows the model to batch all of the inputs together.

This is to keep things consistent. If two different inputs are in two different valid formats, they both become the same format and are easy to batch. In the case of LanguageModels, regardless of whether the input are string prompts, pre-processed dictionaries, or input ids, the batched input is only input ids. On exiting the Tracer context, the Tracer object should use the information and inputs provided to it to carry out the execution of the model.

The Invoker class is what actually accepts inputs to the model/graph, and it updates its parent Tracer object with the appropriate information about respective inputs. On entering the invoker context with some input, the invoker leverages the model to pre-process and prepare the input to the model. Using the prepared inputs, it updates its Tracer object with a batched version of the input, the size of the batched input, and the current generation index. It also runs a ‘meta’ version of the input through the model’s meta_model. This updates the sizes/dtypes of all of the module’s Envoys inputs and outputs based on the characteristics of the input.

nnsight comes with an implementation of a Tracer, Runner, which enables both local and remote execution.

class nnsight.contexts.Runner.Runner(*args, blocking: bool = True, remote: bool = False, **kwargs)

The Runner object manages the intervention tracing for a given model’s _execute method locally or remotely.

remote

If to use the remote NDIF server for execution of model and computation graph. (Assuming it’s running/working)

Type:

bool

blocking

If when using the server option, to hang until job completion or return information you can use to retrieve the job result.

Type:

bool

class nnsight.contexts.Tracer.Tracer(model: NNsight, validate: bool = True, **kwargs)

The Tracer class creates a nnsight.tracing.Graph.Graph around the ._model of a nnsight.models.NNsightModel.NNsight which tracks and manages the operations performed on the inputs and outputs of said model.

_model

nnsight Model object that ths context manager traces and executes.

Type:

nnsight.models.NNsightModel.NNsight

_graph

Graph which traces operations performed on the input and output of modules’ Envoys are added and later executed.

Type:

nnsight.tracing.Graph.Graph

_args

Positional arguments to be passed to function that executes the model.

Type:

List[Any]

_kwargs

Keyword arguments to be passed to function that executes the model.

Type:

Dict[str,Any]

_batch_size

Batch size of the most recent input. Used by Envoy to create input/output proxies.

Type:

int

_batch_start

Batch start of the most recent input. Used by Envoy to create input/output proxies.

Type:

int

_batched_input

Batched version of all inputs involved in this Tracer.

Type:

Any

_invoker

Currently open Invoker.

Type:

Invoker

apply(target: Callable, *args, validate: bool = False, **kwargs) → InterventionProxy

Helper method to directly add a function to the intervention graph.

Parameters:

• target (Callable) – Function to apply

• validate (bool) – If to try and run this operation in FakeMode to test it out and scan it.

Returns:

Proxy of applying that function.

Return type:

InterventionProxy

invoke(*inputs: Any, **kwargs) → Invoker

Create an Invoker context dor a given input.

Raises:

Exception – If an Invoker context is already open

Returns:

Invoker.

Return type:

Invoker

next(increment: int = 1) → None

Increments call_iter of all module Envoys. Useful when doing iterative/generative runs.

Parameters:

increment (int) – How many call_iter to increment at once. Defaults to 1.

class nnsight.contexts.Invoker.Invoker(tracer: Tracer, *inputs: Any, scan: bool = True, **kwargs)

An Invoker is meant to work in tandem with a nnsight.contexts.Tracer.Tracer to enter input and manage intervention tracing.

tracer

Tracer object to enter input and manage context.

Type:

nnsight.contexts.Tracer.Tracer

inputs

Initially entered inputs, then post-processed inputs from model’s ._prepare_inputs(…) method.

Type:

tuple[Any]

scan

If to execute the model using FakeTensor in order to update the potential sizes/dtypes of all modules’ Envoys’ inputs/outputs as well as validate things work correctly. Scanning is not free computation wise so you may want to turn this to false when running in a loop. When making interventions, you made get shape errors if scan is false as it validates operations based on shapes so for looped calls where shapes are consistent, you may want to have scan=True for the first loop. Defaults to True.

Type:

bool

kwargs

Keyword arguments passed to the model’s _prepare_inputs method.

Type:

Dict[str,Any]