Info

Last Execution: 2026-09-03

Package	Version
nnsight	0.8.0
Python	3.12.13
torch	2.13.0+cu126
transformers	5.15.0.dev0

Multiple token generation¶

model.generate(...) runs multi-token, autoregressive generation: one forward pass per new token, with the generated token ids coming back on tracer.result. Interventions in the block run against every forward the decode loop makes, so nnsight gives you a way to say which of those forwards you mean:

• for step in tracer.iter[...] binds the loop body to the steps you name, as a slice, an int, or a list of indices.
• tracer.all() is tracer.iter[:]: every step, however many the model turns out to run.

One rule governs both, and this notebook is arranged around it: a loop must not ask for a step the run does not make. A loop that stays inside the run keeps whatever you wrote after it. A loop that outruns the run is cut short at the loop, and everything below it in the block is discarded.

Setup¶

In [1]:

import nnsight
from nnsight.modeling.transformers import TransformersModel

model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)

import nnsight from nnsight.modeling.transformers import TransformersModel model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)

Basic generation¶

Use .generate() instead of .trace() for multi-token generation, and pass max_new_tokens to bound it. The generated token ids are on tracer.result, a [batch, seq] tensor holding the whole prompt plus the completion.

In [2]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    ids = tracer.result.save()

print(ids.shape)
print(model.tokenizer.decode(ids[0]))

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: ids = tracer.result.save() print(ids.shape) print(model.tokenizer.decode(ids[0]))

torch.Size([1, 15])
The Eiffel Tower is in the city of Paris, and the E

Greedy by default

generate goes through the model, using the checkpoint's own generation settings, so it is greedy by default — repeat runs give identical ids. Ask for sampling explicitly with model.generate(..., do_sample=True, top_k=50); any keyword is forwarded to the model's generate. If you want the task pipeline's decoded records (text, labels) instead of token ids, use model.pipe(...).

Looping over steps¶

To intervene on individual steps, or collect a value from each one, loop with for step in tracer.iter[:N]. step is the real integer step index. The 5 below is the number of steps this generation makes, so the loop finishes on its own and tracer.result after it still runs.

To gather per-step values, save a container once with nnsight.save([]) and append the raw values inside the loop. Do not call .save() on each individual value.

In [3]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    tokens = nnsight.save([])

    for step in tracer.iter[:5]:
        tokens.append(model.lm_head.output[0, -1].argmax(dim=-1))

    ids = tracer.result.save()

for i, t in enumerate(tokens):
    print(f"Step {i}: {model.tokenizer.decode(t)}")

print("Full output:", model.tokenizer.decode(ids[0]))

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: tokens = nnsight.save([]) for step in tracer.iter[:5]: tokens.append(model.lm_head.output[0, -1].argmax(dim=-1)) ids = tracer.result.save() for i, t in enumerate(tokens): print(f"Step {i}: {model.tokenizer.decode(t)}") print("Full output:", model.tokenizer.decode(ids[0]))

Step 0:  Paris
Step 1: ,
Step 2:  and
Step 3:  the
Step 4:  E
Full output: The Eiffel Tower is in the city of Paris, and the E

Targeting specific steps: slice, int, and list¶

Subscript tracer.iter to name the steps you want:

• a slice, tracer.iter[1:3], runs the body for steps 1 and 2 (stop is exclusive);
• an int, tracer.iter[0], runs just that one step (0 is the prefill step, which processes the whole prompt);
• a list, tracer.iter[[0, 2, 4]], runs only those steps.

Each form is a claim that the run reaches those steps. Start with the slice:

In [4]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    tokens = nnsight.save([])

    for step in tracer.iter[1:3]:
        tokens.append(model.lm_head.output[0, -1].argmax(dim=-1))

    ids = tracer.result.save()

print(f"Full output: {model.tokenizer.decode(ids[0])}")
print(f"Collected {len(tokens)} tokens from steps 1-2")

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: tokens = nnsight.save([]) for step in tracer.iter[1:3]: tokens.append(model.lm_head.output[0, -1].argmax(dim=-1)) ids = tracer.result.save() print(f"Full output: {model.tokenizer.decode(ids[0])}") print(f"Collected {len(tokens)} tokens from steps 1-2")

Full output: The Eiffel Tower is in the city of Paris, and the E
Collected 2 tokens from steps 1-2

