envoy¶

envoy ¶

The Envoy — nnsight's window into a running PyTorch model.

An Envoy wraps a torch.nn.Module and mirrors its submodule tree, so every module in the model has a matching envoy reachable by the same attribute path (model.transformer.h[0].mlp). Envoys are the objects you interact with when tracing: they expose each module's live input/output during a forward pass, let you overwrite those values, read gradients, skip whole modules, and reach individual operations inside a forward via source.

You open a trace with with model.trace(x): and, inside the block, read or write envoy attributes as if the forward pass had paused at each module for you. Capture a value with .save() to use it after the trace:

.. code-block:: python

from nnsight.intervention.envoy import Envoy

model = Envoy(my_module)

with model.trace(x):
    hidden = model.layer1.output.save()   # captured mid-forward
    model.layer2.output[:] = 0            # overwrite layer2's output in place

print(hidden.shape)                        # available after the block

Gradients are available the same way. Call .backward() on a captured value as a context manager and, inside it, read .grad on tensors you captured earlier in the forward — you can edit gradients too:

.. code-block:: python

with model.trace(x):
    a1 = model.fc1.output
    loss = model.output.sum()
    with loss.backward():
        g = a1.grad.save()     # gradient flowing into fc1's output
        a1.grad = a1.grad * 2  # and it can be edited in place of autograd's

Locations must be read in execution order: asking for an earlier module's output after a later one has already run raises OutOfOrderError.

Envoy ¶

Envoy(module: Module, path: str = 'model', interleaver: Interleaver | None = None, rename: dict[str, str | list[str]] | None = None, envoys: dict | None = None)

Wraps a torch.nn.Module to expose and edit its values during a trace.

One envoy mirrors one module and reads or overwrites that module's live input/output (and gradients) as the forward pass runs, driving those interventions through a shared Interleaver. The child envoys mirror the module's submodule tree, so the whole model is reachable by attribute path from the root envoy. See the module docstring for the mental model.

ATTRIBUTE	DESCRIPTION
`path`	The module's dotted location in the tree, e.g. `"model.transformer.h.0"`. Every location the interleaver reads (`{path}.output`, `{path}.skip`) is derived from it.
`interleaver`	The `Interleaver` shared across the whole tree; it installs the controllers and routes values.
`_module`	The wrapped `torch.nn.Module`.
`_edits`	Default interventions registered by `edit`, replayed on every trace (a list of `Mediator`). TYPE: `list[Mediator]`
`_children`	The child envoys this envoy owns, in module order — one per module, so a module the tree already wraps elsewhere is not in it. TYPE: `list[Envoy]`
`_child_map`	Every entry of the wrapped module's `_modules`, by name and in module order, so an entry sharing its module with another keeps its own name. What indexing, iteration and the repr walk. TYPE: `dict[str, Envoy]`

path `instance-attribute` ¶

path = path

interleaver `instance-attribute` ¶

interleaver = interleaver if interleaver is not None else Interleaver()

OVERLOAD_PREFIX `class-attribute` `instance-attribute` ¶

OVERLOAD_PREFIX = 'E_'

iter `property` ¶

iter

Deprecated: use tracer.iter.

An alias for tracer.iter[...]; the iteration API lives on the tracer.

source `property` ¶

source: Source

Get the source code representation of the module.

Examples:

>>> model = TransformersModel("openai-community/gpt2", dispatch=True)
>>> print(model.transformer.h[0].attn.source)   # list the operations
>>> with model.trace("Hello World"):
...     attn = model.transformer.h[0].attn.source.attention_interface_1.output.save()

RETURNS	DESCRIPTION
`Source`	A `Source` exposing operation-level access.

device `property` ¶

device: device | None

The device of the module's first parameter, or None if it has none.

devices `property` ¶

devices: set[device]

The set of devices the module's parameters live on (empty if it has none).

setstate ¶

__setstate__(state)

getstate ¶

__getstate__() -> dict

trace ¶

trace(*args: Any, fn: Any = None, backend: Any = None, tracer_cls: type[InterleavingTracer] | None = None, trace: bool = True, **kwargs: Any) -> InterleavingTracer

Open a trace: a with block that runs the module and lets you read and edit its intermediate values.

Inside the block, read an envoy's .output/.input to capture a value, or assign to it to overwrite what the module passes on. Mark a value with .save() to keep it past the block.

Examples:

>>> model = TransformersModel("openai-community/gpt2", dispatch=True)
>>> with model.trace("Hello World"):
...     model.transformer.h[0].attn.output[0][:] = 0   # zero the attn output
...     output = model.output.save()
>>> print(output)

