Mechanistic Interpretability Transformers Tooling

Poking Inside a Language Model with NNsight

By Aayushman

September 12, 2026 · 9 min read

Most mechanistic interpretability questions have the same shape: does a specific piece of the network — this attention head, this MLP neuron, this residual stream position at this layer — causally matter for a specific behavior? Answering that with raw PyTorch means registering forward hooks, remembering to remove them, threading activations captured at one point in the network into an intervention at another, and rewriting most of that plumbing every time you change models. NNsight, from the NDIF project, replaces the plumbing with a single abstraction — a tracing context that lets you read and write a model's internals as if the forward pass had already happened, even though nothing has run yet.

The abstraction: a lazy intervention graph

The entry point is a context manager around a forward pass: 

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The capital of France is") as tracer:
    hidden = model.transformer.h[6].output[0].save()
    logits = model.lm_head.output.save()

print(hidden.shape, logits.shape)

Nothing inside the with block executes immediately. Writing model.transformer.h[6].output doesn't hand you a tensor — it hands you a proxy, an object standing in for "the output of this module, once the trace actually runs." Indexing it, calling .save() on it, doing arithmetic with it: all of that gets recorded onto an intervention graph rather than executed against real data. Only when the with block exits does NNsight run the actual forward pass once, feed real tensors through that recorded graph, and materialize whatever you called .save() on. This is the whole trick, and it's what makes everything downstream easy: you get to write intervention code that reads like it's manipulating activations directly, while what you're actually doing is specifying a graph that gets compiled against the model exactly once.

Why the laziness matters beyond convenience

Building the graph symbolically before running anything has a consequence that isn't just ergonomic. Because the trace is a description of what to compute rather than a sequence of side effects on a live Python object, the same code can target a model running on someone else's hardware. NDIF hosts large open-weight models remotely and executes submitted intervention graphs against them, sending back only the small activations you asked to save — not the full hidden state at every layer, and not the model's weights to your machine. You write the same model.trace(...) block whether model is a small model on your own GPU or a large one served remotely; the interpretability code doesn't change, only where it executes.

A minimal logit lens

One of the oldest mechanistic-interpretability tricks is the logit lens: take the residual stream at an intermediate layer, run it through the model's own unembedding matrix as if it were the final layer, and read off what the model would "say" if it stopped computing right there. With a tracing context this is a few lines rather than a custom forward hook that has to know exactly where the final layer norm and unembedding live: 

with model.trace("The capital of France is") as tracer:
    layer_logits = []
    for layer in model.transformer.h:
        normed = model.transformer.ln_f(layer.output[0])
        layer_logits.append(model.lm_head(normed).save())

# layer_logits[i] now holds the vocabulary distribution the model would
# produce if generation stopped after block i, for every layer at once.

Decoding the top token from each saved distribution and watching it change across layers is a cheap, qualitative way to see roughly where in the network a factual answer or a syntactic decision seems to firm up — before you've committed to a hypothesis worth a more expensive, targeted experiment.

Activation patching in one trace

The logit lens only reads activations. The genuinely causal move — the one behind causal tracing and ROME-style locating-and-editing work — is activation patching: run a "clean" prompt and a "corrupted" prompt, and splice an activation captured from the clean run into the corrupted run to see how much of the clean behavior comes back. NNsight's invoker lets a single trace cover both runs, so the patch doesn't need to leave Python at all: 

clean_prompt = "The Eiffel Tower is located in the city of"
corrupt_prompt = "The Colosseum is located in the city of"

with model.trace() as tracer:
    with tracer.invoke(clean_prompt):
        clean_act = model.transformer.h[6].output[0].save()

    with tracer.invoke(corrupt_prompt):
        model.transformer.h[6].output[0][:] = clean_act
        patched_logits = model.lm_head.output.save()

If patching layer 6's output is enough to make the corrupted run's top prediction shift toward "Paris," that's evidence layer 6 is carrying (or has already committed to) the information the final answer depends on. Sweeping this same patch across every layer and token position, and plotting how much of the clean-minus-corrupt logit gap each patch restores, is exactly how causal-tracing figures in interpretability papers get made — and it's a loop over the same six-line pattern above, not six lines per layer.

Ablation as a cheaper causal test

Patching needs a matched pair of prompts. Sometimes the simpler question is just "does this component matter at all," which a direct edit answers more cheaply: 

with model.trace(clean_prompt) as tracer:
    # zero out one attention head's contribution at a chosen layer
    heads = model.transformer.h[6].attn.output[0]
    heads[:, :, head_index, :] = 0
    ablated_logits = model.lm_head.output.save()

Comparing ablated_logits against the unmodified run's logits for the token you care about is a minimal necessity test: if zeroing a head barely moves the prediction, that head almost certainly isn't doing the work your hypothesis assigns it, at least not on this input.

What this doesn't give you for free

NNsight removes the mechanical friction of writing an intervention, and that's a real thing to remove — a working ablation or patch that used to take an afternoon of hook-registration bookkeeping is now a few lines. It doesn't remove the actual interpretability work. You still need a hypothesis about which layer, head, or position plausibly matters before you write the trace; the tool tests a specific causal claim, it doesn't generate one. A single successful patch or a suggestive logit-lens trajectory is evidence, not proof — the standard worry in this literature is that an intervention can restore a metric for a reason unrelated to the mechanism you think you're isolating, which is why serious causal-tracing work follows up a single patch with sweeps across layers and positions, path patching to isolate which downstream computation actually uses the patched value, and ablations that check necessity as well as sufficiency. And remote execution through NDIF is only as useful as the specific models it hosts at a given time — it widens what's practical to try, it doesn't make every model equally available.

This is an introductory walkthrough of the NNsight API, not a write-up of a specific research result. Code in this post follows the library's documented public API at nnsight.net; I haven't independently benchmarked the patches shown here against a published causal-tracing result, and the logit-lens and patching examples are meant to illustrate the pattern rather than report findings. An LLM was used only for organizing this write-up.

References