Info
Last Execution: 2026-09-03
| Package | Version |
|---|---|
| nnsight | 0.8.0 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Model Editing¶
Interventions inside model.trace() are temporary: they apply during that single forward pass and are gone. model.edit() stores the block on the model instead, and every later tracing call replays it before your own interventions run.
"Later tracing call" is the exact boundary. model.trace(), model.generate(), model.pipe() and model.trace(**ids, trace=False) all replay stored edits. Calling the wrapped module directly, as model(**ids) or model._module(**ids), does not: that path never reaches the interleaver, so anything holding the underlying module (a HuggingFace Trainer, an eval harness) sees the model unedited.
Setup¶
import torch
import torch.nn as nn
from nnsight import TransformersModel
model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
Creating an Edited Model¶
Use model.edit() to define persistent interventions. By default (inplace=False) it stores the edit on a shallow copy of the model — the original is left clean. Entering the block binds a (tracer, edited) tuple: write your interventions against edited, and later trace through edited to replay them.
# First, capture the last-layer hidden states that produce "Paris"
with model.trace("The Eiffel Tower is in the city of"):
paris_hs = model.transformer.h[-1].output[:, -1, :].save()
# Create an edited model that always injects the "Paris" hidden states at the last layer
with model.edit() as (tracer, edited):
edited.transformer.h[-1].output[:, -1, :] = paris_hs
# Original model still works normally
with model.trace("Vatican is in the city of"):
original = model.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(f"Original: {model.tokenizer.decode(original[0, -1])}")
print(f"Edited: {edited.tokenizer.decode(modified[0, -1])}")
Original: Rome Edited: Paris
The edit persists across every trace through the edited model:
prompts = [
"The Colosseum is in the city of",
"Big Ben is in the city of",
"The Statue of Liberty is in the city of",
]
for prompt in prompts:
with edited.trace(prompt):
tokens = edited.lm_head.output.argmax(dim=-1).save()
print(f"{prompt} → {edited.tokenizer.decode(tokens[0, -1])}")
The Colosseum is in the city of → Paris Big Ben is in the city of → Paris The Statue of Liberty is in the city of → Paris
Every prompt comes back " Paris", because this edit overwrites the last position's hidden state with the captured one no matter what the input was. That is the edit doing exactly what it says, and it is also the reason an edit meant to change one fact has to be checked on prompts it was not meant to touch. Changing the target is the easy half.
How editing works
model.edit() captures your interventions and stores them on the envoy instead of running them once. On every later trace, the stored edits replay first — before your invoke's own interventions — so their effects are visible to the rest of your trace. The edit context uses the same syntax as model.trace(); you access .output and .input on modules the same way.
Non-inplace edit() stores the edit on a shallow copy (edited). The underlying torch.nn.Module, interleaver, and children are shared with model — only the list of stored edits is independent — so no weights are duplicated.
In-Place Editing¶
By default, model.edit() leaves the original untouched. To modify the original model directly, pass inplace=True. The block then binds only the tracer (there is no separate edited copy), and you write against model itself.
with model.edit(inplace=True) as tracer:
model.transformer.h[-1].output[:, -1, :] = paris_hs
# Now the original model itself is edited — every trace replays the edit
with model.trace("Vatican is in the city of"):
tokens = model.lm_head.output.argmax(dim=-1).save()
print(f"In-place edited: {model.tokenizer.decode(tokens[0, -1])}")
In-place edited: Paris
Use inplace=True with caution
In-place edits affect all subsequent forward passes through the model, including every model.trace() call. If you're experimenting, prefer the default (non-inplace) mode so the original model stays clean.
Clearing Edits¶
Use .clear_edits() to drop all stored edits and restore the model to its original behavior.
model.clear_edits()
with model.trace("Vatican is in the city of"):
tokens = model.lm_head.output.argmax(dim=-1).save()
print(f"After clear_edits(): {model.tokenizer.decode(tokens[0, -1])}")
After clear_edits(): Rome
Attaching Custom Modules¶
You can attach your own PyTorch modules to the model — like an SAE or LoRA adapter — and wire them into the forward pass with model.edit(). Calling the attached module with hook=True runs its full __call__, so its submodules become fully instrumented: you can access its .output in a trace just like any other module.
# Define a low-rank adapter
class Adapter(nn.Module):
def __init__(self, dim, rank=16):
super().__init__()
self.down = nn.Linear(dim, rank, bias=False)
self.up = nn.Linear(rank, dim, bias=False)
nn.init.zeros_(self.up.weight) # Start as identity (no effect)
def forward(self, x):
return self.up(self.down(x))
# Attach it to layer 5 (on the same device as the layer)
device = next(model.transformer.h[5]._module.parameters()).device
torch.manual_seed(0) # the adapter's `down` init draws from the RNG -- seed it for reproducible outputs
model.transformer.h[5].adapter = Adapter(768, rank=16).to(device)
# Wire the adapter into the forward pass: layer output += adapter(layer input)
with model.edit() as (tracer, edited):
h5_input = edited.transformer.h[5].inputs[0][0]
adapter_out = edited.transformer.h[5].adapter(h5_input, hook=True)
edited.transformer.h[5].output[:] = edited.transformer.h[5].output + adapter_out
With zero-initialized weights, the adapter has no effect:
with model.trace("The Eiffel Tower is in the city of"):
orig_logits = model.lm_head.output.save()
with edited.trace("The Eiffel Tower is in the city of"):
adapter_result = edited.transformer.h[5].adapter.output.save()
edited_logits = edited.lm_head.output.save()
print(f"Adapter output shape: {adapter_result.shape}")
print(f"Adapter norm: {adapter_result.norm():.4f}")
print(f"Original prediction: {model.tokenizer.decode(orig_logits[0, -1].argmax(dim=-1))}")
print(f"Adapted prediction: {edited.tokenizer.decode(edited_logits[0, -1].argmax(dim=-1))}")
Adapter output shape: torch.Size([1, 10, 768]) Adapter norm: 0.0000 Original prediction: Paris Adapted prediction: Paris
After training or modifying the adapter weights, it changes the model's behavior:
# Simulate trained weights by setting non-zero values
torch.manual_seed(0)
nn.init.normal_(model.transformer.h[5].adapter._module.up.weight, std=10.0)
with edited.trace("The Eiffel Tower is in the city of"):
adapter_result = edited.transformer.h[5].adapter.output.save()
adapted_logits = edited.lm_head.output.save()
print(f"Adapter output norm: {adapter_result.norm():.4f}")
print(f"Adapted prediction: {edited.tokenizer.decode(adapted_logits[0, -1].argmax(dim=-1))}")
assert adapter_result.norm() > 0
Adapter output norm: 61750.8477 Adapted prediction: Hands
You can also intervene on the adapter itself during a trace — for example, zeroing out its contribution:
with edited.trace("The Eiffel Tower is in the city of"):
# Zero out the adapter's output, neutralizing it for this run
edited.transformer.h[5].adapter.output[:] = 0
logits = edited.lm_head.output.save()
print(f"Adapter zeroed: {edited.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")
Adapter zeroed: Paris
Edits Across Generation Steps¶
A stored edit applies at each location's first occurrence. Under generate that is prefill: the edit shapes the prompt's pass, and the remaining decoding steps run unedited. Put the block under the tracer's iter to reach every step. Bound the loop to the number of steps the run makes; an open iter[:] ends by asking for a step that never happens, which warns and drops anything written after the loop.
model.clear_edits()
prompt = "The Eiffel Tower is in the city of"
with model.edit(inplace=True):
model.transformer.h[9].output[:] = 0
with model.generate(prompt, max_new_tokens=4, do_sample=False) as tracer:
once = tracer.cache(modules=[model.transformer.h[9]])
print("plain edit, zeroed per step: ",
[bool((step == 0).all()) for step in once["model.transformer.h.9"].output])
model.clear_edits()
with model.edit(inplace=True) as tracer:
for _ in tracer.iter[:4]:
model.transformer.h[9].output[:] = 0
with model.generate(prompt, max_new_tokens=4, do_sample=False) as tracer:
every = tracer.cache(modules=[model.transformer.h[9]])
zeroed = [bool((step == 0).all()) for step in every["model.transformer.h.9"].output]
print("under tracer.iter[:4], zeroed:", zeroed)
assert all(zeroed)
model.clear_edits()
plain edit, zeroed per step: [True, False, False, False] under tracer.iter[:4], zeroed: [True, True, True, True]
A cache over the edited module is the cheapest way to check which of the two you have: it records one entry per step the module ran, so the list above is the edit's own coverage.
Which Model an Edit Belongs To¶
An edit is stored on the envoy you called .edit() on, and a trace replays only the edits of the envoy it is rooted at. So an edit stored on a layer does nothing when you trace the model, and the model's clear_edits() does not remove it either.
with model.transformer.h[9].edit(inplace=True):
model.transformer.h[9].output[:] = 0
with model.trace(prompt):
tokens = model.lm_head.output.argmax(dim=-1).save()
print("root trace:", model.tokenizer.decode(tokens[0, -1]))
print("model.clear_edits() leaves h[9] holding", len(model.transformer.h[9]._edits), "edit")
model.clear_edits()
assert len(model.transformer.h[9]._edits) == 1
model.transformer.h[9].clear_edits() # this is what clears it
assert len(model.transformer.h[9]._edits) == 0
root trace: Paris model.clear_edits() leaves h[9] holding 1 edit
To scope an edit to one layer of a whole-model run, store it on the model and write the layer's path inside the block, which is what every example above does.
When to use editing vs tracing
- Use
model.trace()for one-off interventions during a single forward pass - Use
model.edit()when you need the same intervention applied repeatedly across many forward passes - Use
model.edit()to wire custom modules (adapters, SAEs, probes) into the model's forward pass, making them automatically instrumented for tracing