Introducing NNsight 0.6
By Jaden Fiotto-Kaufman
NNsight is releasing its sixth major version, focused on addressing user feedback about common hurdles with the library.
Wait, What is NNsight Again?
If you are a new user or you haven't used NNsight in a while, here's a small refresher! NNsight is a Python library for interpreting and intervening on the internals of PyTorch models. You wrap a model, open a tracing context, and read or write activations at any layer. While the tool supports any Pytorch model, we also provide first-class support for popular architectures such as
transformers and diffusers.
from nnsight import LanguageModel
model = LanguageModel("openai-community/gpt2", device_map="auto", dispatch=True)
with model.trace("The Eiffel Tower is in"):
# Read the hidden states at layer 5
hidden = model.transformer.h[5].output[0].save()
# Zero out the MLP output at layer 0
model.transformer.h[0].mlp.output[:] = 0
Under the hood, NNsight uses deferred execution. When you enter with model.trace(...), your code AST is extracted, compiled into a function, and run in a worker thread. When that thread accesses .output, it waits until the model's forward pass reaches the selected module, extracting the desired output tensor through a PyTorch hook. This means your intervention code is fully aligned with the forward pass—no proxies, no fake tensors. You're working with real PyTorch values!