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 |
Cache¶
tracer.cache() collects activations from many modules at once during a trace, without writing a .save() for each one. Every module's input and output already passes through a single handoff where interventions are applied; the cache keeps the values for the modules you named as they go by, after any intervention.
Use it when you want the same value from many modules (every layer's hidden state for a logit lens or a probe), or activations from every generation step. It costs what the equivalent .save() calls cost. What it saves you is the loop.
Setup¶
import torch
import nnsight
from nnsight.modeling.transformers import TransformersModel
model = TransformersModel("openai-community/gpt2", device_map="auto", dispatch=True)
Caching All Modules¶
Calling tracer.cache() with no arguments captures the output of every module the run reaches. The returned object is a CacheView — a path- and attribute-addressable view over the recorded values. It is already .save()d for you, so it survives past the trace.
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache()
print(f"Modules cached: {len(cache.keys())}")
print(f"First few keys: {list(cache.keys())[:5]}")
print(f"Layer 0 output shape: {cache['model.transformer.h.0'].output.shape}")
Modules cached: 151 First few keys: ['model.transformer.wte', 'model.transformer.wpe', 'model.transformer.drop', 'model.transformer.h.0.ln_1', 'model.transformer.h.0.attn.c_attn'] Layer 0 output shape: torch.Size([1, 10, 768])
cache.keys() lists the cached module paths (all of them, at the root), in the order the run reaches them rather than the order you passed them. A GPT-2 block returns its hidden-state tensor directly, so cache[path].output is that tensor.
Open the cache before you read or write any activation. Opening one afterwards raises
ValueError: tracer.cache() must be declared before reading or modifying a model value. The routes are fixed before the model starts, so the ordering is checked rather than assumed: a cache never quietly captures fewer modules than you asked for.
Selective Caching¶
Pass a list of modules (Envoy references or path strings) to cache only what you need. This keeps memory usage low and is the typical pattern for real experiments.
# By Envoy reference
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache(modules=[
model.transformer.h[0],
model.transformer.h[5],
model.transformer.h[11],
model.lm_head,
])
print(f"Cached: {list(cache.keys())}")
Cached: ['model.transformer.h.0', 'model.transformer.h.5', 'model.transformer.h.11', 'model.lm_head']
# By path string — useful when iterating programmatically
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache(modules=[f"model.transformer.h.{i}.attn" for i in range(12)])
print(f"Cached {len(cache.keys())} attention modules")
Cached 12 attention modules
Outputs, Inputs, or Both¶
By default, cache() records module outputs. Set include_inputs=True to also capture each module's inputs, or set include_output=False to capture only inputs.
The view exposes three read accessors:
.output— the module's forward output.inputs— a(args, kwargs)tuple of all inputs (Noneunlessinclude_inputs=True).input— shorthand for the first positional (or first keyword) argument
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache(
modules=[model.transformer.h[0], model.transformer.h[1]],
include_inputs=True,
)
# Layer 1's first input should equal layer 0's output
print("h[0].output == h[1].input:",
torch.equal(cache["model.transformer.h.0"].output,
cache["model.transformer.h.1"].input))
h[0].output == h[1].input: True
Use Case: Logit Lens in One Pass¶
A classic mech-interp pattern: project every layer's hidden state through the final layer norm and lm_head to see what token the model would predict at each layer. With cache(), you grab all 12 hidden states in a single forward pass.
prompt = "The Eiffel Tower is in the city of"
with model.trace(prompt) as tracer:
cache = tracer.cache(modules=[layer for layer in model.transformer.h])
# Apply the logit lens *outside* the trace — cheap arithmetic on the saved
# tensors. The cache stores them on CPU by default, so read the small final
# modules on the same device.
device = next(model.transformer.ln_f.parameters()).device
print(f"Prompt: {prompt!r}\n")
with torch.no_grad():
for i in range(12):
hs = cache[f"model.transformer.h.{i}"].output.to(device)
logits = model.lm_head(model.transformer.ln_f(hs))
top = logits[0, -1].argmax(dim=-1)
print(f"Layer {i:2d}: {model.tokenizer.decode(top)!r}")
Prompt: 'The Eiffel Tower is in the city of' Layer 0: ' the' Layer 1: ' the' Layer 2: ' the' Layer 3: ' the' Layer 4: ' the' Layer 5: ' the' Layer 6: ' East' Layer 7: ' Ing' Layer 8: ' Rome' Layer 9: ' London' Layer 10: ' Paris' Layer 11: ' Paris'
Caching During Generation¶
A module reached once per step accumulates one entry per step. len(cache[path]) is the visit (step) count, and cache[path].output becomes a list — one value per step.
with model.generate("The Eiffel Tower is in", max_new_tokens=4) as tracer:
cache = tracer.cache(modules=[model.transformer.h[-1]])
n_steps = len(cache["model.transformer.h.11"])
outputs = cache["model.transformer.h.11"].output # list of per-step outputs
print(f"Number of generation steps cached: {n_steps}")
for i, out in enumerate(outputs):
print(f" step {i}: hidden state shape {out.shape}")
Number of generation steps cached: 4 step 0: hidden state shape torch.Size([1, 7, 768]) step 1: hidden state shape torch.Size([1, 1, 768]) step 2: hidden state shape torch.Size([1, 1, 768]) step 3: hidden state shape torch.Size([1, 1, 768])
The first step contains all prompt tokens; subsequent steps contain only the newly generated token (when KV caching is active).
Single vs Multiple Visits¶
cache[path].output unwraps automatically: a single visit returns the value directly, multiple visits return a list. len(cache[path]) is the visit count — so a plain forward gives you the value itself (whatever the module returns), and a generation loop gives you a list of them.
# Single forward -> one visit -> the value directly (print the type to see what this module returns)
with model.trace("Hello") as tracer:
cache = tracer.cache(modules=[model.transformer.h[-1]])
print("single visit ->", type(cache["model.transformer.h.11"].output).__name__,
"| len:", len(cache["model.transformer.h.11"]))
# Generation -> N visits -> a list
with model.generate("Hello", max_new_tokens=3) as tracer:
cache = tracer.cache(modules=[model.transformer.h[-1]])
print("multiple visits ->", type(cache["model.transformer.h.11"].output).__name__,
"| len:", len(cache["model.transformer.h.11"]))
single visit -> Tensor | len: 1
multiple visits -> list | len: 3
Verifying an Intervention¶
The cache observes post-intervention values, so it's a convenient way to confirm a patch landed.
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache(modules=[model.transformer.h[5]])
# Zero out layer 5's hidden state
model.transformer.h[5].output[:] = 0
is_zero = torch.all(cache["model.transformer.h.5"].output == 0).item()
print(f"Layer 5 output is all zeros: {is_zero}")
Layer 5 output is all zeros: True
Multiple Caches in One Trace¶
You can create separate caches for different module groups in the same trace. They observe independently and are returned as separate views.
with model.trace("The Eiffel Tower is in the city of") as tracer:
attn_cache = tracer.cache(modules=[layer.attn for layer in model.transformer.h])
mlp_cache = tracer.cache(modules=[layer.mlp for layer in model.transformer.h])
print(f"Attention modules cached: {len(attn_cache.keys())}")
print(f"MLP modules cached: {len(mlp_cache.keys())}")
print(f"Sample attn key: {list(attn_cache.keys())[0]}")
print(f"Sample mlp key: {list(mlp_cache.keys())[0]}")
Attention modules cached: 12 MLP modules cached: 12 Sample attn key: model.transformer.h.0.attn Sample mlp key: model.transformer.h.0.mlp
Caching Inside an Invoke¶
A cache opened inside tracer.invoke(...) records that invoke's rows. It records them at the whole batch's padded length, and nothing in the view marks which positions are padding.
with model.trace() as tracer:
with tracer.invoke("the cat"): # 2 tokens
short = tracer.cache(modules=[model.transformer.h[0]])
with tracer.invoke("a much longer prompt here about many things"): # 8 tokens
pass
out = short["model.transformer.h.0"].output
print("shape:", tuple(out.shape))
print("per-position norms:", [round(float(n), 1) for n in out[0].norm(dim=-1)])
assert out.shape[1] == 8
shape: (1, 8, 768) per-position norms: [192.4, 192.4, 192.4, 192.4, 192.4, 192.4, 136.7, 56.7]
The prompt is two tokens long, so six of those eight positions are padding, and at block 0 each of them carries a larger norm than either real token. A mean or a max over the sequence axis of that capture is mostly padding. Index from the right ([:, -1] is stable, since the padding is on the left) or mask before you reduce. Batching covers where the padding comes from.
To capture the whole combined batch instead, open the cache in an empty tracer.invoke(), which sees every row:
with model.trace() as tracer:
with tracer.invoke("the cat"):
pass
with tracer.invoke("a much longer prompt here about many things"):
pass
with tracer.invoke():
batch = tracer.cache(modules=[model.transformer.h[0]])
full = batch["model.transformer.h.0"].output
print("whole batch:", tuple(full.shape))
assert full.shape[0] == 2
whole batch: (2, 8, 768)
Memory Management with device and dtype¶
For large models or long sequences, the activations themselves can dominate memory. cache() accepts device and dtype arguments that are applied to every captured tensor — letting you offload to CPU and/or downcast to a smaller dtype as values are recorded.
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache(
modules=[layer for layer in model.transformer.h],
device=torch.device("cpu"), # move captured tensors here (default: CPU)
dtype=torch.float16, # downcast in the same step (default: keep)
)
sample = cache["model.transformer.h.0"].output
print(f"device: {sample.device}")
print(f"dtype: {sample.dtype}")
print(f"shape: {sample.shape}")
device: cpu dtype: torch.float16 shape: torch.Size([1, 10, 768])
The defaults (device=cpu, detach=True, non_blocking=False) are memory-friendly and safe: tensors are detached from the autograd graph and copied off the compute device as they're captured. Pass device=None to leave them where they are.
non_blocking=True makes that copy asynchronous and is not safe on its own — nothing synchronises before you read the cache, so a capture off a CUDA device can be read while the copy is still in flight (the window scales with copy size, so it is invisible on small examples and corrupts corpus-scale runs). Only pass it if you call torch.cuda.synchronize() yourself before reading.
Access Patterns¶
The view supports both dictionary-style (by path) and attribute-style (by navigating the tree) access. They return the exact same data — pick whichever reads better. Tree navigation resolves against the model's envoy tree, so ModuleList indices and renamed modules work the same way they do on the model.
with model.trace("The Eiffel Tower is in the city of") as tracer:
cache = tracer.cache(modules=[layer for layer in model.transformer.h] + [model.lm_head])
# By path
a = cache["model.transformer.h.0"].output
# By navigation (the leading model name is optional: cache.transformer.h[0] also works)
b = cache.model.transformer.h[0].output
print("Same tensor:", torch.equal(a, b))
print("LM head shape:", cache.model.lm_head.output.shape)
Same tensor: True LM head shape: torch.Size([1, 10, 50257])
How Cache Works¶
Every module input and output already flows through the interleaver's single handoff point, where interventions are applied first. tracer.cache() registers on that handoff and keeps the values for the modules you selected as the run passes them.
That is the same handoff a .output read uses, so a cache and a hand-written save loop over the same modules take the same time: 19.0 ms against 18.8 ms for GPT-2's twelve blocks at batch 32 x 64 tokens, A and B interleaved, minimum of five runs. Reach for the cache because you do not have to write the loop, and because it absorbs the cases a loop makes tedious: modules named programmatically, inputs as well as outputs, one entry per generation step.
Pass modules= on anything real. Caching all 151 of GPT-2's modules takes 381 ms and 1.1 GiB against 19.0 ms and 36 MiB for its twelve blocks.
You can mix the two in the same trace: cache the bulk of your activations, and use .output / .save() where you need to edit a value or branch on it.
When to use cache
- Bulk activation collection — grabbing hidden states from every layer for probing, logit lens, or SAE training data
- Selective sweeps — caching just attention or just MLP outputs across all layers
- Memory-constrained runs —
device="cpu"+dtype=torch.float16keeps compute-device memory flat while you record - Intervention verification — confirm a patch took effect by reading the post-intervention value out of the cache
- Generation tracking — collect per-step activations across decoding without writing per-step code