source¶
Expose, edit, and skip a module's forward at the operation level.
A module's input/output are the two locations its controller hands to the
Interleaver. Everything in
between — the individual operations a forward performs — is invisible to it
because it isn't a submodule with a controller of its own.
This module makes those intermediates observable, editable, and skippable without
the interleaver knowing anything about source. The interleaver runs on one
primitive: a provider-location string and Interleaver.handle, which
serves a value to interventions and returns it back edited if one wrote to that
location. Inputs/outputs are just the two locations a controller emits; here we
add more, mid-forward:
- Parse the module's
forwardand rewrite every callfn(*args, **kwargs)into__nnsight_op__("source.{name}_{n}", fn, *args, **kwargs). At run timemake_opbrackets the call withInterleaver.handleon its.input(before) and.output(after) — both readable/replaceable — and a.skipgate that can bypass the call entirely. nameis the called function's dotted path joined with_(self.act(...)→self_act,torch.relu(...)→torch_relu,dropout(...)→dropout);nis a per-name counter in execution order (nested calls run inner-first, so the inner call is_0), which is the order the interleaver serves values.- Rewrite every assignment
x = valueintox = __nnsight_op__("source.x_{n}", __nnsight_bind__, value)— the same bracket around an identity, so a value that is not a call's return (aq @ kproduct, a running stateS = S * decay + updateinside a loop) is addressable too.nis the same per-name counter calls use, so a name that is bound and then called (attention_interface = ...; attention_interface(...)) isattention_interface_0at the binding andattention_interface_1at the call.
A decorated forward is instrumented through its decorators: a wrapper that
calls the function it closes over is peeled and rebuilt around the instrumented
function (decorator_chain, rewrap); one that doesn't — a dispatcher that
hands the function to a lookup and calls the result — is instrumented as it is,
its closure intact (compile_source), so the call that actually runs is
the operation to drill into.
Installation is permanent. When an envoy is built its module's forward is
replaced by a Controller over a single per-module State (see STATE): it
hands off .input, gates on .skip, runs the body — the module's own
forward, or the source-instrumented one once .source is used — and
hands off .output. The controller is inert outside a trace, so later runs
work regardless of request order, and source and skip compose on one wrapper. A
module wrapped by several envoys routes to whichever interleaver is running
(State.active).
An Envoy exposes operations as envoy.source.{name}_{n}, whose
.input/.inputs/.output/.skip mirror an Envoy's own, one level finer.
Bases: Exception
There is no Python source to instrument.
A builtin or C function, a function compiled from a string, a call into a
submodule (which has its own .source), or an assignment (which has no
callee). A decorated forward is not one of these: it is peeled,
instrumented, and rebuilt.
Per-module source/skip state, stored at module.__dict__[STATE].
Created when a module is first sourced or skipped. The controller and op probes read it live on every call, so re-wrapping — or wrapping the same module in more than one Envoy at once — just registers another interleaver here.
routes lists each interleaver that instrumented this module with the
path it addresses the module by; active picks the one whose trace is
currently running (there is at most one).
Because this state lives on the module (module.__dict__[STATE]), it holds
neither the module nor its interleavers strongly — the module would sit in a
reference cycle and never be freed by refcounting. The interleavers are held by
weakref (a finished local wrapper's interleaver drops out on its own; a
server's persistent interleaver stays, so the same module serves request after
request), and body is the unbound forward
(a plain function taking self) rather than a bound method that would pin it —
which a module carrying its own instance-level forward is the one exception to
(module_body): what it holds is already bound to that module, and only garbage
collection frees it.
class-attribute
instance-attribute
¶
What a copy or a pickle of this module carries: the body as it was before instrumentation, and no routes.
The registry is weakrefs, which pickle refuses outright ("cannot pickle
'weakref.ReferenceType' object" was the whole reason a wrapped module could
never be pickled again), and a copy belongs to whatever wraps it next, not
to the traces the original is in. The instrumented body goes with them: it
and its Compiled are built at run time — a function pickle can't name and
a code object it can't write at all — and install_source rebuilds them
from original the next time the copy is sourced.
Record that interleaver reaches this module at path.
The (interleaver, path, locations) whose trace is running now and has workers, or Nones — the gate every module call and every op passes.
A plain list walked on every module call, so it has to stay cheap: one
entry in all but the shared-module case, and a dead weakref costs a test.
busy is False for a run with no workers (a vLLM step with no nnsight
requests in it), where every handoff would be a no-op.
Bases: NodeTransformer
Rewrite every call into __nnsight_op__(location, fn, *args, **kwargs) and
every assignment into target = __nnsight_op__(location, __nnsight_bind__, value).
Numbers occurrences in execution order: a call's arguments (and an
assignment's value) are visited before the node is assigned its counter, so
f(f(x)) gives the inner f f_0 and the outer f_1, and
h = relu(x) gives relu_0 then h_0. Calls and assignments share
one counter per name.
staticmethod
¶
The attribute chain of expr and whether it is rooted in a name.
self.a.b → (['self', 'a', 'b'], True); x[i].y → (['x', 'y'], True)
(a subscript says where in the object, not what the object is called);
(a @ b).sum → (['sum'], False).
__nnsight_op__("source.{name}_{n}", fn, *args, **keywords) at node.
value routed through the identity under the target's name.
a, b = e1, e2 binds each name its own value, so each element gets its
own op (the tuple is still built before any name is bound, so a, b = b, a
still swaps). Any other unpacking, and a target with no name to label
(f()[0] = v), is left as it is.
The forward installed on an instrumented module: the module's handoff.
This is where a module's .input, .skip and .output reach the
interleaver -- the same three handles run_op emits for an operation,
one level up. Being the forward rather than a hook keeps the module on
PyTorch's fast call path, and it runs inside the module's own hooks, so a
runtime that keeps collectives in them sees the pre-collective value here --
which is what its Fragments describes.
State.body is the (unbound) original forward or, once sourced, the
instrumented one, and a .skip bypasses it entirely (and, if sourced, all
its ops).
An object rather than a closure so that copying a module can rebind it: a
function is atomic to copy and to pickle, so a closure came through
copy.deepcopy still pointing at the module it was built for -- the copy
computed with the original's weights, silently and outside any trace.
Holds the module by weakref: the module owns this controller (as its
forward), so a strong back-reference would cycle.
functools.update_wrapper preserves the forward's signature, which
generate() introspects to decide whether to pass
attention_mask/position_ids.
Rebind to the copied module, so the copy runs on its own weights.
Copying a module registers the copy in memo before it copies the
__dict__ this controller lives in, so by the time this runs the copy is
already there to bind to; deep-copying the module here is what finds it.
Pickle as the module and its state, rebuilt on the far side.
The weakref is what pickle chokes on, and it is only a way to reach the
module -- which is already memoized by the time its __dict__ is
written, so naming it here costs nothing.
A single operation inside a module's forward, e.g. source.torch_relu_0.
You never construct one directly; you reach it by indexing a Source
with an operation's {callable}_{occurrence} name (envoy.source.torch_relu_0).
It is the operation-level analogue of an
Envoy: where an Envoy exposes a
submodule's .input/.output, a SourceEnvoy exposes those same
handles for a single call the forward makes — one level finer.
Inside a trace, each handle both reads and writes the live value:
output— the operation's return value.input— the operation's first argument.inputs— the operation's full(args, kwargs).skip— bypass the call, substituting a value for its output.source— drill into the called function, exposing its own operations one level deeper (recursively).
Reading returns the value; assigning replaces it for the rest of the forward.
These handles are only meaningful inside a with envoy.trace(...): block.
To use a captured value after the trace, call .save() on it.
Examples:
Capture and edit an intermediate operation (for a forward that runs
h = torch.relu(self.fc1(x)))::
with model.trace(x):
pre = model.layer1.source.torch_relu_0.output.save() # capture
model.layer1.source.torch_relu_0.output = pre * 2 # and rescale it
On a real transformer, reach the activation inside an MLP::
with model.trace(prompt):
act = model.transformer.h[0].mlp.source.self_act_0.output.save()
Drill into a called function to reach an operation inside it::
with model.trace(prompt):
attn = model.transformer.h[0].attn.source
out = attn.attention_interface_1.source.attn_output_transpose_0.output.save()
property
¶
Drill into the called function, exposing its operations recursively.
Returns a Source over the function this operation calls, so its
internal operations become addressable as
...source.{name}.source.{inner} — with the same
.input/.output/.inputs/.skip/.source handles, to any
depth.
Only available inside a trace: the called function is resolved from the
live value flowing through the call at run time (a call target is often a
local variable, e.g. an attention implementation, so it can't be found
statically). Raises SourceNotAvailable if the target has no
recoverable Python source (a builtin/C function) or is itself a submodule
(call .source on that submodule directly instead).
Examples:
The operation's return value.
Read it to capture the value the call produced; assign to it to replace
that value for the remainder of the forward (downstream operations
see the replacement). In-place edits work too.
Examples:
The operation's arguments as an (args, kwargs) pair.
Use this when you need every argument (or the keyword arguments); for the
common case of a single leading argument, input is more direct.
Assigning a new (args, kwargs) pair replaces the arguments the call
runs with.
Examples:
Skip this operation, using replacement as its output.
The call never runs; replacement takes the place of its return value
and flows on to whatever consumed it. Use it to short-circuit expensive
or unwanted compute, or to splice in a value of your own.
Reading .output for a skipped operation returns replacement.
| PARAMETER | DESCRIPTION |
|---|---|
replacement
|
Value substituted for the operation's return value.
TYPE:
|
Examples:
A window of the forward source around this operation's call site.
The call site is flagged with --> / <-- so you can confirm you
indexed the operation you meant, with a few surrounding lines for context
(.... marks truncation above or below). Falls back to the operation's
dotted path when the source text is unavailable.
Examples:
def forward(self, x):
--> h = torch.relu(self.fc1(x)) <--
return self.fc2(h)
A module's forward decomposed into its individual operations.
Reached as envoy.source (e.g. model.layer1.source). Every call the
forward makes becomes an operation named {callable}_{occurrence},
where callable is the called function's dotted path joined with _:
self.fc1(x) → self_fc1_0, torch.relu(...) → torch_relu_0. The
occurrence counter is per name and runs in execution order (nested calls run
inner-first, so the inner call gets _0); two torch.relu(...) calls are
therefore torch_relu_0 and torch_relu_1. Every assignment is an
operation too, named {target}_{occurrence} — h = q @ k gives
h_0, whose .output is the assigned value — so a value that is not
a call's return (a matmul, a running state inside a loop) is reachable by the
name the forward gives it. Index in with that name to get a
SourceEnvoy::
model.layer1.source.torch_relu_0 # -> SourceEnvoy for the relu call
model.layer1.source.h_0 # -> SourceEnvoy for `h = ...`
You rarely need to memorize the names: print(model.layer1.source) renders
the whole forward with each operation labelled at its call site, and
print(model.layer1.source.torch_relu_0) zooms in on one. Iterating a
Source yields its operations in execution order.
Source values are only meaningful inside a trace: outside one every operation
calls straight through, which costs an idle model a few percent and changes
nothing about what it computes. Requesting an operation on a forward with no
recoverable Python source — a builtin or C function — raises
SourceNotAvailable; a
decorated forward is peeled and instrumented.
A Source also decomposes a called function — reached as
some_op.source (see SourceEnvoy.source) — the same way, one level
deeper. In that nested form the operations live under the drilled-into op's
path rather than the module's forward.
Examples:
Inspect, then capture and edit, an intermediate operation::
print(model.layer1.source) # list the operations
with model.trace(x):
h = model.layer1.source.torch_relu_0.output.save() # capture
model.layer1.source.self_fc2_0.input = h * 0 # edit a later op's input
Iterate every operation::
for op in model.layer1.source:
print(op.name)
__getattr__(name: str) -> SourceEnvoy
Resolve source.<name> to its SourceEnvoy.
Raises AttributeError for an unknown operation, listing the
available names so a mistyped or wrong-occurrence label is easy to fix.
Examples:
__iter__() -> Iterator[SourceEnvoy]
Iterate the operations in execution order, yielding a SourceEnvoy each.
Examples:
The whole forward source with every op labelled at its call site.
Each source line is shown with the operations that occur on it in a left
gutter (the def line is marked *); when several operations share a
line, the extras appear as + continuations. This is the map from
source code to {callable}_{occurrence} names, so you never have to
guess an occurrence number.
Examples:
The callee of an assignment operation: x = e runs as
x = __nnsight_op__("source.x_n", __nnsight_bind__, e), so the bracket
run_op puts around every call serves the assigned value as .output.
Parse the function code was compiled from: its module AST, first line, and text.
From the code object, not the function: given a function, inspect follows
__wrapped__ and hands back the decorated function's source instead of the
wrapper's. Raises SourceNotAvailable when there is nothing to read.
compile_source(func: Callable) -> Compiled
Parse, instrument, and compile a Python func, or raise.
The definition is compiled inside a shell function whose parameters are
func's free variables: recompiled at module level they would become
globals and break; as a nested definition they compile as free variables
again, and instrument attaches the original cells. The caller has peeled
func's decorators and rebuilds them around the result, so the @ lines
(which getsourcelines includes) are dropped rather than doubled.
Index of the closure cell holding the function wrapper decorates, or None.
A decorator's wrapper keeps the function it decorates in a closure cell and
calls it, so the cell is found from the wrapper's own source: the free names
it calls directly (fn(*args, **kwargs)) that hold a Python function.
Exactly one is the decorated function. None means the wrapper doesn't call
what it closes over — a dispatcher that hands the function to a lookup and
calls the result (transformers' experts wrapper, which runs a fused kernel
instead of the eager loop it wraps) — and several is ambiguous; either way
the wrapper is instrumented as it is, so the call that actually runs is what
shows up. Matching by __wrapped__ would peel the dispatcher too.
Peel func's decorators: the innermost function and the (wrapper, cell)
chain, outermost first, that rewrap rebuilds around its replacement.
Rebuild chain's decorators around innermost, inside out.
Each wrapper is rebuilt with a fresh closure rather than having its cell assigned: the wrapper is the class's attribute, shared by every instance in the process, so mutating its cell in place would redirect models nobody is tracing.
A new function with fn's globals, defaults and names but code and closure.
instrument(fn: Callable, op: Callable) -> tuple[Callable, Compiled]
A source-instrumented replacement for fn, calling op per operation,
and its Compiled — or raise SourceNotAvailable.
Peels fn's decorators, instruments the function they wrap, and rebuilds
them around it so their behaviour still runs. The instrumented copy shares
the function's closure cells, matched by name (the shell can order them
differently), so a wrapper keeps reaching what it closed over. A bound method
is rebuilt from its function and re-bound to the same instance.
A callable instance is instrumented through its __call__, which is where
its Python source is: a diffusers attention processor is a plain object, and
asking the instance for a code object said there was no source to read.
Bracket one operation at location base: input, skip, (recursive) run, output.
Reports/replaces .input, honors a .skip gate, and reports/replaces
.output — the same three handles a module's controller emits, one level finer.
Between them, if a worker asked to drill into this op (base is a key in
Interleaver.sourced), the raw fn is offered over {base}.fn so
the worker can hand back a source-instrumented copy (cached in
Interleaver.sourced for later fires); that copy runs in place of fn,
making its operations addressable under {base}.source.* — recursively.
Build the __nnsight_op__ an instrumented function calls at each operation.
locate answers "which trace is running, and under what path?" — for a
module's forward, its live State (so re-wrapping and multiple wrappers just
work); for a drilled-into callable, the interleaver and op path it was drilled
from. With no trace running the op calls straight through; inside one it
brackets the call via run_op under {path}.{location}.
Run the module's body, honouring accelerate's device-alignment hook.
accelerate.add_hook_to_module installs alignment by replacing
module.forward (instance __dict__) and keeping the real forward on
module._old_forward. We install our controller into that same slot, so its
wrapper is gone and pre_forward/post_forward would never run -- which
silently breaks any model sharded across devices, because the inter-module
tensor moves are exactly what those do. _hf_hook stays attached either way,
so the omission is invisible.
Bracketing the body here restores it, and works for both the original forward and the source-instrumented one.
The body the controller runs: the module's own forward if it has one,
otherwise its class's.
The controller takes the instance slot, so a forward already sitting there is
one it would otherwise destroy — self.forward = self._fast picked in
__init__, a monkeypatched layer, torch.compile's OptimizedModule —
which left the module raising NotImplementedError, or quietly running a
different implementation, for the rest of the process.
An instance forward is already bound, so it is wrapped to take (and ignore) the module the controller passes; where it is a method of this module its function is kept instead, which does the same without pinning the module the state lives on.
install_controller(envoy: 'Envoy') -> State
Install the controller forward on envoy's module once; (re)bind and return
its State.
Installed directly into the module's __dict__ (shadowing the class method
for __call__) and left there permanently — inert outside a trace. The body
is whatever the module's forward was (module_body);
install_source upgrades it.