tracer¶
The tracer behind with model.trace(...):.
InterleavingTracer is the context manager a user gets from
Envoy.trace. It captures the body of the with block instead of
running it inline, then — on exit — runs that body as intervention code
interleaved with the model's forward pass (see
nnsight.intervention.interleaver).
Concretely, the tracer turns::
with envoy.trace(x):
l1 = envoy.l1.output.save()
into: run envoy(x) while a worker executes the block, parking on
envoy.l1.output until the model's l1 module produces its output, then
resuming with that value. Saved values survive once the with exits.
The capture/compile/execute plumbing lives in the base
Tracer; this module only overrides
execute to drive the interleaver, and adds the
trace-body API (stop,
result).
InterleavingTracer(envoy: Envoy, fn: Callable | str, *args: Any, backend: Backend | None = None, **kwargs: Any)
Bases: Tracer
Runs a with block as interventions interleaved with a model call.
Constructed by Envoy.trace; the fn and forward-pass arguments are
remembered here and only executed when the with block exits, at which
point execute hands the captured body to Envoy.interleave to
run alongside fn(*args, **kwargs) (usually the model's forward).
Inside the block, code reads and edits activations through Envoy
properties (envoy.l1.output, envoy.l2.input, ...) and can use the
tracer's own members:
result— the valuefnreturned (e.g. the model's output).stop— halt the model run early.
| ATTRIBUTE | DESCRIPTION |
|---|---|
envoy |
The
|
fn |
The callable to run interleaved, or its name on the module. A name
is left unresolved until
|
args |
Positional arguments forwarded to
|
kwargs |
Keyword arguments forwarded to
|
property
¶
iter: Iterations
Target specific occurrences of a location across a repeated run.
When the traced call reaches a module more than once — e.g. each step of
a generation loop — looping over tracer.iter binds reads and writes in
the loop body to a chosen range of those occurrences:
.. code-block:: python
with model.generate("...", max_new_tokens=10) as tracer:
for step in tracer.iter[:3]:
hidden = model.transformer.h[0].output.save() # steps 0, 1, 2
See Iterations for how the range is selected.
Halt the model run as soon as this point is reached.
Everything captured before the stop is kept; the model does not run past it, so later locations become unreachable::
with model.trace("The Eiffel Tower is in") as tracer:
hidden = model.transformer.h[0].output.save()
tracer.stop()
The value the traced call returned — e.g. the model's output.
For generate this is the token
ids, so read it inside the block and .save() it to keep it::
with model.generate("Madison Square Garden is in", max_new_tokens=3) as tracer:
ids = tracer.result.save()
barrier(n: int) -> Barrier
A meeting point for n of this trace's blocks.
The blocks of a trace run in the order the model reaches what each asked for, so one that hands something to another needs a point both agree on. Every block that holds the barrier calls it; the last to arrive releases them all, so everything above a barrier has happened before anything below one.
Examples:
>>> with model.generate(max_new_tokens=3) as tracer:
... barrier = tracer.barrier(2)
... with tracer.invoke("Madison Square Garden is in"):
... embeddings = model.transformer.wte.output
... barrier()
... with tracer.invoke("_ _ _ _ _"):
... barrier()
... model.transformer.wte.output = embeddings
| PARAMETER | DESCRIPTION |
|---|---|
n
|
How many blocks will call this barrier. Fewer than this and it never releases; the blocks left waiting report it when the run ends.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Barrier
|
A |
all() -> Iterations
Every occurrence — a shorthand for tracer.iter[:].
.. code-block:: python
with model.generate("...", max_new_tokens=10) as tracer:
for step in tracer.all():
hidden = model.transformer.h[0].output.save() # every step
cache(modules: list[Any] | None = None, device: Any = device('cpu'), dtype: Any = None, detach: bool = True, include_output: bool = True, include_inputs: bool = False, non_blocking: bool = False) -> CacheView
Record activations of many modules at once during the run.
Returns a CacheView that fills in as
the model runs; read captured values from it after the trace::
with model.trace(prompt) as tracer:
cache = tracer.cache() # every module's output
cache["model.transformer.h.0"].output # by path
cache.transformer.h[0].output # or by navigation
Declare the cache before reading or modifying a model value. Values are
recorded post-intervention; in a generation loop a module is captured once
per step (len(cache[path]) is the step count).
| PARAMETER | DESCRIPTION |
|---|---|
modules
|
Envoys or path strings to capture;
TYPE:
|
device
|
Device to move captured tensors to (default CPU);
TYPE:
|
dtype
|
Optional dtype to cast captured tensors to.
TYPE:
|
detach
|
Detach captured tensors from the autograd graph.
TYPE:
|
include_output
|
Capture each module's output.
TYPE:
|
include_inputs
|
Capture each module's inputs.
TYPE:
|
non_blocking
|
Use an async (non-blocking) device transfer (default
TYPE:
|
Add a batched input inside a with model.trace() as tracer: block.
Each with tracer.invoke(x): block contributes its input as one group of
rows in a single combined forward, and its body's interventions see only
those rows. Give trace() no input and use one invoke per prompt::
with model.trace() as tracer:
with tracer.invoke("the cat"):
a = model.transformer.h[0].output.save()
with tracer.invoke("a much longer prompt here"):
b = model.transformer.h[0].output.save()
An empty tracer.invoke() sees the whole batch (no row scoping).
Run the captured trace body interleaved with the (possibly batched) model call.
Two shapes, distinguished by whether trace() itself got input:
- direct input (
trace(x)) — the whole block is one implicit invoke; the block runs as a single worker overx. - invoke mode (
trace()) — the block is run now to collect itstracer.invoke(...)sub-blocks (each registers its input and captured body); their inputs are combined into one batched forward and each sub-block runs as a worker scoped to its rows.
Either way the workers are handed to Envoy.interleave to run
alongside fn(*combined_input), then results are pushed back with
save-gating (see Tracer.execute).
Return a clean traceback for an error raised inside the trace body.
When a worker raises, Mediator.switch stashes the intervention-only
traceback on the exception as __intervention_tb__ before the model and
model frames pile on top during unwinding. Prefer that stashed trace, fall
back to the live one, and filter nnsight's own frames out so the user sees
just their own code.
ScanningTracer(envoy: Envoy, fn: Callable | str, *args: Any, backend: Backend | None = None, **kwargs: Any)
Bases: InterleavingTracer
Like InterleavingTracer, but runs the forward under fake tensors.
Constructed by scan. The block
still reads activations through Envoy properties, but the model runs
inside a FakeTensorMode, so operations
only propagate tensor metadata (shape/dtype/device) — no real compute, no
real weights. This lets you inspect activation shapes on an undispatched
(meta-weight) model without loading it: Meta.interleave sees the active
fake mode and skips dispatch, leaving parameters on the meta device, and
Envoy.interleave moves the inputs onto that same device so shapes
propagate consistently.
The values read inside a scan are fake tensors: they carry metadata and never
hold data. Read .shape / .dtype and save those. A fake tensor saved
out of the block keeps working — .sum(), .mean() + 1, .cpu() all
hand back more fake tensors — so nothing tells you it is empty until something
asks for a real number (float(...), .numpy(), or an op mixing it with a
real tensor), which is usually far from where it came from.
Run InterleavingTracer.execute under a fake-tensor mode.
Deferring to super().execute means scan goes through the same
_batch_size/_batch preprocessing as trace (a string prompt is
tokenized, invokes are batched); wrapping it in a
FakeTensorMode makes the forward
propagate only shapes/dtypes — no real compute, no dispatch.
assume_static_by_default keeps shapes concrete rather than symbolic, and
allow_non_fake_inputs lets the meta-device parameters take part without
being faked first.
Bases: Tracer
One with tracer.invoke(x): block inside a with model.trace() as tracer:.
Captures its body and registers it — together with the batch group its input
occupies — on the parent InterleavingTracer, which combines all
invokes into a single batched forward (see InterleavingTracer.execute).
Register this invoke's input and body, deferring the run to the parent.
Adds the input to the parent tracer's batcher
(claiming a batch group) and appends the captured body as a worker on the
interleaver. The parent's InterleavingTracer.execute runs them all
together once every invoke has been collected.