nnsight

nnsight gives you access to every intermediate value in a model's forward pass — reading them, replacing them, backpropagating through them — for models you run locally and for models too large to fit on your machine (via NDIF).

When to activate

• The user wants activations, attention patterns, gradients, or any value from inside a model, not just its output
• The user wants to ablate, patch, steer, or otherwise edit a forward pass
• The user is writing or fixing code that imports nnsight
• The user mentions NDIF, remote=True, or running a model they cannot host

Orientation

import nnsight
from nnsight import TransformersModel

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

prompt = "The Eiffel Tower is in the city of"

with model.trace(prompt):
    clean = model.output.logits.save()                     # read the output

with model.trace(prompt):
    resid = model.transformer.h[5].output.save()           # read an activation
    model.transformer.h[8].output[:, -1, :] *= 2           # modify one
    edited = model.output.logits.save()                    # see what it changed

assert tuple(resid.shape) == (1, 10, 768)
assert model.tokenizer.decode(clean[0, -1].argmax()) == " Paris"
assert model.tokenizer.decode(edited[0, -1].argmax()) == " London"

Install: pip install nnsight (needs torch and transformers >= 5).

The five things that break agent-written nnsight code

1. .save() or it never existed. Assignments inside a trace body do not escape it — the body runs in a different frame. Bind what you save (x = ....save(); a bare ....save() on its own line returns nothing), and when collecting, save the container and append raw values:

with model.trace("The Eiffel Tower is in the city of"):
    per_layer = nnsight.save([])                  # save the list…
    for block in model.transformer.h:
        per_layer.append(block.output[0, -1])     # …append raw values

assert len(per_layer) == 12

There is no .value in 0.8 — the saved variable is the tensor.

2. Access modules in forward-pass order. Reading layer 8 then layer 2 raises OutOfOrderError: your code is a worker that parks until the model produces each value, and the model has already gone past. Within a block, the submodules (ln_1, attn, mlp) come before the block's own .output. This binds writes too — an edit at layer 0 goes above a read at layer 11, not below it.

3. Don't guess module paths or output types. model.transformer.h[i] is GPT-2; Llama is model.model.layers[i]; Gemma-3 is model.model.language_model.layers[i]. A block's .output is a plain tensor, but an attention submodule's is a tuple — so .output[0] copied from an old example silently selects batch row 0. Check before indexing. print(model) shows the module tree, and model.scan reports shapes without running the model or even downloading weights:

meta = TransformersModel("openai-community/gpt2")     # no dispatch=True

with meta.scan(prompt):
    hidden_size = nnsight.save(meta.transformer.h[0].output.shape[-1])
    n_layers = nnsight.save(len(meta.transformer.h))
    resid_shape = nnsight.save(tuple(meta.transformer.h[-1].output.shape))

print(hidden_size, n_layers, resid_shape)             # 768 12 (1, 10, 768)

Registration order (what print(model) shows) is not execution order: on Llama, self_attn is registered first but input_layernorm runs first.

4. One trace is one forward pass — structure by input, not by activation. Everything you want from one input comes out of one trace. N traces to fetch N layers is N forward passes (and, remotely, N network round-trips).

5. Old nnsight code is everywhere and it is wrong here. Most nnsight code on the internet predates 0.8. If you are adapting a tutorial or a paper repo, convert it first:

Pre-0.8	Write instead
LanguageModel(repo)	TransformersModel(repo, task="text-generation")
x.save() … x.value	drop .value — the saved name is the value
nnsight.list() / nnsight.dict()	nnsight.save([]) / nnsight.save({}) inside the trace
nnsight.apply(fn, x), nnsight.cond(c), nnsight.iter(...)	plain fn(x), if c:, for … in …:
tracer.next(), module.next()	for step in tracer.iter[:N]:
with tracer.iter[:3]: / with tracer.all():	for step in tracer.iter[:3]:
model.generator.output	tracer.result
nnsight.session()	model.session()

Picking a run method

You want	Use	You get back
one forward pass	model.trace(x)	the model's output object
generated tokens	model.generate(x, max_new_tokens=N)	token ids on tracer.result (the checkpoint's generation_config decides sampling — pass do_sample=False)
decoded text / labels	model.pipe(x, ...)	pipeline records (often sampled — pass do_sample=False)
shapes, no compute	model.scan(x)	fake tensors: shapes and dtypes only, never values
several traces sharing values	model.session()	values flow between traces
a permanent intervention	model.edit(inplace=True)	replayed on every later run

Core moves

Batch a sweep into one pass. Loop inside the trace, one tracer.invoke per variant, no input on trace():

prompt = "The Eiffel Tower is in the city of"
paris = model.tokenizer(" Paris").input_ids[0]

with model.trace() as tracer:
    scores = nnsight.save([])
    for layer in range(len(model.transformer.h)):
        with tracer.invoke(prompt):
            model.transformer.h[layer].output[:, -1, :] = 0
            scores.append(model.output.logits[0, -1, paris])

print([round(s.item(), 2) for s in scores])     # 12 ablations, one forward pass

Inside an invoke, index as if that input were alone — [:, -1, :] means "this invoke's rows, last position".

Grab many modules at once with tracer.cache() (call it first thing; it observes post-intervention values, and returns a list per module when the module runs more than once, as in generation):

with model.trace(prompt) as tracer:
    cache = tracer.cache(modules=[model.transformer.h[0], model.transformer.h[11]])

print(cache["model.transformer.h.0"].output.shape)

Gradients — capture in the forward, read in reverse order inside with metric.backward()::

with model.trace(prompt):
    hidden = model.transformer.h[-1].output
    metric = model.output.logits[0, -1, paris]
    with metric.backward():
        grad = hidden.grad.clone().save()

print(grad.shape)

Generation — a tracer.iter loop must not ask for a step the run does not make. A loop that outruns the run warns, keeps what it saved, and drops the statements after it — the result looks complete while being short. max_new_tokens is an upper bound, so pass min_new_tokens= when the bound has to hold, and check the len() of what you collected:

with model.generate(prompt, max_new_tokens=3, min_new_tokens=3) as tracer:
    picks = nnsight.save([])
    for step in tracer.iter[:3]:
        picks.append(model.output.logits[0, -1].argmax(dim=-1))
    ids = tracer.result.save()

assert len(picks) == 3
print(model.tokenizer.decode(ids[0]))

Running on NDIF

The same trace runs on NDIF against a model you can't fit locally — add remote=True and reduce metrics before saving, since every .save() is a download:

with model.trace(prompt, remote=True):
    last = model.transformer.h[-1].output[:, -1].detach().cpu().save()

print(last.shape)

An API key comes from login.ndif.us and is set with nnsight.CONFIG.set_default_api_key("...") or the NDIF_API_KEY environment variable; nnsight.status() lists deployed models. Anything beyond one trace (loops, sweeps, multi-step experiments) should be a model.session(remote=True) — one job instead of N round-trips.

Before you run an experiment

• Module paths and output types confirmed with print(model) / model.scan, not assumed
• Every value you need after the block is .save()d, containers not elements
• Module accesses in forward order; sweeps batched into invokes
• A control condition in the same trace (unmodified run) to compare against
• For generation: bounded tracer.iter[:N], and do_sample=False if you want reproducibility from pipe
• For large models: reduce metrics inside the trace before saving

Further reading

• Documentation and tutorials: https://nnsight.net/
• Per-technique skills (logit lens, activation patching, steering, probing, SAEs, vLLM, tensor parallelism, …) with every example executed in CI: https://github.com/ndif-team/skills
• NDIF, remote access to large models: https://ndif.us/