nnsight – Peek Inside and Edit Deep‑Learning Models

What it is
nnsight is a Python library that lets you trace the forward pass of any PyTorch model (e.g., Hugging Face Transformers, Diffusers, vLLM) and read or modify intermediate tensors as the model runs.
It works locally or, with the NDIF service, on very large models you can’t host yourself.

How it works
You wrap a model with nnsight (e.g., TransformersModel) and open a with model.trace(...): block. Inside the block you write ordinary Python that accesses .output of layers, reads values with .save(), or overwrites them. The library injects these operations into the real forward pass, so the model computes using your edited tensors without you having to add hooks manually.

Key features

• Read/write internals – inspect hidden states, attention maps, gradients; replace them on‑the‑fly.
• Generation hooks – intervene at each generation step, collect logits, or run a logit‑lens.
• Batch‑wise tracing – handle multiple prompts in one pass with tracer.invoke.
• Gradient capture – compute gradients w.r.t. any internal tensor.
• Remote execution – set remote=True and run the same trace on NDIF‑hosted models.
• Support for many runtimes – wrappers for Hugging Face, Diffusers, vLLM, plus a generic NNsight wrapper for any torch.nn.Module.

Typical usage

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

with model.trace("The Eiffel Tower is in the city of"):
    # Zero the first transformer block’s output
    model.transformer.h[0].output[:] = 0
    # Save a hidden state deeper in the network
    hidden = model.transformer.h[6].output.save()
    # Keep the final logits for later
    logits = model.output.logits.save()

print(hidden.shape)          # torch.Size([1, 10, 768])
print(logits.argmax(-1))     # token predictions

The same pattern works for generation, gradient extraction, or remote runs.

Who might use it

• Researchers probing model internals, doing mechanistic interpretability, or building “logit‑lens” visualisations.
• Developers building debugging or editing tools for LLMs.
• Agents that need to run custom interventions on foundation models (the README even lists ways to expose nnsight to Claude, OpenAI Codex, etc.).

Where to learn more

• Full docs: https://www.nnsight.net
• Design paper (arXiv 2407.14561) and the NNsight.md implementation guide in the repo.
• Interactive Colab walkthrough linked in the README.

Installation

pip install nnsight

---

All details are taken directly from the repository’s README.