envoy¶

envoy ¶

Envoy ¶

Envoy(module: Module, interleaver: Optional[Interleaver] = None, path: Optional[str] = 'model', rename: Optional[Dict[str, Union[str, List[str]]]] = None)

Bases: Batchable

A proxy class that wraps a PyTorch module to enable intervention during execution.

This class provides access to module inputs and outputs during forward passes, and allows for modification of these values through an interleaving mechanism. It serves as the primary interface for inspecting and modifying the behavior of neural network modules during execution.

ATTRIBUTE	DESCRIPTION
`path`	The module's location in the model hierarchy. Example: "model.encoder.layer1" indicates this module is the first layer of the encoder in the model. TYPE: `str`
`_module`	The underlying PyTorch module TYPE: `Module`
`_source`	Source code representation of the module TYPE: `Optional[EnvoySource]`
`_interleaver`	Interleaver for managing execution flow TYPE: `Optional[Interleaver]`
`_default_mediators`	List of default mediators created with .edit TYPE: `List[List[str]]`
`_children`	List of child Envoys TYPE: `List[Envoy]`
`_alias`	Aliaser object for managing aliases TYPE: `Aliaser`

PARAMETER	DESCRIPTION
`module`	The PyTorch module to wrap TYPE: `Module`
`interleaver`	Optional interleaver for managing execution flow TYPE: `Optional[Interleaver]` DEFAULT: `None`
`path`	Optional path string representing the module's location in the model hierarchy TYPE: `Optional[str]` DEFAULT: `'model'`
`rename`	Optional dictionary mapping module names to alias names. Example: {"layer1": "first_layer", "layer2": "second_layer"} Example: {".model.layers": ".layers"} <-- Mounts .layers to the root model. Example: {".transformer": ["model", "mdl"]} <-- Allows access of .transformer as .model or .mdl TYPE: `Optional[Dict[str, Union[str, List[str]]]]` DEFAULT: `None`

path `instance-attribute` ¶

path = path

interleaving `property` ¶

interleaving: bool

Check if the Envoy is currently nterleaving.

RETURNS	DESCRIPTION
`bool`	True if the Envoy is interleaving, False otherwise

output `property` `writable` ¶

output: Object

Get the output of the module's forward pass.

This property allows access to the return values produced by the module during the forward pass.

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True) with model.trace("Hello World"): attn = model.transformer.h[0].attn.output[0].save() print(attn)

RETURNS	DESCRIPTION
`Object`	The module's output values

inputs `property` `writable` ¶

inputs: Tuple[Tuple[Object], Dict[str, Object]]

Get the inputs to the module's forward pass.

This property provides access to all input values passed to the module during the forward pass.

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True) with model.trace("Hello World"): args, kwargs = model.transformer.h[0].attn.inputs

RETURNS	DESCRIPTION
`Tuple[Tuple[Object], Dict[str, Object]]`	The module's input values as a tuple of positional and keyword arguments. i.e (args, kwargs)

input `property` `writable` ¶

input: Object

Get the first input to the module's forward pass.

This is a convenience property that returns just the first input value from all inputs passed to the module. So first positional argument, or first keyword argumetn if there are no positional arguments.

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True) with model.trace("Hello World"): hidden_states = model.transformer.h[0].attn.input.save() print(hidden_states)

RETURNS	DESCRIPTION
`Object`	The first input value

source `property` ¶

source: EnvoySource

Get the source code representation of the module.

This property provides access to the module's source code with operations highlighted, allowing for inspection and intervention at specific points.

Example:

>>> model = LanguageModel("gpt2", device_map='auto', dispatch=True)

>>> # We can print to see the formward method of the module and names associated with the operations within.
>>> print(model.transformer.h[0].attn.source)

                                       60
                                       61     if using_eager and self.reorder_and_upcast_attn:
  self__upcast_and_reordered_attn_0 -> 62         attn_output, attn_weights = self._upcast_and_reordered_attn(
                                       63             query_states, key_states, value_states, attention_mask, head_mask
                                       64         )
                                       65     else:
  attention_interface_0             -> 66         attn_output, attn_weights = attention_interface(
                                       67             self,
                                       68             query_states,
                                       69             key_states,
                                       70             value_states,
                                       71             attention_mask,
                                       72             head_mask=head_mask,
                                       73             dropout=self.attn_dropout.p if self.training else 0.0,
                                       74             is_causal=is_causal,
                                       75             **kwargs,
                                       76         )
                                       77
  attn_output_reshape_0             -> 78     attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
  contiguous_0                      ->  +     ...
  self_c_proj_0                     -> 79     attn_output = self.c_proj(attn_output)
  self_resid_dropout_0              -> 80     attn_output = self.resid_dropout(attn_output)
                                       81
                                       82     return attn_output, attn_weights
                                       83