An int targets a single step. tracer.iter[0] is the prefill step — the one forward pass that ingests the whole prompt:

In [5]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    prefill_pick = nnsight.save([])

    for step in tracer.iter[0]:
        prefill_pick.append(model.lm_head.output[0, -1].argmax(dim=-1))

    ids = tracer.result.save()

print("Step 0 predicts:", model.tokenizer.decode(prefill_pick[0]))

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: prefill_pick = nnsight.save([]) for step in tracer.iter[0]: prefill_pick.append(model.lm_head.output[0, -1].argmax(dim=-1)) ids = tracer.result.save() print("Step 0 predicts:", model.tokenizer.decode(prefill_pick[0]))

Step 0 predicts:  Paris

A list targets an explicit set of steps — here just steps 0, 2, and 4:

In [6]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    picks = nnsight.save([])

    for step in tracer.iter[[0, 2, 4]]:
        picks.append(model.lm_head.output[0, -1].argmax(dim=-1))

    ids = tracer.result.save()

print("Steps 0, 2, 4:", [model.tokenizer.decode(t) for t in picks])

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: picks = nnsight.save([]) for step in tracer.iter[[0, 2, 4]]: picks.append(model.lm_head.output[0, -1].argmax(dim=-1)) ids = tracer.result.save() print("Steps 0, 2, 4:", [model.tokenizer.decode(t) for t in picks])

Steps 0, 2, 4: [' Paris', ' and', ' E']

Conditional interventions per step¶

Because step is a plain integer, an ordinary Python if lets you apply different interventions at different points in generation.

In [7]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    for step in tracer.iter[:5]:
        if step == 0:
            # Only intervene on the first (prefill) step
            model.transformer.h[0].output[:] = 0

    ids = tracer.result.save()

print(model.tokenizer.decode(ids[0]))

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: for step in tracer.iter[:5]: if step == 0: # Only intervene on the first (prefill) step model.transformer.h[0].output[:] = 0 ids = tracer.result.save() print(model.tokenizer.decode(ids[0]))

The Eiffel Tower is in the city of, but I'm not

Applying interventions to every step¶

To run the same intervention on every step, loop over the whole range. Here we zero-ablate layer 0 on every step while collecting the last layer's hidden state each time.

In [8]:

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

    hidden_states = nnsight.save([])

    for step in tracer.iter[:5]:
        model.transformer.h[0].output[:] = 0
        hidden_states.append(model.transformer.h[-1].output)

    ids = tracer.result.save()

print(f"Collected {len(hidden_states)} hidden states")
print(f"Shapes: {[tuple(h.shape) for h in hidden_states]}")
print(f"Output: {model.tokenizer.decode(ids[0])}")

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: hidden_states = nnsight.save([]) for step in tracer.iter[:5]: model.transformer.h[0].output[:] = 0 hidden_states.append(model.transformer.h[-1].output) ids = tracer.result.save() print(f"Collected {len(hidden_states)} hidden states") print(f"Shapes: {[tuple(h.shape) for h in hidden_states]}") print(f"Output: {model.tokenizer.decode(ids[0])}")

Collected 5 hidden states
Shapes: [(1, 10, 768), (1, 1, 768), (1, 1, 768), (1, 1, 768), (1, 1, 768)]
Output: The Eiffel Tower is in the city of,,,,,

Why the first hidden state is larger

The first step processes the full prompt (10 tokens), so its hidden state has shape [1, 10, 768]. Each later step processes a single new token, giving [1, 1, 768]. A GPT-2 block's .output is the plain hidden-state tensor, batch dimension included, so read and write it whole rather than indexing [0].

When a loop outruns the run¶

A bound is a claim about how many steps the run makes. When the run makes fewer, the loop is left waiting on a step that never arrives, and nnsight cuts the loop short there: what it saved is kept, the statements after the loop are discarded, and the only signal is a warning. The result looks complete — five entries, no error — while holding fewer steps than the bound named, so check the len() of what you collected:

In [9]:

import warnings

ids = None
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")

    with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

        picks = nnsight.save([])

        for step in tracer.iter[:8]:          # 8 steps asked of a 5-step generation
            picks.append(model.lm_head.output[0, -1].argmax(dim=-1))

        ids = tracer.result.save()

