ndif-team/nnsight

nnsight – Interpreting & Editing the Inside of PyTorch Models

What it is – nnsight is a Python library (installable via pip install nnsight) that lets you peek inside any PyTorch model, read the hidden‑state tensors at any layer, change them on‑the‑fly, compute gradients of intermediate values, and even create permanent edits to a model. It works locally on your own GPU/CPU and can also run on ND Institute’s remote infrastructure for very large models.

Why it matters – Modern foundation models (GPT‑2, LLaMA, etc.) are black boxes. Researchers who want to understand why a model makes a particular prediction, test causal hypotheses, or prototype model‑editing techniques need a clean way to access and manipulate internal activations without rewriting the model’s code. nnsight provides a high‑level, Pythonic API that abstracts away the hook‑injection and tracing boiler‑plate.

---

Core capabilities (as described in the README)

Feature	How you use it	What you get
Activation access	with model.trace(prompt): hidden = model.transformer.h[5].output[0].save()	A real tensor containing the hidden state of layer 5 for the given prompt.
In‑place intervention	model.transformer.h[0].output[0][:] = 0 inside a trace block	The forward pass continues with the modified activation, letting you test causal effects.
Gradient of intermediate tensors	with loss.backward(): grad = hs.grad.save()	Gradient of any tensor (e.g., a hidden state) with respect to a loss you define.
Batch‑level invocations	with tracer.invoke(prompt): …	Run many prompts in parallel threads; each invoke runs sequentially but shares values via .save().
Generation with step‑wise control	with model.generate(prompt, max_new_tokens=5) as tracer: for step in tracer.iter[:]: …	Autoregressive generation where you can intervene on specific steps.
Model editing	with model.edit() as edited: edited.transformer.h[0].output[0][:] = 0	Produces a new LanguageModel instance that permanently incorporates the edit, leaving the original untouched.
Shape‑only scanning	with model.scan(prompt): dim = nnsight.save(layer.output.shape[-1])	Retrieves tensor shapes without a full forward pass.
Caching & sessions	cache = tracer.cache() or with model.session() as s: …	Re‑uses previously captured activations across multiple traces for efficiency.
Remote execution	CONFIG.set_default_api_key(..); model = LanguageModel('meta-llama/Meta-Llama-3.1-8B'); with model.trace(..., remote=True): …	Runs the tracing code on NDIF’s cloud service, useful for models that don’t fit locally.
vLLM integration	from nnsight.modeling.vllm import VLLM; model = VLLM('gpt2', ...)	High‑throughput inference while still exposing the same tracing API.
Arbitrary PyTorch models	NNsight(net) where net is any torch.nn.Module	The same tracing/intervention tools work on non‑transformer models (e.g., simple feed‑forward nets).

---

Quick‑start example (from the README)

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='auto', dispatch=True)

with model.trace('The Eiffel Tower is in the city of'):
    # Zero out the first layer’s activations
    model.transformer.h[0].output[0][:] = 0
    # Save the final hidden state and the logits
    hidden = model.transformer.h[-1].output[0].save()
    logits = model.output.save()

print(model.tokenizer.decode(logits.logits.argmax(dim=-1)[0]))

The snippet demonstrates loading a model, opening a trace, intervening on a layer, and retrieving the final prediction.

---

Typical use cases

• Mechanistic interpretability – Examine how attention heads or MLP blocks contribute to a token’s prediction.
• Causal probing – Zero, add noise, or replace activations to test hypotheses about information flow.
• Model editing research – Create reversible edits (e.g., “make the model always output ‘Paris’ after ‘Eiffel Tower’”).
• Debugging custom architectures – Use the same API on any torch.nn.Module to verify forward‑pass behavior.
• Efficient batch experiments – Run many prompts in a single forward pass using the invoke/session machinery.

---

Limitations & gotchas (as noted in the README)

• Execution order matters – Inside a trace you must access modules in the exact order they are executed; otherwise you hit an OutOfOrderError deadlock.
• Thread‑based synchronization – The library runs your tracing code in a separate worker thread; values are only available after you call .save().
• Unbounded iterators – tracer.iter[:] never returns, so any code after it will not run unless placed in a separate invoke block.
• Remote execution requires an API key and the model must be available on NDIF’s platform.
• vLLM integration is limited to models supported by vLLM (currently transformer‑style language models).

---

Where to learn more

• Documentation site – <https://www.nnsight.net
• Paper – NNsight and NDIF: Democratizing Access to Foundation Model Internals (arXiv 2407.14561)
• Discord & forum – Community support and discussion channels linked in the README.
• Colab walkthrough – Interactive notebook for hands‑on exploration.

---

TL;DR

nnsight gives researchers a concise, Pythonic way to trace, read, modify, and persistently edit the hidden states of any PyTorch model, with support for batching, gradient extraction, remote execution, and high‑performance back‑ends like vLLM. It is a genuine, research‑oriented tool for mechanistic work on modern foundation models.