PARAMETER	DESCRIPTION
`*args`	Arguments to pass to the tracer TYPE: `Any` DEFAULT: `()`
`fn`	What the trace runs, defaulting to the module's `"__call__"`. Either a method name on this envoy or a bound callable. A subclass that drives the model through its own method points at it here — which is also how a plain `NNsight` subclass runs a driver that passes values through an attached standalone child. The `traceable` decorator sets it, so `model.generate(...)` is `trace(..., fn=self.generate)`. TYPE: `Any` DEFAULT: `None`
`backend`	What running the captured block means, defaulting to local execution. Set by `remote=` on the model classes that support it; see `Backend`. TYPE: `Any` DEFAULT: `None`
`tracer_cls`	Tracer class to construct instead of the default `InterleavingTracer` — an extension point for a custom tracer. TYPE: `type[InterleavingTracer] \| None` DEFAULT: `None`
`trace`	If `False`, bypass tracing — run the module directly on the inputs and return its output. A one-shot forward with no intervention: input prep, dispatch, and device placement still happen (via `interleave`), but no `with` block is captured. TYPE: `bool` DEFAULT: `True`
`**kwargs`	Keyword arguments to pass to the tracer TYPE: `Any` DEFAULT: `{}`

RETURNS	DESCRIPTION
`InterleavingTracer`	An InterleavingTracer for this module, or — when `trace=False` —
`InterleavingTracer`	the module's output directly.

edit ¶

edit(*, inplace: bool = False, backend: Any = None) -> EditingTracer

Open an editing tracer: capture interventions and store them as defaults.

The block is not executed against a live forward; instead the interventions it captures are stored and replayed on every later trace of the (edited) model. Clear them with clear_edits.

Examples:

>>> model = TransformersModel("openai-community/gpt2", dispatch=True)
>>> # The first layer's attention output will always be zeroed.
>>> with model.edit() as (tracer, edited_model):
...     edited_model.transformer.h[0].attn.output[0][:] = 0

>>> with model.trace("Hello World"):
...     output = model.output.save()          # original model, unedited
>>> print(output)

>>> with edited_model.trace("Hello World"):
...     edited_output = edited_model.output.save()   # edit applied
>>> print(edited_output)

PARAMETER	DESCRIPTION
`inplace`	If `False` (default), store the edit on a shallow copy and leave this envoy clean, so the edited behavior is opt-in through the copy. If `True`, store it on this envoy itself. TYPE: `bool` DEFAULT: `False`
`backend`	Backend for the underlying trace. TYPE: `Any` DEFAULT: `None`

RETURNS	DESCRIPTION
`EditingTracer`	An `EditingTracer`. Entering it binds the tracer — for its `iter` API — and, when `inplace=False`, the edited copy as well: `with model.edit() as (tracer, edited):`. With `inplace=True` only the tracer is bound.

clear_edits ¶

clear_edits() -> None

Drop every edit stored on this envoy, restoring its unedited behavior.

session ¶

session(backend: Any = None, tracer_cls: type[Tracer] | None = None) -> Tracer

Open a session: a scope enclosing several traces that share values.

Inside with model.session(): you can open multiple with model.trace(...) blocks and pass values between them — a value read in one trace is available in a later trace without an explicit .save(), because the session (not each individual trace) is the save boundary. Only values marked with nnsight.save survive past the session itself. Ordinary Python — loops, conditionals, building lists — runs natively in the session body.

Examples:

>>> with model.session():
...     with model.trace(x):
...         hidden = model.layer1.output      # no .save() needed
...     with model.trace(x):
...         out = (hidden * 2).save()         # `hidden` flows in
>>> print(out)

