interleaver¶
Interleaving intervention code with a model's forward pass.
nnsight lets you read and edit a model's intermediate values from ordinary
Python written inside a with model.trace(...): block. To make that work,
the intervention code and the model's forward pass have to run in lockstep:
the intervention pauses whenever it asks for a value the model hasn't produced
yet, the model runs until it reaches that value, hands it over, and the
intervention resumes — possibly editing the value on the way back in.
This module implements that dance with greenlets <https://greenlet.readthedocs.io>_
(cooperative, single-threaded coroutines), not OS threads:
-
Each block of intervention code becomes a
Mediator, which runs the code in its own greenlet — the "worker". The worker drives the interaction: it runs until it needs a value, then parks, switching control back to the greenlet that started it (the "parent", i.e. the model side). -
The worker parks by naming a location — a provider string such as
"model.layer1.output"or the run's"result". It parks to read a location (Mediator.value), to replace one (Mediator.swap), or to skip a gated computation (Mediator.skip). It can also park on no location at all, waiting on the other workers rather than the model (Mediator.barrier). -
An
Interleaverinstalls a controller on each of the model's modules. As the forward pass reaches each location, the controller callsInterleaver.handle(location, value), which offers the value to the workers and caches interested in that location. A worker waiting on it is served the value (read) or has its replacement substituted in (swap); the possibly edited value is returned back into the model's execution.
Because a worker and the model take strict turns on one thread, there are no
locks or queues — only greenlet switches. Each Mediator holds at most
one pending event at a time (the location it is currently parked on). A worker
must request locations in the order the model reaches them; asking for a
location the model already ran past raises OutOfOrderError.
Bases: Enum
What a parked worker is asking for.
A worker parks by switching a tuple (Event, location, ...) to its parent;
Mediator.handle inspects the first element to decide how to serve it.
See Mediator.value, Mediator.swap, Mediator.skip, and
Mediator.barrier for how each is raised from intervention code.
BARRIER is the odd one: it names no location, so the model side never
serves it — another worker does, on its way past the same barrier.
Bases: NamedTuple
What a worker is parked on, waiting for the model to reach.
The occurrence is its own field so that matching a visit is a plain
comparison; printing rejoins them ('model.layers.16.output.i2'), which is
the form worth reading in an error.
| ATTRIBUTE | DESCRIPTION |
|---|---|
event |
What the worker wants done at
TYPE:
|
provider |
The location, undecorated, or
TYPE:
|
iteration |
Which occurrence of
TYPE:
|
value |
The replacement a swap or skip carries; unused by a read.
TYPE:
|
Bases: Exception
Raised by an intervention to halt the model run early.
Thrown into the model's execution (e.g. via tracer.stop()) to unwind the
forward pass immediately. Interleaver.__exit__ swallows it, since the
early stop was intentional rather than a genuine error.
Bases: Exception
An intervention requested a location the model already ran past.
Workers must ask for locations in the order the model reaches them. If the
run finishes (or moves past a location) while a worker is still parked
waiting for it, Interleaver.check_dangling_mediators throws this into
the worker so the traceback points at the exact line that was waiting.
Mediator(code: Any, glbls: dict, lcls: dict, copy: bool = False, node: Any = None, shared: dict | None = None)
Runs one block of intervention code as a greenlet, in step with the model.
A mediator wraps one captured block — the body of a with block, or one
registered edit — and runs it inside a greenlet, the "worker". The worker
drives the interaction: it runs until the intervention asks for a value, then
parks, recording that pending request in pending and switching control
back to the parent greenlet (the model side). The parent later resumes it
through switch / handle.
The classmethods value, swap, skip, and
barrier are the API the intervention code calls to park
(Envoy properties like .output and .input are thin wrappers
over them). start, switch, and handle are the
parent-side machinery that runs and feeds the worker. current is how
code inside a worker finds the mediator driving it.
The block and its scope travel; everything the run builds does not — see
__getstate__, which is how an edit rides to a remote server.
| ATTRIBUTE | DESCRIPTION |
|---|---|
code |
The captured block, compiled. Executed by the worker.
|
glbls |
The globals the block was written against.
|
lcls |
The
|
copy |
Whether to exec against a fresh copy of
|
node |
The block's AST node, kept so the mediator can serialize.
|
interleaver |
The run this worker belongs to, set in
TYPE:
|
batch_group |
This worker's
TYPE:
|
worker |
The greenlet running
TYPE:
|
pending |
What the worker is currently parked on — a
TYPE:
|
iteration |
Which occurrence of a location the worker currently wants —
the occurrence its pending request is matched under.
TYPE:
|
base |
The interleaver's per-location counts when this worker started;
its own occurrence of a location is the interleaver's count minus this
(see
TYPE:
|
caches |
The caches this worker's
TYPE:
|
property
¶
Whether the worker exists and still has intervention code left to run.
False before start (no worker yet) and after the worker
finishes — a greenlet is falsy once it has run to completion — so this is
only True while the worker is parked mid-intervention.
classmethod
¶
The mediator whose worker is running now.
Only intervention code has one, and intervention code only runs while
interleaving — so no worker means what was asked for outside a run,
and there is nothing to park on and nothing to answer with.
Raised as a ValueError rather than the AttributeError that
reaching for the absent worker gives: from a property like
output, an AttributeError is
taken for "no such attribute" and comes back out of __getattr__ as one,
naming the property instead of the reason.
classmethod
¶
event(event: Event, location: str, *rest: Any) -> Any
Raise an event from inside a worker and return what's sent back.
Called on the worker side (from intervention code): switch to the parent
greenlet — the interleaver driving the model — handing it the event tuple,
and block until the parent switches a value back in. This is the
counterpart to switch, which drives the worker from the parent.
location is tagged .i{n} with the occurrence the worker wants, so
handle can bind it with a single match. When pinned
(iteration is an int), that's the pinned step. When relaxed
(None), it's the mediator's current count for this location — the next
occurrence the model hasn't handled yet — so the request follows the model
sequentially.
How many times the model has reached location since this worker started.
classmethod
¶
Read the value at location from inside a worker.
Parks until the interleaver reaches location — a module input/output
path (e.g. "model.h.0.output") or the run's "result" — then returns
the value produced there.
classmethod
¶
Replace the value at location from inside a worker.
Parks like value, but when the interleaver reaches location it
substitutes value for what the model produced (see handle), then
resumes the worker. Reading then swapping the same location works — both
events are drained in one handle.
classmethod
¶
Skip the computation gated at location, using value as its result.
Parks like swap, but targets a .skip gate that a module's (or
op's) forward wrapper queries before running — so value is returned
in place of running the computation, not after. A distinct event from SWAP
so skip-specific behavior can hang off it later.
classmethod
¶
Create the worker greenlet and run it up to its first park.
Switching into a fresh greenlet runs the captured block until it first
parks on a location (or finishes). Whatever it parks on becomes
pending, ready for handle to serve once the model reaches
that location. Per-run counters are reset here so a stored edit mediator,
replayed on a later trace, starts clean. interleaver is the run this
worker belongs to (it reads batch scoping from interleaver.batcher).
Resume the worker with args; return the next event it parks on.
Switches control into the worker greenlet, handing it args as the
return value of whatever park call it was blocked in, and blocks until
the worker parks again (returning its new event tuple) or finishes
(returning None). If the worker raises, its traceback is stashed on
the exception as __intervention_tb__ — a clean, intervention-only
trace captured before the re-raise unwinds the model's own stack on top —
and the exception propagates, halting the run.
Re-point the worker's parent at whoever is switching in now, so both its
return paths — parking (worker.parent.switch(...)) and finishing (a
greenlet auto-returns to its parent) — go back here rather than to a fixed
greenlet. This keeps the chain correct when a worker is served from inside
another worker's greenlet, e.g. an Envoy.__call__(hook=True) adapter
run whose submodule controllers serve a second worker mid-call.
Drain the worker's events parked on this visit to provider; return the value.
A read (Event.VALUE) is served value; a swap
(Event.SWAP) replaces value with the worker's. The worker may
do both in turn (read a location, then assign it), so loop until it parks
somewhere else or finishes. The returned value flows back up through
Interleaver.handle to the controller, which substitutes it into the run.
A location can be reached many times in one run — e.g. a module revisited
on every step of a generation loop. This visit is the
occurrence(provider)-th, so it serves the workers waiting for that
occurrence of it. A worker parks already carrying the occurrence it wants
(see event) — pinned to a step, or resolved to the next
occurrence when relaxed — so this is a location match and an integer
match: a request pinned to a later step doesn't match yet and waits while
earlier visits pass by. With no tracer.iter the occurrence is always
0, so every request binds to the first visit. Once a pinned non-zero
step is hit, the mediator is relaxed to None so the rest of that
step's requests follow the model sequentially rather than re-forcing the
index.
Drives the model side of interleaving: model handoffs in, workers served.
An interleaver owns the module controllers that turn a model's forward pass
into a stream of handle calls, and the list of Mediator workers
those calls feed. One interleaver is shared across an Envoy tree, so
every module's controller reports into the same set of workers.
Lifecycle of a run (see Envoy.interleave):
- A
Mediatoris appended tomediatorsfor each intervention block and each registered edit. - Entering the interleaver (
with interleaver:) flipsinterleavingon andstart\ s every worker so each parks on its first requested location. - The model runs. Each module's controller (see
instrument) callshandle, serving reads and applying swaps for any worker parked there, and returns the (possibly edited) value into the forward pass. check_dangling_mediatorssurfaces any worker still waiting for a location the model never reached (OutOfOrderError), andcancelclears the workers so the next run starts clean.
| ATTRIBUTE | DESCRIPTION |
|---|---|
mediators |
The workers to serve this run.
TYPE:
|
batcher |
The
TYPE:
|
interleaving |
|
sourced |
Op-location -> the instrumented callable a worker drilled into
(see
TYPE:
|
fragments |
A
TYPE:
|
instance-attribute
¶
instance-attribute
¶
observers: dict[str, list[tuple[Mediator, Any, tuple[str, str]]]] = {}
Rebuild the parked set, cache routes and busy from mediators.
Called by __enter__ once the workers have started, and by a driver that
replaces the list while a run is in progress (a scheduler reshuffling
which workers are in the batch), since nothing else notices the swap.
__enter__() -> Interleaver
Begin interleaving: arm the controllers and start each not-yet-started worker.
Only a worker with no greenlet yet (worker is None) is started; one that
already has a worker is left as is — parked mid-run on a re-entered interleaver.
The gate tests worker rather than alive so that a worker
whose block has finished is also left alone: a finished greenlet is falsy, so
an alive gate would take it for never-started and rerun its whole block.
End interleaving, swallowing an intentional early stop.
Returning True for EarlyStopException suppresses it: an
intervention asked to halt the run and has already unwound the model, so
it is not an error. Any other exception propagates.
instrument(envoy: Envoy) -> None
Route an envoy's module through this interleaver.
Installs the module's controller (see
install_controller), which
hands the module's input and output to handle under the locations
"{path}.input" and "{path}.output" while
interleaving, and gates
.skip. No forward hooks: a module with none is called on PyTorch's
fast path, which is what keeps an instrumented model's cost near zero
outside a trace. Also registers this interleaver on the module, so a
module can be skipped or source-drilled by this trace — and by another
envoy sharing the module at the same time.
Route value to this provider's consumers; return it, edited if any
intervention wrote to this location.
The provider index picks only workers parked here and caches that selected it. The visit is counted once, on the interleaver, after everyone parked on it has been served — so a worker that arrives here during the visit (released from a barrier by one that was served) asks for this occurrence and is served in it too, whichever order the workers were written in.
An EarlyStopException
from a worker is held back and re-raised on the way out, once the visit is
finished, so a tracer.stop() halts what follows the location it fires
at rather than the location itself: it is still counted, still assembled,
still recorded by the caches observing it, and the other workers parked on
it — a sibling invoke of the same batched run — are still served. Any other
exception propagates immediately, unless defer_exceptions is set, in
which case it is recorded on its own worker and the run continues.
Surface any worker still parked after the run.
Called once the model has finished. A worker that is still alive
was waiting for a location the model never reached.
dangling_unwind decides
what that means — an out-of-order read, which stands as an error, or a
tracer.iter loop that outran the run, bounded and open alike, which
warns — and this is the local half of it: what stands is raised, out of
the worker's own frame, and what is expected is warned about.
Drop all mediators and the batcher so the next run starts clean.
A worker still alive —
parked mid-intervention because the model's forward raised before it
reached the location — is unwound first. Dropping the reference does not
end a greenlet: a parked one keeps its frame, the frame keeps the block's
scope, and the scope keeps the model, so a run that errors would hold the
weights forever. The throw runs the block's finally blocks on the way
out; an exception raised there only warns, because cancel runs in the
driver's finally with the error that ended the run already in flight,
and that error is the one worth surfacing.
Each mediator's worker greenlet is released too, so a stored edit mediator
replayed on a later trace is seen as never-started (worker is None) and
restarts fresh rather than being skipped for still holding its finished
greenlet. Surfacing dangling mediators is a separate concern (see
check_dangling_mediators), handled by the driver after a run.
How a worker still parked when its run is over should be ended.
Returns the error to throw into the worker — the throw unwinds it, running its
finally blocks, and points the traceback at the line that was waiting — and
the warning to emit instead of surfacing that error, or None when the
error stands. Two drivers ask, and they differ only in what surfacing means:
Interleaver.check_dangling_mediators
raises the error, and vLLM's Requests.finish_dangling carries it home as the
request's deferred error. The policy is one copy here because two copies of it
can disagree, and a block that is an error locally and a warning on the engine
is a block whose meaning depends on where it runs.
Two cases:
- Out of order — a plain request (
iteration == 0) for a location the model already ran past, or never called. A real error. - A loop that outran the run —
tracer.iter[...]asked for a step the run did not make. An opentracer.iter[:]/tracer.all()has no end of its own, so outrunning the model is how it finishes; a bounded loop cut short by an EOS or a stop string ends the same way. That one warns: values from the steps that were reached are already saved, and the statements after the loop do not run.
Call this before the throw: unwinding the worker's loop restores both the pin and the loop's shape it reads.