>>> # We can print out one of these to see the only the operation and a few operations before and after.
>>> print(model.transformer.h[0].attn.source.attention_interface_0)

.transformer.h.0.attn.attention_interface_0:

     ....

         if using_eager and self.reorder_and_upcast_attn:
             attn_output, attn_weights = self._upcast_and_reordered_attn(
                 query_states, key_states, value_states, attention_mask, head_mask
             )
         else:
     -->     attn_output, attn_weights = attention_interface( <--
                 self,
                 query_states,
                 key_states,
                 value_states,
                 attention_mask,
                 head_mask=head_mask,
     ....

>>> with model.trace("Hello World"):
>>>     # Now we can access it like we would any other Envoy with .input or .output to grab the intermediate value.
>>>     attn = model.transformer.h[0].attn.source.attention_interface_0.output.save()

>>> print(attn)

RETURNS	DESCRIPTION
`EnvoySource`	An EnvoySource object containing the module's source code and operations

iter `property` ¶

iter

device `property` ¶

device: Optional[device]

Get the device the module is on. Finds the first parameter and return its device.

devices `property` ¶

devices: Optional[set[device]]

Get the devices the module is on. Finds all parameters and return their devices.

getitem ¶

__getitem__(key: str) -> Envoy

Access a child Envoy by index for Module Lists.

PARAMETER	DESCRIPTION
`key`	The index of the child Envoy to retrieve TYPE: `str`

RETURNS	DESCRIPTION
`Envoy`	The child Envoy at the specified index

call ¶

__call__(*args, hook: bool = False, **kwargs)

trace ¶

trace(*args, fn: Optional[Callable] = None, trace: bool = None, tracer_cls: Type[InterleavingTracer] = InterleavingTracer, **kwargs)

Create a tracer for this module.

This method returns a tracer that can be used to capture and modify the execution of the module.

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True) with model.trace("Hello World"): model.transformer.h[0].attn.output[0][:] = 0
output = model.output.save()
print(output)

PARAMETER	DESCRIPTION
`*args`	Arguments to pass to the tracer DEFAULT: `()`
`**kwargs`	Keyword arguments to pass to the tracer DEFAULT: `{}`

RETURNS	DESCRIPTION
	An InterleavingTracer for this module

scan ¶

scan(*args, **kwargs)

Just like .trace() but runs the model in fake tensor mode to validate operations and inspect tensor shapes.

This method returns a tracer that runs the model in fake tensor mode to validate operations and inspect tensor shapes without performing actual computation. This is useful for: - Validating that operations will work with given input shapes - Inspecting the shapes and types of tensors that would flow through the model - Debugging shape mismatches or other tensor-related issues.

Note this will not dispatch the model if not dispatched.

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True)

Value error as the fake inputs and outputs have not been scanned in.¶

print(model.transformer.h[0].mlp.output.shape)

Scan the model to validate operations and inspect shapes¶

with model.scan("Hello World"): # Access fake inputs/outputs to inspect shapes attn_input = model.transformer.h[0].attn.input.save() attn_output = model.transformer.h[0].attn.output[0].save() print(f"Attention input shape: {attn_input.shape}") print(f"Attention output shape: {attn_output.shape}") print(model.transformer.h[0].mlp.output.shape)

PARAMETER	DESCRIPTION
`*args`	Arguments to pass to the tracer DEFAULT: `()`
`**kwargs`	Keyword arguments to pass to the tracer DEFAULT: `{}`

RETURNS	DESCRIPTION
	A ScanningTracer for this module

edit ¶

edit(*, inplace: bool = False)

Create an editing tracer for this module. Allows for setting default interventions. This means this tracer won't execute the module, but will instead set default interventions that are applied on all future executions.

Edits can be cleared with Envoy.clear_edits().

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True)

Now the first layer attention output will always be 0.¶

with model.edit() as edited_model: edited_model.transformer.h[0].attn.output[:] = 0

