Info
Last Execution: 2026-11-04
| Package | Version |
|---|---|
| nnsight | 0.8.2 |
| Python | 3.12.13 |
| torch | 2.15.0+cpu |
| transformers | 5.22.1 |
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, or a whole sweep of interventions, 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])
Sweeps¶
A sweep runs one intervention many times with one thing changed: the layer, the head, the neuron. Batched, each variant becomes a row instead of a forward pass of its own, and the two ways to batch give two ways to build one.
One invoke per variant. Loop inside the trace and open an invoke for each variant, with no input on trace(). Each invoke runs its own intervention code, so the variants can touch different modules. Here every block's MLP is switched off in turn, beside a clean baseline:
import nnsight
prompt = "The Eiffel Tower is in the city of"
paris = model.tokenizer.encode(" Paris")[0]
with model.trace() as tracer:
p_paris = nnsight.save({}) # one entry per variant
with tracer.invoke(prompt): # baseline: nothing switched off
p_paris["clean"] = model.lm_head.output[0, -1].softmax(-1)[paris]
for layer in range(len(model.transformer.h)):
with tracer.invoke(prompt): # this row: block `layer`'s MLP off
model.transformer.h[layer].mlp.output[:] = 0
p_paris[layer] = model.lm_head.output[0, -1].softmax(-1)[paris]
print(f"clean P(Paris) = {p_paris['clean']:.3f}")
for layer in range(len(model.transformer.h)):
print(f"MLP {layer:2d} off P(Paris) = {p_paris[layer]:.3f}")
clean P(Paris) = 0.070 MLP 0 off P(Paris) = 0.000 MLP 1 off P(Paris) = 0.058 MLP 2 off P(Paris) = 0.035 MLP 3 off P(Paris) = 0.043 MLP 4 off P(Paris) = 0.067 MLP 5 off P(Paris) = 0.100 MLP 6 off P(Paris) = 0.080 MLP 7 off P(Paris) = 0.092 MLP 8 off P(Paris) = 0.070 MLP 9 off P(Paris) = 0.053 MLP 10 off P(Paris) = 0.054 MLP 11 off P(Paris) = 0.128
Collect the results in a dict created with nnsight.save({}) before the loop and keyed by the variant. Invokes resume in the order the model reaches what each one asked for, so a list appended to from several invokes can come back in a different order than the loop ran. A key does not depend on that order.
One row per variant. When the variants differ only in an index, such as which neuron or which position, a single list input is enough. Repeat the prompt once per variant and give each row its own index: with rows = torch.arange(n), the edit act[rows, :, neurons] = 0 zeroes neuron neurons[i] in row i, at every position. This ranks all 3,072 neurons of the MLP whose removal hurt most above, 512 rows per forward pass:
LAYER = min(range(len(model.transformer.h)), key=lambda layer: p_paris[layer])
N_NEURONS = 3072 # GPT-2's MLP width, 4 x 768
CHUNK = 512 # rows per forward pass
rows = torch.arange(CHUNK)
effect = torch.empty(N_NEURONS, device=model.device)
with torch.no_grad():
for start in range(0, N_NEURONS, CHUNK):
neurons = torch.arange(start, start + CHUNK)
with model.trace([prompt] * CHUNK, logits_to_keep=1):
model.transformer.h[LAYER].mlp.act.output[rows, :, neurons] = 0 # row i: neuron start + i off
p = model.lm_head.output[:, -1].softmax(-1)[:, paris].save()
effect[neurons] = p - p_paris["clean"]
for n in effect.argsort()[:5].tolist():
print(f"MLP {LAYER} neuron {n:4d} off: P(Paris) {effect[n]:+.3f}")
MLP 0 neuron 2733 off: P(Paris) -0.068 MLP 0 neuron 2106 off: P(Paris) -0.026 MLP 0 neuron 2111 off: P(Paris) -0.008 MLP 0 neuron 2496 off: P(Paris) -0.007 MLP 0 neuron 390 off: P(Paris) -0.006
Each row is the same computation as making that one edit on its own, which is easy to check:
n = effect.argmin().item()
with model.trace(prompt):
model.transformer.h[LAYER].mlp.act.output[:, :, n] = 0
alone = model.lm_head.output[0, -1].softmax(-1)[paris].save()
print(f"neuron {n}, row of the sweep: {effect[n]:+.4f}")
print(f"neuron {n}, edited alone: {alone - p_paris['clean']:+.4f}")
neuron 2733, row of the sweep: -0.0676 neuron 2733, edited alone: -0.0676
How big a sweep fits. Every row costs the memory of a full forward pass, and a trace holds its whole batch at once. Invokes count the same way: 768 invokes in one trace are a 768-row batch, just like a 768-prompt list. What keeps a large sweep small:
- Save one reduced value per row, such as a probability, rather than whole activations or logits.
- Run under
torch.no_grad(), so saved values do not keep the autograd graph alive. - Pass
logits_to_keep=1totrace()when you read only the last position. The model then computes logits for that position instead of for every position. - Split the variants into chunks with one trace each, as above.
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.