print(f"Steps collected: {len(picks)}")       # 5, not the 8 the bound named
print(f"ids bound: {ids is not None}")        # the save after the loop never ran
print("Warning:", caught[0].message)

import warnings ids = None with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: picks = nnsight.save([]) for step in tracer.iter[:8]: # 8 steps asked of a 5-step generation picks.append(model.lm_head.output[0, -1].argmax(dim=-1)) ids = tracer.result.save() print(f"Steps collected: {len(picks)}") # 5, not the 8 the bound named print(f"ids bound: {ids is not None}") # the save after the loop never ran print("Warning:", caught[0].message)

Steps collected: 5
ids bound: False
Warning: 'model.lm_head.output.i5' was never reached: the loop asked for a step the run did not make, so it was cut short — values saved inside the loop are kept, and the statements after it did not run. An open `tracer.iter[:]` / `tracer.all()` loop ends this way by design. To hold a generation to a bounded loop's count, pass `min_new_tokens=` on transformers or `min_tokens=` / `ignore_eos=True` on vLLM; put what follows the loop in a separate `tracer.invoke()`.

The warning names the first step the run did not make — .i5, the sixth ask — which is usually enough to fix the bound.

Getting the number right is harder than it looks, because max_new_tokens is an upper bound. An end-of-sequence token stops generation wherever it appears, and a loop bound to max_new_tokens then outruns the run. Below, gpt2's greedy continuation is " Paris, and the Eiff"; treating " and" as the end-of-sequence token makes the run stop after three steps, the way a real EOS would.

In [10]:

early_eos = model.tokenizer.encode(" and")[0]

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")

    with model.generate("The Eiffel Tower is in the city of",
                        max_new_tokens=6, eos_token_id=early_eos) as tracer:

        picks = nnsight.save([])

        for step in tracer.iter[:6]:
            picks.append(model.lm_head.output[0, -1].argmax(dim=-1))

print(f"Steps collected: {len(picks)}")       # the EOS ended the run after 3
print("Warned:", any("was never reached" in str(w.message) for w in caught))

early_eos = model.tokenizer.encode(" and")[0] with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") with model.generate("The Eiffel Tower is in the city of", max_new_tokens=6, eos_token_id=early_eos) as tracer: picks = nnsight.save([]) for step in tracer.iter[:6]: picks.append(model.lm_head.output[0, -1].argmax(dim=-1)) print(f"Steps collected: {len(picks)}") # the EOS ended the run after 3 print("Warned:", any("was never reached" in str(w.message) for w in caught))

Steps collected: 3
Warned: True

min_new_tokens=N suppresses the end-of-sequence token until N tokens have been generated, which is what turns a bound of N into something the run has to honour. The continuation changes as a result: with " and" held back for six steps, the model writes something else instead.

In [11]:

with model.generate("The Eiffel Tower is in the city of",
                    max_new_tokens=6, min_new_tokens=6, eos_token_id=early_eos) as tracer:

    picks = nnsight.save([])

    for step in tracer.iter[:6]:
        picks.append(model.lm_head.output[0, -1].argmax(dim=-1))

    ids = tracer.result.save()

print(f"Steps collected: {len(picks)}")
print(model.tokenizer.decode(ids[0]))

with model.generate("The Eiffel Tower is in the city of", max_new_tokens=6, min_new_tokens=6, eos_token_id=early_eos) as tracer: picks = nnsight.save([]) for step in tracer.iter[:6]: picks.append(model.lm_head.output[0, -1].argmax(dim=-1)) ids = tracer.result.save() print(f"Steps collected: {len(picks)}") print(model.tokenizer.decode(ids[0]))

Steps collected: 6
The Eiffel Tower is in the city of Paris, France.

When the step count is unknown: tracer.all()¶

tracer.all() is shorthand for the open-ended tracer.iter[:]. Rather than naming a step count, it keeps handing out step indices until the model stops generating, which makes it the right form when you cannot know the length ahead of time.

Ending that way is the same over-run as above — the same warning, the same cut — but for an open loop it is not a mistake: it has no bound of its own to be wrong about, and outrunning the run is how it finishes. Below we capture the warning to print it; normally it goes to stderr.

In [12]:

import warnings

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")

    with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer:

        hidden_states = nnsight.save([])

        for step in tracer.all():
            model.transformer.h[0].output[:] = 0
            hidden_states.append(model.transformer.h[-1].output)

        after_the_loop = nnsight.save("this line never runs")