with model.trace("Hello World"): output = model.output.save()

The orignal model will have the default output.¶

print(output)

with edited_model.trace("Hello World"): edited_output = edited_model.output.save()

The edited model will have the output after our intervention.¶

print(edited_output)

PARAMETER	DESCRIPTION
`inplace`	Whether to edit in place. Defaults to False. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`EditingTracer`	An EditingTracer for this module

clear_edits ¶

clear_edits()

Clear all edits for this Envoy.

export_edits ¶

export_edits(name: str, export_dir: Optional[str] = None, variant: str = '__default__')

TODO

PARAMETER	DESCRIPTION
`name`	description TYPE: `str`
`export_dir`	description. Defaults to None. TYPE: `Optional[str]` DEFAULT: `None`
`variant`	description. Defaults to 'default'. TYPE: `str` DEFAULT: `'__default__'`

RAISES	DESCRIPTION
`ValueError`	description

import_edits ¶

import_edits(name: str, export_dir: Optional[str] = None, variant: str = '__default__')

TODO

PARAMETER	DESCRIPTION
`name`	description TYPE: `str`
`export_dir`	description. Defaults to None. TYPE: `Optional[str]` DEFAULT: `None`
`variant`	description. Defaults to 'default'. TYPE: `str` DEFAULT: `'__default__'`

session ¶

session(*args, tracer_cls: Type[Tracer] = Tracer, **kwargs)

all ¶

all()

next ¶

next(step: int = 1)

skip ¶

skip(replacement: Any)

Skips the execution of this module duting execution / interleaving. Behavior is the module will not be executed and will return a replacement value instead.

Example

model = LanguageModel("gpt2", device_map='auto', dispatch=True) with model.trace("Hello World"): # Skip the first layer and replace it with the input to the layer. model.transformer.h[0].skip((model.transformer.h[0].input, None)) output = model.output.save() print(output)

PARAMETER	DESCRIPTION
`replacement`	The replacement value to replace the module's output with. TYPE: `Any`

wait_for_input ¶

wait_for_input()

Wait for the input to the module to be available.

wait_for_output ¶

wait_for_output()

Wait for the output to the module to be available.

to ¶

to(device: device)

Move the module to a specific device.

This method moves the underlying PyTorch module to the specified device.

PARAMETER	DESCRIPTION
`device`	The device to move the module to TYPE: `device`

RETURNS	DESCRIPTION
	Self, for method chaining

cpu ¶

cpu(*args, **kwargs)

Move the module to the CPU.

cuda ¶

cuda(*args, **kwargs)

Move the module to the GPU.

modules ¶

modules(include_fn: Callable[[Envoy], bool] = None, names: bool = False) -> List[Envoy]

Get all modules in the Envoy tree.

This method returns all Envoys in the tree, optionally filtered by an inclusion function.

PARAMETER	DESCRIPTION
`include_fn`	Optional function to filter modules TYPE: `Callable[[Envoy], bool]` DEFAULT: `None`
`names`	Whether to include module names in the result TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`List[Envoy]`	A list of Envoys or (name, Envoy) tuples

named_modules ¶

named_modules(*args, **kwargs) -> List[Tuple[str, Envoy]]

Returns all Envoys in the Envoy tree along with their name/module_path.

This is a convenience method that calls modules() with names=True.

PARAMETER	DESCRIPTION
`include_fn`	Optional function to be ran against all Envoys to check if they should be included in the final collection of Envoys. Defaults to None. TYPE: `Callable`
`args, *kwargs`	Additional arguments to pass to modules()

RETURNS	DESCRIPTION
`List[Tuple[str, Envoy]]`	List[Tuple[str, Envoy]]: Included Envoys and their names/module_paths.

get ¶

get(path: str) -> Object

Gets the Envoy/Proxy via its path.

e.x: model = nnsight.LanguageModel("openai-community/gpt2")

module = model.get('transformer.h.0.mlp')

with model.trace("Hello"):
    value = model.get('transformer.h.0.mlp.output').save()

PARAMETER	DESCRIPTION
`path`	'.' separated path. TYPE: `str`

RETURNS	DESCRIPTION
`Object`	Union[Envoy, InterventionProxyType]: Fetched Envoy/Proxy

interleave ¶

interleave(fn: Union[Callable, str], *args, **kwargs)

len ¶

__len__()

Get the length of the Envoy.

iter ¶

__iter__()