PARAMETER	DESCRIPTION
`backend`	What runs the captured session block. Defaults to running it in place (which executes the nested traces as it reaches them). TYPE: `Any` DEFAULT: `None`
`tracer_cls`	Tracer class to construct instead of the default `Tracer` — an extension point for a custom session tracer. TYPE: `type[Tracer] \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`Tracer`	A `Tracer` acting as the session scope.

inputs ¶

inputs(value: Any) -> Any

The module's forward inputs as an (args, kwargs) tuple.

Read or replace the whole input during a trace::

with model.trace("Hello World"):
    args, kwargs = model.transformer.h[0].attn.inputs

input ¶

input(value: Any) -> Any

output ¶

output(value: Any) -> Object

The module's forward output — read or replace it during a trace::

with model.trace("Hello World"): attn = model.transformer.h[0].attn.output[0].save()

all ¶

all()

Deprecated: use tracer.all().

skip ¶

skip(replacement: Any) -> None

Skip this module's execution, returning replacement as its output.

The module's forward is not run; replacement is used in its place.

Examples:

>>> model = TransformersModel("openai-community/gpt2", dispatch=True)
>>> with model.trace("Hello World"):
...     # Skip the first layer, passing its input through as its output.
...     model.transformer.h[0].skip(model.transformer.h[0].input)
...     output = model.output.save()
>>> print(output)

PARAMETER	DESCRIPTION
`replacement`	The value to use as the module's output; must match the shape the module would return. Read it from anywhere — including this module's own `.input`, which is offered before the skip gate. TYPE: `Any`

RETURNS	DESCRIPTION
`None`	Nothing; the skip takes effect when the module runs.

to ¶

to(device: device) -> Envoy

Move the wrapped module to device (in place).

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

RETURNS	DESCRIPTION
`Envoy`	This envoy, for method chaining.

cpu ¶

cpu(*args: Any, **kwargs: Any) -> Envoy

Move the wrapped module to the CPU (in place); return this envoy.

cuda ¶

cuda(*args: Any, **kwargs: Any) -> Envoy

Move the wrapped module to the GPU (in place); return this envoy.

interleave ¶

interleave(fn: Any, *args: Any, batcher: 'Batcher | None' = None, **kwargs: Any) -> Any

Run fn interleaved with the interleaver's registered workers.

This is the low-level driver behind trace; you rarely call it directly. The workers (edits + per-invoke interventions, with their batch groups) are set up on the interleaver first; this runs fn(*args, **kwargs) alongside them and clears them afterward.

PARAMETER	DESCRIPTION
`fn`	The callable to run, or a method name resolved against the module. TYPE: `Any`
`*args`	Positional inputs for a direct (untraced) call. They're wrapped as a single implicit invoke and assembled like a one-invoke trace (tokenized/collated), then moved to `device`. Ignored when `batcher` is given (the inputs come from assembling it). TYPE: `Any` DEFAULT: `()`
`batcher`	The per-invoke inputs to combine into one call. When given, it's assembled into `(args, kwargs)` and registered on the interleaver so `handle` can narrow each worker to its own rows; any `kwargs` passed alongside (trace-level params like `max_new_tokens`) override the assembled ones. TYPE: `'Batcher \| None'` DEFAULT: `None`
`**kwargs`	Keyword inputs — part of the direct call's single invoke when no `batcher` is given, else trace-level params for the assembled call. TYPE: `Any` DEFAULT: `{}`

RETURNS	DESCRIPTION
`Any`	Whatever `fn` returned (typically the module's forward output). TYPE: `Any`

getattr ¶

__getattr__(name: str) -> Any

call ¶

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

Run this module's forward — applying it ad hoc, out of execution order.

Inside a trace you can feed a module any input to compute with it away from its place in the forward pass — e.g. the logit lens, running lm_head on an intermediate hidden state::

with model.trace(prompt):
    hidden = model.transformer.h[-1].output
    logits = model.lm_head(model.transformer.ln_f(hidden))

While interleaving the module is called the ordinary way, with this trace stood down for the duration. Its own hooks and wrappers still fire — a runtime that keeps collectives around the forward (transformers tensor parallelism wraps module.forward with them) needs them to — while nnsight serves nothing and spends no occurrence, for this module or anything under it. That keeps the call from switching into the very worker greenlet making it, and leaves the module's real place in the forward pass untouched: calling a whole layer ad hoc does not consume the occurrence its children's real visit was going to fill. Outside a trace it is an ordinary module call.

Pass hook=True to let the trace watch the call too. Use it for a module attached to the tree that isn't part of the real forward pass — an adapter, LoRA, or SAE applied in an edit — so its internals become observable at .submodule.output::

model.transformer.h[0].adapter = MyAdapter()
with model.edit() as (tracer, edited):
    acts = edited.transformer.h[0].output
    edited.transformer.h[0].output =                     edited.transformer.h[0].adapter(acts, hook=True)
with edited.trace(prompt):
    inner = edited.transformer.h[0].adapter.inner.output.save()

The block above applies the adapter once, at the first time the layer runs. To apply it every time — each step of a generation loop — put the passthrough under the edit tracer's iter::

with model.edit(inplace=True) as tracer:
    for _ in tracer.iter[:]:
        acts = model.transformer.h[0].output
        model.transformer.h[0].output =                         model.transformer.h[0].adapter(acts, hook=True)

PARAMETER	DESCRIPTION
`*args`	Inputs to run the module's forward on. TYPE: `Any` DEFAULT: `()`
`hook`	Whether nnsight watches the call. If `False` (default) the module runs normally, its own hooks and all, but this trace is stood down for the duration — so neither this module nor anything under it serves a value or spends an occurrence, and its real place in the forward pass is untouched. If `True`, the trace watches too, so the call's internals are addressable at `.submodule.output`; use it for a module attached to the tree rather than one the forward pass already runs. TYPE: `bool` DEFAULT: `False`
`**kwargs`	Keyword inputs to the module's forward. TYPE: `Any` DEFAULT: `{}`

RETURNS	DESCRIPTION
`Any`	The module's output for these inputs.

setattr ¶

__setattr__(name: str, value: Any) -> None

iter ¶

__iter__() -> Iterator[Envoy]

Iterate over this envoy's direct children.

Yields each immediate child envoy — e.g. the blocks of a ModuleList, so for layer in model.model.layers: walks the layers. This is not recursive; use modules to walk the whole subtree.

YIELDS	DESCRIPTION
`Envoy`	Each direct child envoy, in order. TYPE:: `Envoy`

Example

::

for layer in model.model.layers:
    print(layer.path)

getitem ¶

__getitem__(key: Any) -> Envoy

Index into direct child envoys, e.g. for a ModuleList.

An int or str key resolves by name, the way the wrapped module indexes it — layers[2] is the module layers holds at "2", even when an earlier entry is a module the tree already wraps elsewhere and so has no envoy of its own here.

PARAMETER	DESCRIPTION
`key`	An index the wrapped module accepts (an int, or a str), or a slice over this envoy's children in module order. TYPE: `Any`

RETURNS	DESCRIPTION
`Envoy`	The child envoy at `key` (e.g. `model.layers[0]` for the TYPE: `Envoy`
`Envoy`	first block of a `ModuleList`).

len ¶

__len__() -> int

The number of entries in the wrapped module (e.g. a ModuleList's length).

get ¶

get(path: str) -> Any

Resolve a dotted path from this envoy, e.g. "transformer.h.0.mlp".

A programmatic alternative to attribute access, useful when the path is built at runtime. Outside a trace it returns the descendant envoy; inside one, a trailing .output/.input resolves through to the live value.

Examples:

>>> model = TransformersModel("openai-community/gpt2", dispatch=True)
>>> module = model.get("transformer.h.0.mlp")
>>> with model.trace("Hello"):
...     value = model.get("transformer.h.0.mlp.output").save()

PARAMETER	DESCRIPTION
`path`	A `.`-separated attribute path relative to this envoy. TYPE: `str`

RETURNS	DESCRIPTION
`Any`	The resolved child `Envoy`, or the live value when `path`
`Any`	ends in an intervention attribute during a trace.

modules ¶

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

Flatten this envoy's whole subtree (children first, then self) into a list.

PARAMETER	DESCRIPTION
`include_fn`	Optional predicate on an envoy; only those it returns `True` for are kept. TYPE: `Callable[[Envoy], bool] \| None` DEFAULT: `None`
`names`	If `True`, yield `(path, envoy)` tuples instead of envoys. TYPE: `bool` DEFAULT: `False`

RETURNS	DESCRIPTION
`list[Any]`	A list of `Envoy` (or `(path, Envoy)` tuples when `names`).

named_modules ¶

named_modules(include_fn: Callable[[Envoy], bool] | None = None) -> list[tuple[str, Envoy]]

Flatten the subtree into (path, envoy) tuples: modules with names=True.

PARAMETER	DESCRIPTION
`include_fn`	Optional predicate on an envoy; only those it returns `True` for are kept. TYPE: `Callable[[Envoy], bool] \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`list[tuple[str, Envoy]]`	A list of `(path, Envoy)` tuples for the included envoys.

