Info
Last Execution: 2026-07-27
| Package | Version |
|---|---|
| nnsight | 0.8 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0 |
NNsight Walkthrough¶
The API for a transparent science on black-box AI¶
In this era of large-scale deep learning, the most interesting AI models are massive black boxes that are hard to run. Ordinary commercial inference service APIs let us interact with huge models, but they do not let us access model internals.
The nnsight library is different: it provides full access to all neural
network internals. When using nnsight together with a remote service like the
National Deep Inference Fabric
(NDIF), it is possible to run complex experiments on huge open models easily
with fully transparent access.
Through NDIF and NNsight, our team wants to enable entire labs and independent researchers alike, as we believe a large, passionate, and collaborative community will produce the next big insights on this profoundly important field.
This walkthrough teaches you nnsight from the ground up, starting with the core mental model and building to advanced features. The local cells below were executed on GPT-2 to produce the outputs you see; the remote example is illustrative and runs on NDIF.
Table of Contents¶
- Getting Started - Setup and wrapping models
- Intervening - Accessing and modifying activations
- Language Models -
TransformersModel, batching, and multi-token generation - Gradients - Accessing and modifying gradients
- Advanced Features - Source tracing, caching, early stopping, skipping, scanning
- Model Editing - Persistent modifications
- Remote Execution - Running on NDIF
Installation¶
nnsight is on PyPI:
pip install nnsight
See the installation guide for optional extras.
# Install nnsight (uncomment on a fresh environment, e.g. Colab)
!pip install nnsight
from IPython.display import clear_output
clear_output()
A Tiny Model¶
To demonstrate the core functionality and syntax of nnsight, we'll define and use a tiny two-layer neural network.
Our little model here is composed of two submodules – linear layers layer1 and layer2. We specify the sizes of each of these modules and create some complementary example input.
from collections import OrderedDict
import nnsight
import torch
input_size = 5
hidden_dims = 10
output_size = 2
net = torch.nn.Sequential(
OrderedDict([
("layer1", torch.nn.Linear(input_size, hidden_dims)),
("layer2", torch.nn.Linear(hidden_dims, output_size)),
])
).requires_grad_(False)
# random input
input = torch.rand((1, input_size))
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Wrapping with NNsight¶
The core object of the nnsight package is NNsight. This wraps around any PyTorch model to enable investigation of its internal parameters.
from nnsight import NNsight
model = NNsight(net)
Printing a PyTorch model shows a named hierarchy of modules, which is very useful for knowing how to access sub-components directly. NNsight reflects the same hierarchy:
print(model)
Sequential( (layer1): Linear(in_features=5, out_features=10, bias=True) (layer2): Linear(in_features=10, out_features=2, bias=True) )
Python Contexts¶
Before we actually get to using the model, let's talk about Python contexts.
Python contexts define a scope using the with statement and are often used to create some object, or initiate some logic, that you later want to destroy or conclude.
The most common application is opening files:
with open('myfile.txt', 'r') as file:
text = file.read()
Python uses the with keyword to enter a context-like object. This object defines logic to be run at the start of the with block, as well as logic to be run when exiting. When using with for a file, entering the context opens the file and exiting the context closes it. Being within the context means we can read from the file.
Simple enough! Now we can discuss how nnsight uses contexts to enable intuitive access into the internals of a neural network.
The Tracing Context¶
The main tool in nnsight is a context for tracing. We enter the tracing context by calling model.trace(<input>) on an NNsight model, which defines how we want to run the model. Inside the context, we can customize how the neural network runs. The model actually runs upon exiting the tracing context:
input = torch.rand((1, input_size))
with model.trace(input):
# Your intervention code goes here.
# The model runs when the context exits.
pass
But where's the output? To get it, we'll have to learn how to request it from within the tracing context.
The .input and .output Properties¶
When we wrapped our neural network with the NNsight class, this added a couple of properties to each module in the model (including the root model itself). The two most important ones are .input and .output:
model.input # The input to the model
model.output # The output from the model
They correspond to the inputs and outputs of their respective modules during a forward pass. We can use these attributes inside the with block to access values at any point in the network.
Let's try accessing the model's output:
with model.trace(input):
output = model.output
print(output)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
...
NameError: name 'output' is not defined
Oh no, an error!
Why doesn't our output have a value? A traced with block doesn't run in place — its body runs alongside the forward pass and only returns the values you explicitly mark. Values accessed inside a trace exist during the trace; they persist afterward only if you call .save() on them. This keeps memory costs down — we only keep what we explicitly ask for.
Saving Values with .save()¶
Adding .save() fixes the error:
with model.trace(input):
output = model.output.save()
print(output)
tensor([[ 0.6376, -0.0140]])
Success! We now have the model output. We just completed our first intervention using nnsight.
The .save() method tells nnsight "I want to use this value after the trace ends." The value comes back bound to the variable you assign it to, so always write output = model.output.save() — a bare model.output.save() on its own line is marked but has no variable to return under, and is silently lost.
💡 Tip:
nnsight.save(value)is an equivalent form that works on any value, even one without a.save()method (a plain list, for example):output = nnsight.save(model.output)In 0.8, saving only makes sense inside a trace — calling
.save()outside awith model.trace(...)block raises an error rather than silently doing nothing.
Accessing Submodule Outputs¶
Just like we saved the model's output, we can access any submodule's output. Remember when we printed the model earlier? That showed us layer1 and layer2 - we can access those directly:
with model.trace(input):
layer1_output = model.layer1.output.save()
layer2_output = model.layer2.output.save()
print("Layer 1 output:", layer1_output)
print("Layer 2 output:", layer2_output)
Layer 1 output: tensor([[-0.0971, 0.1888, 0.0936, 0.4946, -0.2718, -0.5696, -0.8946, -0.2102,
0.1412, -0.3034]])
Layer 2 output: tensor([[ 0.6376, -0.0140]])
Accessing Module Inputs¶
We can also access the inputs to any module using .input:
| Property | Returns |
|---|---|
.output |
The module's return value |
.input |
The first positional argument to the module |
.inputs |
All inputs as (args_tuple, kwargs_dict) |
with model.trace(input):
layer2_input = model.layer2.input.save()
print("Layer 2 input:", layer2_input)
print("(Notice it equals layer1's output!)")
Layer 2 input: tensor([[-0.0971, 0.1888, 0.0936, 0.4946, -0.2718, -0.5696, -0.8946, -0.2102,
0.1412, -0.3034]])
(Notice it equals layer1's output!)
Operations on Values¶
The values you access are real tensors, so you can apply any PyTorch operation to them inside the trace:
with model.trace(input):
layer1_out = model.layer1.output
# These are real tensor operations, computed during the forward pass.
max_idx = torch.argmax(layer1_out, dim=1).save()
total = (model.layer1.output.sum() + model.layer2.output.sum()).save()
print("Max index:", max_idx)
print("Total:", total)
Max index: tensor([3]) Total: tensor(-0.8050)
The Core Paradigm: Interleaving¶
When you write intervention code inside a with model.trace(...) block, here's what actually happens:
- Your code is captured - nnsight reads the source of the code inside the
withblock. - The code is compiled into an executable function.
- Your code runs alongside the model - as the model executes its forward pass, your intervention code runs next to it.
- Your code waits for values - when you access
.output, your code pauses until the model reaches that point. - The model provides values via hooks - PyTorch hooks hand values to your waiting code.
- Your code can modify values - before the forward pass continues, you can change activations.
This process is called interleaving - your intervention code and the model's forward pass take turns executing, synchronized at specific points (module inputs and outputs).
┌─────────────────────────────────────────────────────────────────────┐
│ Forward Pass (main) Intervention Code (your code) │
│ ───────────────────── ───────────────────────────── │
│ │
│ model(input) # Your code starts │
│ │ │ │
│ ▼ ▼ │
│ layer1.forward() hs = model.layer1.output │
│ │ │ │
│ │──── hook provides value ──────────►│ │
│ │ │ │
│ │◄─── your code continues ────────── │ │
│ │ (can modify value) │ │
│ ▼ ▼ │
│ layer2.forward() out = model.layer2.output │
│ │ │ │
│ ▼ ▼ │
│ return output # Your code finishes │
└─────────────────────────────────────────────────────────────────────┘
Key insight:
Because your code waits for values as the forward pass progresses, you must access modules in the order they execute.
✅ Correct: Access layer 0, then layer 5
with model.trace("Hello"):
layer0_out = model.layers[0].output.save() # Waits for layer 0
layer5_out = model.layers[5].output.save() # Then waits for layer 5
❌ Wrong: Access layer 5, then layer 0
with model.trace("Hello"):
layer5_out = model.layers[5].output.save() # Waits for layer 5
layer0_out = model.layers[0].output.save() # ERROR! Layer 0 already executed
When you try to access a module that has already executed, nnsight raises an OutOfOrderError: the forward pass has already moved past that point, so you missed your chance to intercept that value. The same rule applies to submodules — a block's attention and MLP run before the block's own output, so access them in that order.
Modification¶
Not only can we view intermediate states of the model, we can modify them and see the effect on the output.
Use slice-assignment with [:] for in-place modifications:
with model.trace(input):
# Save the original (clone first, since we'll modify in place).
before = model.layer1.output.clone().save()
# Zero out the first feature.
model.layer1.output[:, 0] = 0
# Save the modified value.
after = model.layer1.output.save()
print("Before:", before)
print("After: ", after)
Before: tensor([[-0.0971, 0.1888, 0.0936, 0.4946, -0.2718, -0.5696, -0.8946, -0.2102,
0.1412, -0.3034]])
After: tensor([[ 0.0000, 0.1888, 0.0936, 0.4946, -0.2718, -0.5696, -0.8946, -0.2102,
0.1412, -0.3034]])
Replacement¶
You can also replace an output entirely by assigning a new value to it:
with model.trace(input):
original = model.layer1.output.clone()
# Add noise to the activation.
noise = 0.1 * torch.randn_like(original)
model.layer1.output = original + noise
modified = model.layer1.output.save()
print("Modified output:", modified)
Modified output: tensor([[-0.0746, 0.1170, 0.0206, 0.4548, -0.5263, -0.5447, -0.9830, -0.1980,
0.3068, -0.2758]])
Error Handling and Debugging¶
If you make a mistake (like invalid indexing), nnsight surfaces the error at the line inside your block that caused it:
# This fails because hidden_dims=10, so valid indices are 0-9.
try:
with model.trace(input):
model.layer1.output[:, hidden_dims] = 0 # Index 10 is out of bounds!
except IndexError as e:
print("Caught error:", e)
Caught error: index 10 is out of bounds for dimension 1 with size 10
Debugging tips:
- Use
print()inside traces - it works normally and prints values as they're computed. - Use
breakpoint()to drop intopdband inspect values interactively. - Toggle internal frames with
nnsight.CONFIG.APP.DEBUG = Trueto see NNsight's internal execution in tracebacks (helpful when the default traceback isn't clear enough).
with model.trace(input):
out = model.layer1.output
print("Layer 1 shape:", out.shape) # Works!
breakpoint() # Drops into pdb - inspect `out`, etc.
3. Language Models¶
Now that we have the basics under our belt, let's scale up to a real language model and combine the techniques we've learned into more interesting experiments.
The bare NNsight class we used in Part 2 wraps a model as-is and does no pre-processing on the inputs. For HuggingFace language models, nnsight provides TransformersModel, a subclass that adds a lot of convenience:
- Automatic tokenization - pass strings directly, no manual tokenization needed.
- HuggingFace integration - load any model from the HuggingFace Hub by its id.
- Generation support - built-in multi-token generation with
.generate(). - Batching - efficiently process multiple inputs in one forward pass.
More on loading models
TransformersModel is the primary wrapper for HuggingFace models in nnsight 0.8. Pass a repo id; dispatch=True loads the weights right away instead of on the first trace, and device_map="auto" lets HuggingFace Accelerate place them on the best available device (a GPU if present, otherwise CPU) and shard a model too large for one GPU across several.
(LanguageModel is a deprecated alias that still works but warns on construction.) For the full set of constructor options — meta-model dispatch, passing an already-loaded model, per-device placement — see the Loading a Model tutorial and the transformers model guide.
Let's load GPT-2 and start experimenting!
from nnsight import TransformersModel
llm = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
print(llm)
GPT2LMHeadModel(
(transformer): GPT2Model(
(wte): Embedding(50257, 768)
(wpe): Embedding(1024, 768)
(drop): Dropout(p=0.1, inplace=False)
(h): ModuleList(
(0-11): 12 x GPT2Block(
(ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
(attn): GPT2Attention(
(c_attn): Conv1D()
(c_proj): Conv1D()
(attn_dropout): Dropout(p=0.1, inplace=False)
(resid_dropout): Dropout(p=0.1, inplace=False)
)
(ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
(mlp): GPT2MLP(
(c_fc): Conv1D()
(c_proj): Conv1D()
(act): NewGELUActivation()
(dropout): Dropout(p=0.1, inplace=False)
)
)
)
(ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True, bias=True)
)
(lm_head): Linear(in_features=768, out_features=50257, bias=False)
(generator): Generator(
(streamer): Streamer()
)
)
Notice the model structure. GPT-2 has:
transformer.wte- token embeddingstransformer.h- a list of transformer blocks (layers 0-11)lm_head- the output projection to the vocabulary
With TransformersModel you can pass strings directly - tokenization happens automatically:
with llm.trace("The Eiffel Tower is in the city of"):
# Hidden states from the last transformer block.
hidden_states = llm.transformer.h[-1].output.save()
# The model's final output (a HuggingFace output object with a `.logits` field).
output = llm.output.save()
print("Hidden states shape:", hidden_states.shape)
print("Predicted next token:", repr(llm.tokenizer.decode(output.logits[0, -1].argmax())))
Hidden states shape: torch.Size([1, 10, 768]) Predicted next token: ' Paris'
Everything you learned with the tiny model applies here! The same .input, .output, and .save() patterns work — the key difference is that you can pass strings directly.
💡 What does
.outputlook like? It depends on the model and thetransformersversion — a module may return a plain tensor or a tuple, so verify withprint(module),type(...), or.shapebefore indexing. In this GPT-2 build, a transformer block returns a plain(batch, seq, hidden)tensor (so we index it directly, with no[0]), while the attention submodule returns a tuple whose hidden states live at.output[0].
Batching with Invokers¶
So far we've run one input at a time. To process several prompts efficiently, or apply different interventions to each, use invokers. Call .trace() with no input, then open a tracer.invoke(...) block per prompt — every invoke contributes its rows to a single batched forward pass:
with llm.trace() as tracer:
with tracer.invoke("The Eiffel Tower is in the city of"):
paris_logits = llm.lm_head.output[:, -1].save()
with tracer.invoke("The Colosseum is in the city of"):
rome_logits = llm.lm_head.output[:, -1].save()
# Both ran in ONE forward pass.
print("Prompt 1 prediction:", repr(llm.tokenizer.decode(paris_logits.argmax(dim=-1))))
print("Prompt 2 prediction:", repr(llm.tokenizer.decode(rome_logits.argmax(dim=-1))))
Prompt 1 prediction: ' Paris' Prompt 2 prediction: ' P'
How Invokers Execute¶
Invokers are not run strictly one-after-another. Each invoke's body runs in its own cooperative worker; they all start together and resume in the order the model reaches what each one asked for. That's what makes them a batch rather than a sequence.
A consequence: if a later invoke wants to use a value an earlier invoke produced from an activation, the consumer would run before the producer has bound it. tracer.barrier(n) fixes the ordering — everything written above a barrier() call happens before anything written below one. Here we capture the last-layer hidden state from a donor prompt and patch it into another prompt, changing its prediction:
with llm.trace() as tracer:
barrier = tracer.barrier(2) # two participating invokes
# First invoke: read the donor prompt's last-layer hidden state.
with tracer.invoke("The Colosseum is in the city of"):
donor = llm.transformer.h[-1].output
barrier() # signal: donor has been read
# Second invoke: wait, then patch that hidden state in.
with tracer.invoke("The Eiffel Tower is in the city of"):
barrier() # wait until the donor is materialized
llm.transformer.h[-1].output = donor
patched_logits = llm.lm_head.output[:, -1].save()
# The Eiffel prompt no longer predicts " Paris" — it now gives the donor prompt's answer.
print("Patched prediction:", repr(llm.tokenizer.decode(patched_logits.argmax(dim=-1))))
Patched prediction: ' P'
💡 Note: A value defined in the enclosing scope (before the invokes) flows into every invoke with no barrier — only values produced from an activation inside one invoke need synchronizing. Batching, cross-prompt transfer, and activation patching get a full treatment in the Batching tutorial.
Multi-Token Generation¶
So far we've done single forward passes. Language models generate text by running multiple forward passes - one per token - which means the same modules are called multiple times.
Use .generate() instead of .trace() for multi-token generation. It returns the generated token ids, available on tracer.result:
with llm.generate("The Eiffel Tower is in", max_new_tokens=3, do_sample=False) as tracer:
ids = tracer.result.save()
print(repr(llm.tokenizer.decode(ids[0])))
'The Eiffel Tower is in the middle of'
Iterating Over Generation Steps with .iter¶
During generation, modules are called once per token. To intervene or collect data at each step, iterate the steps with tracer.iter. This is essential whenever modules run more than once — generation, diffusion steps, recurrent networks, and so on:
with llm.generate("The Eiffel Tower is in", max_new_tokens=3, do_sample=False) as tracer:
tokens = nnsight.save([])
# Iterate over the (bounded) generation steps.
for step in tracer.iter[:3]:
tokens.append(llm.lm_head.output[0, -1].argmax(dim=-1))
print("Generated tokens:", [llm.tokenizer.decode(t) for t in tokens])
Generated tokens: [' the', ' middle', ' of']
tracer.iter accepts different patterns:
| Pattern | Meaning |
|---|---|
tracer.iter[:N] |
The first N steps |
tracer.iter[0] |
First step only |
tracer.iter[1:3] |
Steps 1 and 2 |
tracer.iter[[0, 2, 4]] |
An explicit list of steps |
The step variable is a plain integer, so an ordinary if lets you apply different logic at different steps:
with llm.generate("Hello", max_new_tokens=5, do_sample=False) as tracer:
for step in tracer.iter[:5]:
# Only intervene on the first (prefill) step.
if step == 0:
llm.transformer.h[0].output[:] = 0 # Zero out layer 0's contribution
ids = tracer.result.save()
print(repr(llm.tokenizer.decode(ids[0])))
"Hello, I'm not sure"
💡 Key takeaway:
.iterworks anywhere modules are called multiple times, not just LLM generation — diffusion denoising steps, RNN time steps, any iterative computation.
Footgun: unbounded iteration
tracer.iter[:] (equivalently tracer.all()) is unbounded — it hands out step indices until the model itself stops generating, so it never knows there is a "last" iteration. Any code written after an unbounded loop never runs:
with model.generate("Hello", max_new_tokens=3) as tracer:
for step in tracer.iter[:]: # unbounded
hidden = model.transformer.h[-1].output
ids = tracer.result.save() # ⚠️ never executes!
Prefer a bounded range (tracer.iter[:N]) when you have code to run afterward, or put that code before the loop. See the Multiple Token Generation tutorial.
Section 3 Summary¶
You've learned the core patterns for working with language models in nnsight:
- TransformersModel - load HuggingFace models with automatic tokenization.
- Invokers - process multiple prompts in one batched forward pass.
- Barriers - move a value from one invoke into another.
- Multi-token generation -
.generate()withtracer.result. - Iteration with
.iter- intervene at each step when modules run multiple times.
These patterns are the foundation for interpretability research.
4. Gradients¶
nnsight supports gradient access and modification through a backward tracing context. This is essential for gradient-based interpretability methods like attribution, saliency maps, and gradient-based steering.
Just as we use with model.trace() to intercept the forward pass, we use with loss.backward(): to intercept the backward pass. The key insight: during backpropagation, gradients flow in reverse order — from the loss back through the model — so you access .grad in the reverse order of how you accessed the tensors during the forward pass.
with llm.trace("The Eiffel Tower is in the city of"):
# FORWARD PASS: grab the tensor we want gradients for, and enable grad on it.
hs = llm.transformer.h[-1].output
hs.requires_grad_(True)
# Compute a scalar loss from a value that comes later in the forward pass.
logits = llm.lm_head.output
loss = logits.sum()
# BACKWARD PASS: gradients flow loss -> logits -> hidden states.
with loss.backward():
grad = hs.grad.save()
print("Gradient shape:", tuple(grad.shape))
print("Gradient norm: ", round(grad.norm().item(), 3))
Gradient shape: (1, 10, 768) Gradient norm: 127693.43
Understanding Gradient Order¶
This is the same interleaving principle as the forward pass, but reversed:
Forward pass order: layer0 → layer1 → ... → layer11 → lm_head → loss
Backward pass order: loss → lm_head → layer11 → ... → layer1 → layer0
If you accessed layer5.output and layer10.output during the forward pass, access their gradients in reverse: layer10.grad first, then layer5.grad.
Rules for gradients:
.gradis only accessible inside awith tensor.backward():context..gradis a property of tensors, not modules.- Get the tensor via
.outputbefore entering the backward context. - Call
.requires_grad_(True)on a tensor the graph should flow from. - Access gradients in reverse order of how you got the tensors.
Modifying Gradients¶
You can modify gradients just like activations — useful for gradient clipping, masking, or steering:
with llm.trace("The Eiffel Tower is in the city of"):
hs = llm.transformer.h[-1].output
hs.requires_grad_(True)
loss = llm.lm_head.output.sum()
with loss.backward():
original_grad = hs.grad.clone().save()
# Modify the gradient in place (here, zero it out).
hs.grad[:] = 0
modified_grad = hs.grad.save()
print("Original grad mean:", round(original_grad.mean().item(), 6))
print("Modified grad mean:", round(modified_grad.mean().item(), 6))
Original grad mean: -1e-06 Modified grad mean: 0.0
For more, see the Gradients tutorial.
5. Advanced Features¶
Let's explore some powerful features that unlock deeper investigations. Each has its own dedicated tutorial linked below.
5.1 Source Tracing¶
.output and .input hook a module's boundaries. To reach operations inside a module's forward — the individual function and tensor calls — use .source, which rewrites the forward method so each call site becomes hookable. Print it to see the labeled operations (this works even outside a trace):
print(llm.transformer.h[0].mlp.source)
* def forward(self, hidden_states: tuple[torch.FloatTensor] | None) -> torch.FloatTensor:
self_c_fc_0 -> 0 hidden_states = self.c_fc(hidden_states)
self_act_0 -> 1 hidden_states = self.act(hidden_states)
self_c_proj_0 -> 2 hidden_states = self.c_proj(hidden_states)
self_dropout_0 -> 3 hidden_states = self.dropout(hidden_states)
4 return hidden_states
5
Each labeled operation exposes the same .input / .output interface as a module. Here we grab the MLP's hidden states right after the GELU activation (self_act_0), before the down-projection:
with llm.trace("Hello"):
post_gelu = llm.transformer.h[0].mlp.source.self_act_0.output.save()
print("Post-GELU activation shape:", tuple(post_gelu.shape))
Post-GELU activation shape: (1, 1, 3072)
See the Intermediate Operations tutorial for reading, replacing, and recursively drilling into source operations.
5.2 Caching Activations¶
tracer.cache() automatically captures every module's output in one pass — no need to .save() each one by hand:
with llm.trace("Hello") as tracer:
cache = tracer.cache()
# Access cached values after the trace, by module path...
print("Layer 0 output shape:", tuple(cache['model.transformer.h.0'].output.shape))
# ...or with attribute-style access.
print("Same thing: ", tuple(cache.model.transformer.h[0].output.shape))
Layer 0 output shape: (1, 1, 768) Same thing: (1, 1, 768)
See the Cache tutorial for filtering which modules are captured.
5.3 Early Stopping¶
If you only need early layers, call tracer.stop() to abort the run at that point and skip the remaining computation. Save what you need before you stop:
with llm.trace("Hello") as tracer:
layer0 = llm.transformer.h[0].output.save()
tracer.stop() # Don't run the remaining layers.
print("Early stop - only ran the first block")
print("Layer 0 shape:", tuple(layer0.shape))
Early stop - only ran the first block Layer 0 shape: (1, 1, 768)
See the Early Stopping tutorial.
5.4 Skipping a Module¶
.skip(replacement) bypasses a module's forward pass entirely, substituting the value you provide as its output. Here we skip block 1, feeding block 0's output through in its place:
with llm.trace("The Eiffel Tower is in the city of"):
layer0_out = llm.transformer.h[0].output
# Skip block 1 - use block 0's output as its output instead.
llm.transformer.h[1].skip(layer0_out)
output = llm.lm_head.output.save()
print("Skipped block 1! Prediction:", repr(llm.tokenizer.decode(output[0, -1].argmax(dim=-1))))
Skipped block 1! Prediction:
' London'
See the Skip Execution tutorial.
5.5 Scanning (Shape Inference)¶
.scan() runs the model on a meta device to compute shapes without a real forward pass — useful for checking dimensions while you write intervention code:
with llm.scan("Hello"):
hidden_dim = llm.transformer.h[0].output.shape[-1].save()
print("Hidden dimension:", hidden_dim)
Hidden dimension: 768
See the Scan tutorial.
6. Model Editing¶
Interventions inside model.trace() are temporary — they apply to that single forward pass. model.edit() creates a persistently modified version of a model whose interventions replay on every subsequent forward pass.
Entering the edit context binds a (tracer, edited) tuple: write your interventions against edited, then trace through edited to replay them. By default the edit is stored on a shallow copy, so the original stays clean.
# First, get hidden states that predict "Paris".
with llm.trace("The Eiffel Tower is in the city of"):
paris_hidden = llm.transformer.h[-1].output[:, -1, :].save()
# Create an edited model that always injects those hidden states at the last position.
with llm.edit() as (tracer, edited):
edited.transformer.h[-1].output[:, -1, :] = paris_hidden
# Original model: normal prediction.
with llm.trace("Vatican is in the city of"):
original = llm.lm_head.output.argmax(dim=-1).save()
# Edited model: always predicts "Paris".
with edited.trace("Vatican is in the city of"):
modified = edited.lm_head.output.argmax(dim=-1).save()
print("Original:", repr(llm.tokenizer.decode(original[0, -1])))
print("Edited: ", repr(edited.tokenizer.decode(modified[0, -1])))
Original: ' Rome' Edited: ' Paris'
Pass inplace=True to edit the original model directly, and use edited.clear_edits() (or model.clear_edits()) to drop all stored edits. See the Model Editing tutorial for attaching your own modules (SAEs, LoRA adapters) this way.
7. Remote Execution (NDIF)¶
nnsight can run your interventions on large models hosted by NDIF. You write the same trace you would run locally, add remote=True, and nnsight ships the traced block to a server that holds the real weights, runs it there, and streams the .save()d results back — no local GPU required.
Setup¶
Get a free API key at login.ndif.us and set it once with CONFIG.set_default_api_key(...). Check which models are deployed on the status page or with nnsight.status().
from nnsight import CONFIG
CONFIG.set_default_api_key("YOUR_API_KEY")
Remote Tracing¶
For remote use, build the model without dispatch=True — nnsight constructs a lightweight skeleton on the meta device (the architecture exists so module paths are hookable, but no weights are downloaded). Then add remote=True to .trace(). The example below runs Llama-3.1-8B on NDIF; the output shown came back from the service:
import os
os.environ["HF_TOKEN"] = "YOUR_HUGGING_FACE_TOKEN"
# Built on the meta device — no local weights, no local GPU needed.
llm = TransformersModel("meta-llama/Meta-Llama-3.1-8B")
# Just add remote=True — everything else is the same as a local trace!
with llm.trace("The Eiffel Tower is in the city of", remote=True):
hidden_states = llm.model.layers[-1].output.save()
output = llm.output.save()
print("Hidden states shape:", hidden_states[0].shape)
Hidden states shape: torch.Size([11, 4096])
The same code scales up to even larger models (Llama-3.1-70B, 405B, DeepSeek, ...) — just swap in the repo id and keep remote=True.
For the full remote workflow — accounts, model availability, running your own helper code on the server, sessions, and the blocking / non-blocking / async submission modes — see the Access LLMs with NDIF getting-started guide and the Remote Execution tutorial.
Next Steps¶
Congratulations! You've learned the core concepts of nnsight:
- Wrapping models with
NNsightandTransformersModel - Accessing activations with
.output,.input, and.save() - Modifying activations with in-place and replacement patterns
- The interleaving paradigm — your code runs alongside the model
- Batching — invokers and barriers for multi-input experiments
- Multi-token generation —
.generate()and.iter - Gradients —
with loss.backward():for gradient access - Advanced features — source tracing, caching, early stopping, skipping, scanning
- Model editing — persistent modifications with
.edit() - Remote execution — running on NDIF with
remote=True
Where to go next:
- Features — a deep dive into each capability above.
- Documentation — comprehensive API reference.
- Tutorials — end-to-end interpretability techniques (probing, steering, causal mediation, and more).
Getting Involved!¶
Both nnsight and NDIF are in active development. Join us:
- Discord: discord.gg/6uFJmCSwW7
- Forum: discuss.ndif.us
- Twitter/X: @ndif_team
- LinkedIn: National Deep Inference Fabric
We'd love to hear about your work using nnsight! 💟