Info
Last Execution: 2026-09-02
| Package | Version |
|---|---|
| nnsight | 0.8.0 |
| Python | 3.12.13 |
| torch | 2.13.0+cu126 |
| transformers | 5.15.0.dev0 |
Accessing Intermediate Operations¶
.output and .input let you hook into a module's inputs and outputs. But what about the operations inside a module's forward pass? .source exposes every call site and every assignment in a module's forward as a hookable operation — function calls, method calls, tensor operations, and the values they are bound to — so you can read, replace, or skip a value that lives between two operations, with no submodule to attach to.
Setup¶
from nnsight.modeling.transformers import TransformersModel
model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
Discovering Operations¶
Print .source on any module to see its forward method with every hookable operation labeled at its call site or assignment (this works outside a trace).
print(model.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)
hidden_states_0 -> + ...
self_act_0 -> 1 hidden_states = self.act(hidden_states)
hidden_states_1 -> + ...
self_c_proj_0 -> 2 hidden_states = self.c_proj(hidden_states)
hidden_states_2 -> + ...
self_dropout_0 -> 3 hidden_states = self.dropout(hidden_states)
hidden_states_3 -> + ...
4 return hidden_states
5
Each labeled line (like self_c_fc_0, self_act_0) is an operation you can access inside a trace. Operation names are the full dotted callee joined with _, plus a per-name occurrence index in execution order — so self.c_fc(...) becomes self_c_fc_0 and a second call to the same callee would be ..._1.
Larger modules have more operations. Here's the attention module:
print(model.transformer.h[0].attn.source)
* def forward(
0 self,
1 hidden_states: tuple[torch.FloatTensor] | None,
2 past_key_values: Cache | None = None,
3 attention_mask: torch.FloatTensor | None = None,
4 encoder_hidden_states: torch.Tensor | None = None,
5 encoder_attention_mask: torch.FloatTensor | None = None,
6 output_attentions: bool | None = False,
7 **kwargs,
8 ) -> tuple[torch.Tensor | tuple[torch.Tensor], ...]:
is_cross_attention_0 -> 9 is_cross_attention = encoder_hidden_states is not None
10 if past_key_values is not None:
isinstance_0 -> 11 if isinstance(past_key_values, EncoderDecoderCache):
past_key_values_is_updated_get_0 -> 12 is_updated = past_key_values.is_updated.get(self.layer_idx)
is_updated_0 -> + ...
13 if is_cross_attention:
14 # after the first generated id, we can subsequently re-use all key/value_layer from cache
curr_past_key_values_0 -> 15 curr_past_key_values = past_key_values.cross_attention_cache
16 else:
curr_past_key_values_1 -> 17 curr_past_key_values = past_key_values.self_attention_cache
18 else:
curr_past_key_values_2 -> 19 curr_past_key_values = past_key_values
20
21 if is_cross_attention:
hasattr_0 -> 22 if not hasattr(self, "q_attn"):
ValueError_0 -> 23 raise ValueError(
24 "If class is used as cross attention, the weights `q_attn` have to be defined. "
25 "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
26 )
self_q_attn_0 -> 27 query_states = self.q_attn(hidden_states)
query_states_0 -> + ...
attention_mask_0 -> 28 attention_mask = encoder_attention_mask
29
30 # Try to get key/value states from cache if possible
31 if past_key_values is not None and is_updated:
key_states_0 -> 32 key_states = curr_past_key_values.layers[self.layer_idx].keys
value_states_0 -> 33 value_states = curr_past_key_values.layers[self.layer_idx].values
34 else:
self_c_attn_0 -> 35 key_states, value_states = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
split_0 -> + ...
shape_kv_0 -> 36 shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
key_states_view_0 -> 37 key_states = key_states.view(shape_kv).transpose(1, 2)
transpose_0 -> + ...
key_states_1 -> + ...
value_states_view_0 -> 38 value_states = value_states.view(shape_kv).transpose(1, 2)
transpose_1 -> + ...
value_states_1 -> + ...
39 else:
self_c_attn_1 -> 40 query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
split_1 -> + ...
shape_kv_1 -> 41 shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
key_states_view_1 -> 42 key_states = key_states.view(shape_kv).transpose(1, 2)
transpose_2 -> + ...
key_states_2 -> + ...
value_states_view_1 -> 43 value_states = value_states.view(shape_kv).transpose(1, 2)
transpose_3 -> + ...
value_states_2 -> + ...
44
shape_q_0 -> 45 shape_q = (*query_states.shape[:-1], -1, self.head_dim)
query_states_view_0 -> 46 query_states = query_states.view(shape_q).transpose(1, 2)
transpose_4 -> + ...
query_states_1 -> + ...
47
48 if (past_key_values is not None and not is_cross_attention) or (
49 past_key_values is not None and is_cross_attention and not is_updated
50 ):
curr_past_key_values_update_0 -> 51 key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
52 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
53 if is_cross_attention:
past_key_values_is_updated_0 -> 54 past_key_values.is_updated[self.layer_idx] = True
55
using_eager_0 -> 56 using_eager = self.config._attn_implementation == "eager"
ALL_ATTENTION_FUNCTIONS_get_interface_0 -> 57 attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
attention_interface_0 -> + ...
58 self.config._attn_implementation, eager_attention_forward
59 )
60
61 if using_eager and self.reorder_and_upcast_attn:
self__upcast_and_reordered_attn_0 -> 62 attn_output, attn_weights = self._upcast_and_reordered_attn(
63 query_states, key_states, value_states, attention_mask
64 )
65 else:
attention_interface_1 -> 66 attn_output, attn_weights = attention_interface(
67 self,
68 query_states,
69 key_states,
70 value_states,
71 attention_mask,
72 dropout=self.attn_dropout.p if self.training else 0.0,
73 scaling=self.scaling,
74 **kwargs,
75 )
76
attn_output_reshape_0 -> 77 attn_output = attn_output.reshape(*attn_output.shape[:-2], -1).contiguous()
contiguous_0 -> + ...
attn_output_0 -> + ...
self_c_proj_0 -> 78 attn_output = self.c_proj(attn_output)
attn_output_1 -> + ...
self_resid_dropout_0 -> 79 attn_output = self.resid_dropout(attn_output)
attn_output_2 -> + ...
80
81 return attn_output, attn_weights
82
Getting an Intermediate Value¶
Access any labeled operation's .output inside a trace — just like you would with a module. Here we grab the hidden states right after the GELU activation, before the down projection.
with model.trace("The Eiffel Tower is in the city of"):
post_gelu = model.transformer.h[0].mlp.source.self_act_0.output.save()
# GPT-2's MLP widens to 4x the hidden size, and this is the value after the GELU.
assert post_gelu.shape[-1] == 4 * model.config.n_embd
print(f"Post-GELU shape: {post_gelu.shape}")
Post-GELU shape: torch.Size([1, 10, 3072])
How source works
On first .source access, nnsight rewrites the module's forward so each call fn(*args, **kwargs) — and each assignment x = value, through an identity — is bracketed by the interleaver. Calls and assignments share one occurrence counter per name. Each operation then exposes .input / .inputs / .output / .skip — the same handles a module has, one level finer. The instrumentation is installed on first .source access and stays on the module for the rest of the process. Outside a trace every operation calls straight through, which is cheap but not free — sourcing every block, attention and MLP of GPT-2 costs about 6% on an untraced forward pass.
Values That Are Not a Call¶
A residual add, a product, a running state — anything a forward assigns but never returns from a call — is an operation named after its target, on the same per-name counter calls use. GPT-2's block adds the attention output back into the residual stream with a plain +. There is no call to hook, but the assignment is hidden_states_1, the residual stream between attention and the MLP:
import torch
block = model.transformer.h[0]
_ = block.source # instrument the block now, not mid-trace
with model.trace("The Eiffel Tower is in the city of"):
resid_pre = block.input # requests go in execution order
attn_out = block.attn.output[0]
resid_mid = block.source.hidden_states_1.output.save() # hidden_states = attn_output + residual
gap = (resid_mid - (attn_out + resid_pre)).abs().max().save()
assert resid_mid.shape == (1, 10, model.config.n_embd)
assert gap.item() < 1e-5 # hidden_states_1 really is the post-attention residual
print(f"resid_mid shape: {resid_mid.shape}")
print(f"max |resid_mid - (attn_out + resid_pre)|: {gap.item()}")
resid_mid shape: torch.Size([1, 10, 768]) max |resid_mid - (attn_out + resid_pre)|: 0.0
Note the bare block.source on its own line, before the trace opens. Reaching for .source is what rewrites the module's forward, and a forward can only be rewritten before it runs. Inside a trace that reads something else first, the model is already past the block by the time the request lands, and you get an OutOfOrderError that reads like an ordering mistake when it is really a timing one.
Touching the attribute outside the trace instruments the module and costs nothing — no forward pass, one line. It is per module, so a sweep over every layer wants one warm-up per layer.
Setting an Intermediate Value¶
You can also modify intermediate values. Assigning to an operation's .output replaces its value for the rest of the forward. Here we zero out the MLP's post-GELU activations at layer 11 to see how it affects the prediction:
with model.trace("The Eiffel Tower is in the city of"):
normal_logits = model.lm_head.output.save()
with model.trace("The Eiffel Tower is in the city of"):
# Zero the MLP's GELU output at layer 11
model.transformer.h[11].mlp.source.self_act_0.output[:] = 0
modified_logits = model.lm_head.output.save()
assert not torch.equal(normal_logits, modified_logits) # the edit reached the logits
print(f"Normal: {model.tokenizer.decode(normal_logits[0, -1].argmax(dim=-1))}")
print(f"Zeroed GELU: {model.tokenizer.decode(modified_logits[0, -1].argmax(dim=-1))}")
Normal: Paris Zeroed GELU: London
Patching Between Layers¶
Transfer an intermediate value from one layer's operation to another:
with model.trace("The Eiffel Tower is in the city of"):
# Capture layer 0's post-GELU activations
gelu_0 = model.transformer.h[0].mlp.source.self_act_0.output
# Patch them into layer 5
model.transformer.h[5].mlp.source.self_act_0.output = gelu_0
logits = model.lm_head.output.save()
assert not torch.equal(logits, normal_logits) # the patch changed the run
print(f"Patched MLP prediction: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")
Patched MLP prediction: London
Recursive Source Tracing¶
.source works recursively. If a labeled operation calls a plain Python function, you can chain .source again to expose its operations. For example, GPT-2's attention module calls an attention_interface function, which internally calls scaled_dot_product_attention — here we drill in and read the query tensor it receives:
with model.trace("The Eiffel Tower is in the city of"):
# Drill into the attention interface -> the SDPA call's input (the query tensor)
sdpa = model.transformer.h[0].attn.source.attention_interface_1.source
query = sdpa.torch_nn_functional_scaled_dot_product_attention_0.input.save()
assert query.shape == (1, model.config.n_head, 10, model.config.n_embd // model.config.n_head)
print(f"Query shape: {query.shape}") # [batch, heads, seq_len, head_dim]
Query shape: torch.Size([1, 12, 10, 64])
Recursive .source is trace-only — the called function is resolved from the live value flowing through the call at run time, so op.source outside a trace raises SourceNotAvailable.
Operations That Never Run¶
A .source listing is the whole forward, branches included — and a branch that this model's configuration never enters still gets its operations listed and named. Ask for one and the request waits at a location the model never reaches, so the trace ends with the same OutOfOrderError you would get by reading two values in the wrong order.
GPT-2's attention lists 50 operations and runs 28 of them. The rest are the cross-attention and cache-hit paths, and their labels sit one character from the live ones.
from nnsight.intervention.interleaver import OutOfOrderError
attn = model.transformer.h[0].attn
_ = attn.source
names = list(attn.source.names)
live = []
for name in names:
try:
with model.trace("The Eiffel Tower is in the city of"):
getattr(attn.source, name).output.save()
live.append(name)
except OutOfOrderError:
pass # the model never reached this operation
assert len(names) == 50 and len(live) == 28
print(f"{len(live)} of {len(names)} listed operations actually run")
print("dead:", [n for n in names if n not in live][:6], "...")
28 of 50 listed operations actually run dead: ['past_key_values_is_updated_get_0', 'is_updated_0', 'curr_past_key_values_0', 'curr_past_key_values_1', 'hasattr_0', 'ValueError_0'] ...
transpose_0 is in that dead list: it is the key transpose on the cross-attention path. The one GPT-2 performs is transpose_2, two occurrences later. Requesting the wrong one of the pair fails with a message about the model having run past the location, which sends you looking at your request order instead of at the branch.
When an operation you can see in the listing reports as run past, read the lines around it. If it sits under an if whose condition is false for your model — is_cross_attention, an is_updated cache hit — it is dead, and the live twin is usually the next occurrence of the same name.
with model.trace("The Eiffel Tower is in the city of"):
key = attn.source.transpose_2.output.save()
head_dim = model.config.n_embd // model.config.n_head
assert key.shape == (1, model.config.n_head, 10, head_dim)
print(f"key states: {key.shape} # [batch, heads, seq, head_dim]")
key states: torch.Size([1, 12, 10, 64]) # [batch, heads, seq, head_dim]
Viewing a specific operation
Print a specific operation to see it highlighted in its surrounding context:
print(model.transformer.h[0].mlp.source.self_c_proj_0)
This shows the operation flagged with --> / <-- and surrounding lines for context.
Don't chain .source into a submodule call
If a .source listing shows a submodule call (like self.c_proj), access that submodule directly — drilling into it with .source raises SourceNotAvailable:
# Wrong — don't chain .source into a submodule call
model.transformer.h[0].mlp.source.self_c_proj_0.source
# Correct — access the submodule directly
model.transformer.h[0].mlp.c_proj.source