repr ¶

__repr__() -> str

traceable ¶

traceable(method: Callable) -> Callable

Make an Envoy method usable as a trace context.

with envoy.method(...): traces the method (runs it interleaved with the block's interventions); envoy.method(...) just calls it. While already interleaving, it always just calls the method (we're inside a trace).

envoy¶

envoy ¶

Envoy ¶

path instance-attribute ¶

interleaver instance-attribute ¶

OVERLOAD_PREFIX class-attribute instance-attribute ¶

iter property ¶

source property ¶

device property ¶

devices property ¶

__setstate__ ¶

__getstate__ ¶

trace ¶

edit ¶

clear_edits ¶

session ¶

inputs ¶

input ¶

output ¶

all ¶

skip ¶

to ¶

cpu ¶

cuda ¶

interleave ¶

__getattr__ ¶

__call__ ¶

__setattr__ ¶

__iter__ ¶

__getitem__ ¶

__len__ ¶

get ¶

modules ¶

named_modules ¶

__repr__ ¶

traceable ¶

path `instance-attribute` ¶

interleaver `instance-attribute` ¶

OVERLOAD_PREFIX `class-attribute` `instance-attribute` ¶

iter `property` ¶

source `property` ¶

device `property` ¶

devices `property` ¶

setstate ¶

getstate ¶

getattr ¶

call ¶

setattr ¶

iter ¶

getitem ¶

len ¶

repr ¶