Open-heart surgery for neural networks, Pythonically

What it does nnsight is a Python library that provides a tracing context for PyTorch models, letting you read and rewrite intermediate activations during forward or generation passes using familiar attribute paths. It supports local execution on standard models and remote runs on large models via the NDIF infrastructure. The library also batches interventions across multiple inputs, computes gradients on intermediate tensors, and supports non-destructive model editing.

The interesting bit Instead of forcing you to register hooks by hand, nnsight captures your code via AST and defers execution, synchronizing worker threads so you can treat layer outputs like local variables inside a context block. It can even rewrite a module’s forward method on the fly to expose internal operations for source-level tracing.

Key highlights

• Read and overwrite hidden states at any layer during a forward pass using dot-notation paths.
• Batch interventions across multiple prompts with serial worker threads inside invoke() blocks.
• Iterate over autoregressive generation steps and apply conditional, per-step activation patches.
• Inspect tensor shapes without running the full model through a scan() mode.
• Execute traces remotely on large models via NDIF, or apply the same interface to arbitrary PyTorch architectures.

Caveats

• Values inside a trace are garbage collected unless explicitly marked with .save(), a sharp edge for newcomers.
• The unbounded generation iterator (tracer.iter[:]) hangs forever; any post-generation logic must move to a separate invoke() block.
• Layer access must follow model execution order, so you cannot backtrack to earlier activations once the trace has moved past them.

Verdict Researchers doing mechanistic interpretability or causal interventions on PyTorch models will find this far more ergonomic than raw forward hooks. If you only need black-box inference, the abstraction is overkill.