Iterate over the Envoy.

str ¶

__str__()

String representation of the Envoy.

RETURNS	DESCRIPTION
	A string representation of the Envoy showing its path

reprlist ¶

__reprlist__()

repr ¶

__repr__()

Representation of the Envoy.

RETURNS	DESCRIPTION
	The string representation of the Envoy

getattr ¶

__getattr__(name: str) -> Union[Module, Envoy, Any]

Get an attribute from the underlying module.

If the attribute is callable, it will be wrapped in a tracer to enable intervention during execution.

PARAMETER	DESCRIPTION
`name`	The name of the attribute to get TYPE: `str`

RETURNS	DESCRIPTION
`Union[Module, Envoy, Any]`	The attribute value, possibly wrapped in a tracer

RAISES	DESCRIPTION
`AttributeError`	If the attribute doesn't exist

setattr ¶

__setattr__(key: Any, value: Any) -> None

Set an attribute on the Envoy.

If the value is a PyTorch module, it will be wrapped in an Envoy to enable intervention during execution.

PARAMETER	DESCRIPTION
`key`	The attribute name TYPE: `Any`
`value`	The attribute value TYPE: `Any`

getstate ¶

__getstate__()

setstate ¶

__setstate__(state)

OperationEnvoy ¶

OperationEnvoy(name: str, source: str, line_number: int, interleaver: Optional[Interleaver] = None)

Represents a specific operation within a module's forward pass.

This class provides access to the inputs and outputs of individual operations within a module's execution, allowing for fine-grained inspection and intervention at the operation level.

PARAMETER	DESCRIPTION
`name`	The fully qualified name of the operation TYPE: `str`
`source`	The source code of the module containing the operation TYPE: `str`
`line_number`	The line number of the operation in the source TYPE: `int`
`interleaver`	Optional interleaver for managing execution flow TYPE: `Optional[Interleaver]` DEFAULT: `None`

name `instance-attribute` ¶

name = name

source_code `instance-attribute` ¶

source_code = source

line_number `instance-attribute` ¶

line_number = line_number

output `property` `writable` ¶

output: Union[Any, Tensor]

Get the output of this operation.

This property provides access to the return value(s) produced by the operation during execution.

RETURNS	DESCRIPTION
`Union[Any, Tensor]`	The operation's output value(s)

inputs `property` `writable` ¶

inputs: Tuple[Tuple[Any, Tensor], Dict[str, Union[Tensor, Any]]]

Get the inputs to this operation.

This property provides access to all input value(s) passed to the operation during execution, structured as a tuple of positional and keyword arguments.

RETURNS	DESCRIPTION
`Tuple[Tuple[Any, Tensor], Dict[str, Union[Tensor, Any]]]`	The operation's input value(s)

input `property` `writable` ¶

input: Union[Any, Tensor]

Get the first input to the operation.

This is a convenience property that returns just the first input value from all inputs passed to the operation.

RETURNS	DESCRIPTION
`Union[Any, Tensor]`	The first input value

source `property` ¶

source: EnvoySource

Get the source code of the operation.

This property provides access to the operation's source code with nested operations highlighted, allowing for inspection and intervention at specific points.

RETURNS	DESCRIPTION
`EnvoySource`	An EnvoySource object containing the operation's source code and nested operations

str ¶

__str__()

String representation showing the operation in context.

This method returns a formatted string showing the operation's source code with surrounding context lines and highlighting the operation line.

RETURNS	DESCRIPTION
	A formatted string showing the operation's source code with context

EnvoySource ¶

EnvoySource(name: str, source: str, line_numbers: dict, interleaver: Optional[Interleaver] = None)

Represents the source code of a module with operations highlighted.

This class provides access to the individual operations within a module's source code, allowing for inspection and intervention at specific points in the code. It serves as a bridge between the source code representation and the runtime execution of operations.

PARAMETER	DESCRIPTION
`name`	The fully qualified name of the module or operation TYPE: `str`
`source`	The source code string TYPE: `str`
`line_numbers`	A dictionary mapping operation names to line numbers TYPE: `dict`
`interleaver`	Optional interleaver for managing execution flow TYPE: `Optional[Interleaver]` DEFAULT: `None`

source `instance-attribute` ¶

source = source

line_numbers `instance-attribute` ¶

line_numbers = line_numbers

operations `instance-attribute` ¶

operations: List[OperationEnvoy] = []

