Info
Last Execution: 2026-09-02
| Package | Version |
|---|---|
| nnsight | 0.8.0 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Batching¶
A single trace() can hold several tracer.invoke(...) blocks. Their inputs are combined into one batched forward, and each block's interventions see only its rows of every activation. This lets you run multiple prompts through the model in a single forward pass — and, when you want, move a value out of one invoke and into another, the basis of techniques like activation patching, where a representation from a clean prompt is transferred into a corrupt one.
How batching works
TransformersModel implements batching, so multiple prompts can share one forward pass. By default a plain NNsight model / Envoy cannot batch — a lone invoke is already the whole batch, and combining two or more raises NotImplementedError. To batch a custom model you implement two methods:
_batch_size(self, *inputs, **kwargs)— return the number of rows a set of inputs represents (0if it carries no data)._batch(self, invokes, fn)—invokesis a list of each invoke's(inputs, kwargs); combine them into the single batched(args, kwargs)thatfnis called with.
TransformersModel overrides both (tokenizing and padding each prompt, then stacking the rows); the base Envoy implements only the one-invoke passthrough.
from nnsight.modeling.transformers import TransformersModel
model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
Invokers¶
When you pass input directly to model.trace(x), an implicit invoke is created over the whole batch. For multiple prompts, open explicit tracer.invoke(...) blocks. Each contributes its rows to the batch, and each block's .output carries only that invoke's row(s).
The two styles are alternatives, not a combination: give the input straight to trace(...), or leave trace() empty and supply every input through invoke(...) blocks. A direct input has already started the run, so opening an invoke underneath it raises ValueError: Cannot invoke while the model is already running.
with model.trace() as tracer:
with tracer.invoke("The Eiffel Tower is in the city of"):
logits1 = model.lm_head.output.save()
with tracer.invoke("The Colosseum is in the city of"):
logits2 = model.lm_head.output.save()
print(f"Prompt 1: {model.tokenizer.decode(logits1[0, -1].argmax(dim=-1))}")
print(f"Prompt 2: {model.tokenizer.decode(logits2[0, -1].argmax(dim=-1))}")
Prompt 1: Paris Prompt 2: P
How invokers really run
Each invoke's body runs in its own greenlet worker (cooperative, single-threaded). Workers do not run strictly one-after-another: they all start together and resume in the order the model reaches what each asked for. That is what makes them a batch rather than a sequence. Within a single invoke, you must still access modules in forward-pass order, or you hit OutOfOrderError.
You can also batch several prompts inside a single invoke by passing a list. The saved activation then has one row per prompt:
with model.trace() as tracer:
with tracer.invoke(["The Eiffel Tower is in the city of", "The Colosseum is in the city of"]):
logits = model.lm_head.output.save()
print(f"Prompt 1: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")
print(f"Prompt 2: {model.tokenizer.decode(logits[1, -1].argmax(dim=-1))}")
Prompt 1: Paris Prompt 2: P
You don't even need an invoker for this — pass the list of prompts straight to trace([...]) and an implicit invoke batches them for you:
with model.trace(["The Eiffel Tower is in the city of", "The Colosseum is in the city of"]) as tracer:
logits = model.lm_head.output.save()
print(f"Prompt 1: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")
print(f"Prompt 2: {model.tokenizer.decode(logits[1, -1].argmax(dim=-1))}")
Prompt 1: Paris Prompt 2: P
Values from the enclosing scope¶
A value bound before or around the invokes flows into every invoke automatically — no synchronization needed, because it is already materialized when the workers start. This is the easy case: the shared value did not come from another invoke's activation.
import torch
steer = torch.zeros(768, device=model.device) # defined in the enclosing scope
with model.trace() as tracer:
with tracer.invoke("The Eiffel Tower is in the city of"):
a = model.lm_head.output.save()
with tracer.invoke("The Colosseum is in the city of"):
# `steer` came from outside the invokes, so it is visible with no barrier
model.transformer.h[5].output[:, -1, :] += steer
b = model.lm_head.output.save()
print(f"Shapes: {a.shape}, {b.shape}")
Shapes: torch.Size([1, 10, 50257]), torch.Size([1, 10, 50257])
Cross-prompt transfer needs a barrier¶
The interesting case is transferring a value one invoke produces from an activation into another invoke.
Each invoke's block runs until it asks for a value the model has not produced yet, which parks it; it resumes when the model reaches that location. A block can read a name a sibling bound once it has parked past the binding. The block below cannot: its first statement is the swap, and an assignment evaluates its right-hand side before the attribute access parks anything, so it reads embeddings having parked nowhere. wte is the first module in the model, so there is no earlier place for it to park either.
tracer.barrier(n) supplies the ordering instead. Everything written above a barrier happens before anything written below one, where n is the number of invokes that call barrier(). Here the first prompt's embeddings are transferred onto a second prompt made only of underscores, which then predicts as if it were the first prompt.
with model.trace() as tracer:
barrier = tracer.barrier(2) # two participating invokes
with tracer.invoke("The Eiffel Tower is in the city of"):
embeddings = model.transformer.wte.output
barrier() # signal: embeddings have been read
with tracer.invoke("_ _ _ _ _ _ _ _ _"):
barrier() # wait until invoke 1 has read its embeddings
model.transformer.wte.output = embeddings
logits = model.lm_head.output.save()
print(f"Prediction from transferred embeddings: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")
Prediction from transferred embeddings: Paris
When is a barrier required?
Whenever the consuming invoke cannot park past the producing one before it reads the value. That covers every consumer whose first statement is a write, and every consumer that has to act at or before the producer's module.
A consumer that only reads can buy its park with one extra .output access on a later module, and then needs no barrier. It is worth knowing, but a barrier is the safer default for anything that writes, since park-past breaks silently into a NameError if a line is inserted above the read.
A value from the enclosing scope, such as steer above, is already materialized when the workers start and needs nothing.
n must equal the number of invokes that actually call barrier(). Too high and the run ends with ValueError: A barrier was never reached by every block it waits for. Too low and the round releases early, and the invoke it lets through raises NameError on the value it was waiting for.
Activation patching¶
A practical application: paste an activation from a clean run into a corrupt run to measure that component's causal contribution. If the corrupt run then predicts the clean answer, the activation carried the deciding information.
Both prompts here are ten tokens long, which is what makes SUBJECT mean the same positions in both. A third, unpatched invoke gives a baseline in the same forward pass. Every invoke is padded out to the batch's longest input, so with prompts of different lengths the same slice would point at different tokens in each row (and at left padding in the shorter one). Indexing from the right, [:, -1], is the one that holds regardless.
Block outputs are plain tensors in current transformers, so we index the tensor directly (no tuple [0]).
clean = "The Eiffel Tower is in the city of" # next token: " Paris"
corrupt = "The Colosseum is in the city of" # next token: " Rome"
paris = model.tokenizer.encode(" Paris")[0]
rome = model.tokenizer.encode(" Rome")[0]
LAYER = 0
SUBJECT = slice(1, 5) # the subject tokens that differ between the two prompts
with model.trace() as tracer:
barrier = tracer.barrier(2)
with tracer.invoke(clean):
clean_hs = model.transformer.h[LAYER].output
barrier() # signal: clean_hs has been read
with tracer.invoke(corrupt):
barrier() # wait until clean_hs is materialized
hs = model.transformer.h[LAYER].output
hs[:, SUBJECT, :] = clean_hs[:, SUBJECT, :]
model.transformer.h[LAYER].output = hs
patched_logits = model.lm_head.output[:, -1, :].save()
with tracer.invoke(corrupt): # no barrier() -> not a participant
baseline_logits = model.lm_head.output[:, -1, :].save()
pb, pp = baseline_logits.softmax(-1), patched_logits.softmax(-1)
print(f"baseline corrupt: P(Paris)={pb[0, paris]:.3f} P(Rome)={pb[0, rome]:.3f}")
print(f"patched corrupt: P(Paris)={pp[0, paris]:.3f} P(Rome)={pp[0, rome]:.3f}")
baseline corrupt: P(Paris)=0.003 P(Rome)=0.014
patched corrupt: P(Paris)=0.064 P(Rome)=0.006
Patching the subject token at an early layer pushes the corrupt run toward the clean answer — the city information is read out early, at the subject position. Patch a late layer or the final position instead and the effect nearly vanishes; the position and layer you choose are the question you are asking.