print(f"Collected {len(hidden_states)} hidden states")
print(f"after_the_loop defined: {'after_the_loop' in dir()}")
for warning in caught:
    print("Warning:", warning.message)

import warnings with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") with model.generate("The Eiffel Tower is in the city of", max_new_tokens=5) as tracer: hidden_states = nnsight.save([]) for step in tracer.all(): model.transformer.h[0].output[:] = 0 hidden_states.append(model.transformer.h[-1].output) after_the_loop = nnsight.save("this line never runs") print(f"Collected {len(hidden_states)} hidden states") print(f"after_the_loop defined: {'after_the_loop' in dir()}") for warning in caught: print("Warning:", warning.message)

Collected 5 hidden states
after_the_loop defined: False
Warning: 'model.transformer.h.0.output.i5' was never reached: the loop asked for a step the run did not make, so it was cut short — values saved inside the loop are kept, and the statements after it did not run. An open `tracer.iter[:]` / `tracer.all()` loop ends this way by design. To hold a generation to a bounded loop's count, pass `min_new_tokens=` on transformers or `min_tokens=` / `ignore_eos=True` on vLLM; put what follows the loop in a separate `tracer.invoke()`.

The values collected inside the loop are all there. The statement after the loop is not: the unwind that ends the loop takes the rest of the block with it, so after_the_loop was never assigned.

Anything that has to happen after an open loop goes in its own tracer.invoke(). An invoke is a separate worker on the same batch, so the loop's unwind does not reach it. The loop still ends the same way, so the warning still appears; what changes is that the second invoke's code survives it.

In [13]:

with model.generate(max_new_tokens=5) as tracer:

    with tracer.invoke("The Eiffel Tower is in the city of"):
        hidden_states = nnsight.save([])
        for step in tracer.all():
            hidden_states.append(model.transformer.h[-1].output)

    with tracer.invoke():
        ids = tracer.result.save()

print(f"Collected {len(hidden_states)} hidden states")
print(f"Output: {model.tokenizer.decode(ids[0])}")

with model.generate(max_new_tokens=5) as tracer: with tracer.invoke("The Eiffel Tower is in the city of"): hidden_states = nnsight.save([]) for step in tracer.all(): hidden_states.append(model.transformer.h[-1].output) with tracer.invoke(): ids = tracer.result.save() print(f"Collected {len(hidden_states)} hidden states") print(f"Output: {model.tokenizer.decode(ids[0])}")

Collected 5 hidden states
Output: The Eiffel Tower is in the city of Paris, and the E

/home/localjadenfk/wd/nnsight/src/nnsight/intervention/interleaver.py:859: UserWarning: 'model.transformer.h.11.output.i5' was never reached: the loop asked for a step the run did not make, so it was cut short — values saved inside the loop are kept, and the statements after it did not run. An open `tracer.iter[:]` / `tracer.all()` loop ends this way by design. To hold a generation to a bounded loop's count, pass `min_new_tokens=` on transformers or `min_tokens=` / `ignore_eos=True` on vLLM; put what follows the loop in a separate `tracer.invoke()`.
  warnings.warn(expected)

What decides whether code after the loop runs

A tracer.iter loop must not ask for a step the run does not make.

• A bound the run meets ends the loop normally, and the rest of the block runs.
• A loop that asks for a step the run does not make — a bound the run cannot meet, or an open tracer.iter[:] / tracer.all() reaching its natural end — warns, keeps everything saved inside the loop, and drops the statements after it.
• A cut-short bounded loop hands back fewer values than its bound named, with nothing raised. Check the len() of what you collected, or hold the run to the count.

max_new_tokens is an upper bound, so pass min_new_tokens= when a bound has to hold, or loop openly and put the trailing code in a separate tracer.invoke(). tracer.result in particular has to be read inside the block: it is served during the run, so reading it below the with block raises instead.

Calling generate directly¶

Without a with block, generate just runs and returns the token ids as a tensor.

In [14]:

ids = model.generate("The Eiffel Tower is in the city of", max_new_tokens=3)

print(model.tokenizer.decode(ids[0]))

ids = model.generate("The Eiffel Tower is in the city of", max_new_tokens=3) print(model.tokenizer.decode(ids[0]))

The Eiffel Tower is in the city of Paris, and