str ¶

__str__()

String representation showing the source code with operations highlighted.

This method returns a formatted string showing the source code with operation names and line numbers, making it easy to identify intervention points.

RETURNS	DESCRIPTION
	A formatted string showing the source code with operation names and line numbers

getattribute ¶

__getattribute__(name: str) -> Union[OperationEnvoy]

Aliaser ¶

Aliaser(rename: Dict[str, Union[str, List[str]]])

PARAMETER	DESCRIPTION
`rename`	Dictionary mapping module names to alias names. Example: {"layer1": "first_layer", "layer2": "second_layer"} Example: {".model.layers": ".layers"} <-- Mounts .layers to the root model. Example: {".transformer": ["model", "mdl"]} <-- Allows access of .transformer as .model or .mdl TYPE: `Dict[str, Union[str, List[str]]]`

ATTRIBUTE	DESCRIPTION
`rename`	Dictionary mapping module names to alias names. TYPE: `Dict[str, Union[str, List[str]]]`
`alias_to_name`	Dictionary mapping alias names to module names. TYPE: `Dict[str, str]`
`name_to_aliases`	Dictionary mapping module names to list of alias names. TYPE: `Dict[str, List[str]]`
`extras`	Dictionary mapping attribute paths (.transformer.h) to list of alias names. Used to show dot seperated attributes in the string representation of the Envoy. TYPE: `Dict[str, List[str]]`

rename `instance-attribute` ¶

rename = rename

alias_to_name `instance-attribute` ¶

alias_to_name = {}

name_to_aliases `instance-attribute` ¶

name_to_aliases = {}

extras `instance-attribute` ¶

extras = {}

build ¶

build(envoy: Envoy)

trace_only ¶

trace_only(fn: Callable)

envoy¶

envoy ¶

Envoy ¶

path instance-attribute ¶

interleaving property ¶

output property writable ¶

inputs property writable ¶

input property writable ¶

source property ¶

iter property ¶

device property ¶

devices property ¶

__getitem__ ¶

__call__ ¶

trace ¶

scan ¶

Value error as the fake inputs and outputs have not been scanned in.¶

Scan the model to validate operations and inspect shapes¶

edit ¶

Now the first layer attention output will always be 0.¶

The orignal model will have the default output.¶

The edited model will have the output after our intervention.¶

clear_edits ¶

export_edits ¶

import_edits ¶

session ¶

all ¶

next ¶

skip ¶

wait_for_input ¶

wait_for_output ¶

to ¶

cpu ¶

cuda ¶

modules ¶

named_modules ¶

get ¶

interleave ¶

__len__ ¶

__iter__ ¶

__str__ ¶

__reprlist__ ¶

__repr__ ¶

__getattr__ ¶

__setattr__ ¶

__getstate__ ¶

__setstate__ ¶

OperationEnvoy ¶

name instance-attribute ¶

source_code instance-attribute ¶

line_number instance-attribute ¶

output property writable ¶

inputs property writable ¶

input property writable ¶

source property ¶

__str__ ¶

EnvoySource ¶

source instance-attribute ¶

line_numbers instance-attribute ¶

operations instance-attribute ¶

__str__ ¶

__getattribute__ ¶

Aliaser ¶

rename instance-attribute ¶

alias_to_name instance-attribute ¶

name_to_aliases instance-attribute ¶

extras instance-attribute ¶

build ¶

trace_only ¶

path `instance-attribute` ¶

interleaving `property` ¶

output `property` `writable` ¶

inputs `property` `writable` ¶

input `property` `writable` ¶

source `property` ¶

iter `property` ¶

device `property` ¶

devices `property` ¶

getitem ¶

call ¶

len ¶

iter ¶

str ¶

reprlist ¶

repr ¶

getattr ¶

setattr ¶

getstate ¶

setstate ¶

name `instance-attribute` ¶

source_code `instance-attribute` ¶

line_number `instance-attribute` ¶

output `property` `writable` ¶

inputs `property` `writable` ¶

input `property` `writable` ¶

source `property` ¶

str ¶

source `instance-attribute` ¶

line_numbers `instance-attribute` ¶

operations `instance-attribute` ¶

str ¶

getattribute ¶

rename `instance-attribute` ¶

alias_to_name `instance-attribute` ¶

name_to_aliases `instance-attribute` ¶

extras `instance-attribute` ¶