Model or dataset
ndif-team/nnsight avatar
ndif-team/nnsight

nnsight: In-Process Intervention on PyTorch Forward Passes

The nnsight package enables interpreting and manipulating the internals of deep learned models.

1,101 stars118 forksPythonMIT

At a glance

What is it?
nnsight lets you read and edit a model's internal tensors from inside a with model.trace(...) block, running locally or forwarded to NDIF for models you cannot host. The design is elegant for research scripts and awkward for anything that needs to run outside a Python process.
Who is it for?
Adopt nnsight if your work is exploratory interpretability on PyTorch models and you want interventions written in execution order rather than as registered hooks. Do not adopt it if you need interventions to survive outside a Python process, or if you cannot accept that the trace block describes rather than executes.
Can I use it commercially?
Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
Is it still maintained?
Yes. The repository last received commits 2 days ago.
What is it written in?
Mainly Python, according to GitHub's language statistics.

Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.

DEEP OPEN-SOURCE ANALYSIS

The problem nnsight targets: intervention without hook registration

Standard PyTorch interpretability work means registering forward hooks, capturing tensors into module-level state, and then reasoning about a call graph that no longer matches the order you wrote. nnsight inverts that. The README states that you open a with model.trace(...) block and write ordinary Python against any internal value, reading it, editing it, or saving it, without registering hooks or refactoring the model. The intervention is written in the order it happens and nnsight runs it interleaved with the real forward pass. That ordering property is the whole product. When you assign model.transformer.h[0].output[:] = 0 inside the block, the model computes downstream on the zeroed value, and the code reads top to bottom the way the computation does. The audience is interpretability researchers and anyone doing activation patching, ablation, or probing who is tired of hook bookkeeping. It is not aimed at production inference serving.

What actually happens inside a trace block

The README is explicit that inside the block you are not running the model, you are describing what to do when it runs. Reading .output gives you the real tensor once the model reaches that module; assigning to it splices your value in. Anything you want after the block must be marked with .save() or nnsight.save(x). That two-phase model has a direct consequence: a variable captured inside the block without .save() is not available outside it, and a print inside the block will not show you a tensor. The repository points to NNsight.md for the underlying machinery, described there as tracing, interleaving, and the envoy tree. One practical wrinkle the README calls out: a GPT-2 block's .output is a plain tensor, but some modules such as attention return a tuple, so you index .output[0]. It suggests print(model) or print(module.source) to inspect the shape. That is the kind of detail that costs an afternoon if you skip it.

Generation, batching, gradients, and out-of-order module calls

The README documents several operations beyond a single forward pass. generate returns token ids on tracer.result, so you save tracer.result and decode afterwards. To reach into each generation step you save a container and append to it inside for step in tracer.iter[:5], with the README warning to use a bounded range so code after the loop still runs. Several prompts can share one pass through tracer.invoke, and the README states each invoke block sees only its own rows. Gradients come from wrapping a backward call, as in with model.output.logits.sum().backward(), then reading hidden.grad.save(). Out-of-order module application is supported, and the README's logit lens example decodes a middle layer through the head by calling model.lm_head(model.transformer.ln_f(hidden)). The README also lists source tracing via .source, persistent edit(), skip(), scan() for shapes, cache(), session(), and the vLLM and diffusion runtimes, with details in docs/ rather than in the README itself.

Getting it running: install, wrapper choice, and the NDIF key

Installation is pip install nnsight. The quick start imports TransformersModel, constructs it with a HuggingFace identifier such as openai-community/gpt2 and dispatch=True, then opens the trace. Wrapping your own network uses NNsight from the same package around any torch.nn.Module, and the README's example wraps a two-layer Sequential and indexes model[0] to zero the first layer. Remote execution is the same trace with remote=True passed to model.trace, preceded by from nnsight import CONFIG and CONFIG.set_default_api_key("YOUR_NDIF_KEY"). The README points to nnsight.net/status for which models are available on NDIF. There is also an agent-facing path: a Claude Code plugin via /plugin marketplace add https://github.com/ndif-team/skills.git followed by /plugin install nnsight@skills, an OpenAI Codex skill-installer command against the same repository, or pointing an MCP client at https://mcp.context7.com/mcp. Handing an agent CLAUDE.md and NNsight.md is the third option listed.

Where the trace-block model gets in the way

The describe-then-execute split is the main limitation and it is structural, not a bug. Because the block is a description, ordinary control flow that depends on a tensor value cannot branch on that value at description time. The README's own workaround for the generation loop is to use a bounded range so code after the loop still runs, which tells you the loop bound cannot come from the model. Anything you forget to .save() is gone. Debugging is therefore indirect: you inspect shapes with scan() or print(module.source) rather than reading values as they flow. The second limitation is remote execution. Models on NDIF are reached through infrastructure the README links to at ndif.us, and using it requires an API key and a network round trip; the README does not describe latency, quotas, or failure behaviour for that path, so treat those as unverified. Third, the README's worked examples are GPT-2 and Llama-shaped, and the tuple-versus-tensor difference between modules means a trace that works on one architecture can silently address the wrong object on another.

nnsight against TransformerLens and plain hooks

The closest comparison is TransformerLens, which takes the opposite approach: it reimplements model architectures with named, canonical hook points such as blocks.0.hook_resid_post, so interventions target a stable vocabulary rather than the model's own module tree. nnsight instead wraps the model as it is and addresses whatever the module tree exposes, which is why the README can say any torch.nn.Module works and why the same trace can run against HuggingFace, diffusers, and vLLM models. The trade is consistency for coverage. With TransformerLens you get hook names that mean the same thing across supported architectures and a smaller set of models; with nnsight you get every PyTorch model and the burden of knowing each one's module layout. Plain forward hooks sit at the other extreme: they work everywhere with no dependency, but the intervention order in your source no longer matches execution order, which is precisely the cost nnsight is built to remove.

Maintenance, releases, and the MIT licence

The project is MIT licensed, which permits commercial and closed-source use with the usual requirement to carry the licence notice; that is a summary of the identifier, not legal advice, and you should read the LICENSE file in the repository. Release cadence visible in the material is roughly every few months: v0.6.3 in March 2026, v0.7.0 in May 2026, and v0.8.0.rc1 in September 2026, with the last push to main in September 2026. The rc suffix on the newest tag means the most recent work is a release candidate, so pinning to v0.7.0 is the conservative choice for anything you need to reproduce. The upgrade surface is the trace API itself: .save(), tracer.iter, tracer.invoke, and the .output accessor are the constructs your scripts are built on, so a change to any of them touches every script you have written. The repository also ships NNsight.md as a design-and-implementation manual and a docs/ directory of task recipes, which is where the maintenance burden of keeping up with API changes is documented. The README's citation block asks academic users to cite the NNsight and NDIF paper, which is a convention rather than a licence term.

Editorial conclusion

Adopt nnsight if your work is exploratory interpretability on PyTorch models and you want interventions written in execution order rather than as registered hooks. Do not adopt it if you need interventions to survive outside a Python process, or if you cannot accept that the trace block describes rather than executes. Before committing, verify three things against your own setup: that the module you intend to edit returns the structure you expect (the README notes attention modules return tuples, so you index .output[0]), that your NDIF key is set via CONFIG.set_default_api_key if you plan to use remote=True, and that a bounded range such as tracer.iter[:5] still lets code after the loop run. The last one is the difference between a script that finishes and one that hangs.

Official sources

  1. License: MIT
  2. ndif-team/nnsight on GitHub
  3. Project website
  4. README
  5. Releases
Community notes

Community notes