Instrumenting an Open Model to Watch It Think

This page builds one instrument: hook every transformer block of a local model, capture what the attention and MLP sublayers each wrote into the residual stream, and print a per-layer table showing how much each contributed and what the model was leaning towards at that depth. It runs on a laptop CPU with a small model.

What you need

• Local weights. Non-negotiable. A hosted API returns tokens, not activations. Start with gpt2 (124M) — it runs on any CPU in a second or two per pass, and every technique here transfers upward unchanged.
• PyTorch and transformers. Nothing else. The whole instrument is register_forward_hook, which has been stable for years, rather than a named interpretability library whose API moves.
• Memory, if you scale up. Capturing every layer at every position costs n_layers × seq_len × d_model × bytes_per_element. For a 32-layer, 4,096-dimension model at fp16 over 2,000 tokens that is 32 × 2,000 × 4,096 × 2 ≈ 524 MB per captured tensor family — and you are capturing three. Capture the last position only, or a slice of layers, unless you need the whole grid.

Finding the module names

Every architecture puts its blocks somewhere different, and this is the step where copied tutorials break. Print the structure once and read it.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("gpt2")

# every module path, one per line — look for the repeated numbered block
for name, _ in model.named_modules():
    print(name)

# GPT-2 output (abridged):
#   transformer
#   transformer.wte
#   transformer.h            <- the block list
#   transformer.h.0
#   transformer.h.0.ln_1
#   transformer.h.0.attn
#   transformer.h.0.mlp
#   ...
#   transformer.ln_f         <- final norm
#   lm_head                  <- unembedding

# Llama-style models instead have:
#   model.layers.N.self_attn / model.layers.N.mlp, model.norm, lm_head

The helper below resolves both families, which is enough for the majority of open decoder models you are likely to load.

def get_parts(model):
    """Returns (blocks, attn_getter, mlp_getter, final_norm, unembed)."""
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return (model.transformer.h,
                lambda b: b.attn,
                lambda b: b.mlp,
                model.transformer.ln_f,
                model.lm_head)
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return (model.model.layers,
                lambda b: b.self_attn,
                lambda b: b.mlp,
                model.model.norm,
                model.lm_head)
    raise ValueError("unrecognised architecture; print named_modules() and add it")

The recorder

One class, three hook types. The only subtlety is that a block’s forward returns a tuple while a sublayer’s usually returns a bare tensor, so the hook normalises both.

import torch

def first_tensor(x):
    """Blocks return tuples, sublayers usually return tensors."""
    return x[0] if isinstance(x, tuple) else x

class Recorder:
    def __init__(self, model):
        self.model = model
        self.blocks, self.get_attn, self.get_mlp, self.norm, self.unembed = get_parts(model)
        self.residual = {}   # layer -> block output
        self.attn_out = {}   # layer -> what attention wrote
        self.mlp_out  = {}   # layer -> what the MLP wrote
        self.handles  = []

    def _store(self, store, i):
        def hook(module, args, output):
            store[i] = first_tensor(output).detach().float().cpu()
        return hook

    def __enter__(self):
        for i, block in enumerate(self.blocks):
            self.handles.append(block.register_forward_hook(
                self._store(self.residual, i)))
            self.handles.append(self.get_attn(block).register_forward_hook(
                self._store(self.attn_out, i)))
            self.handles.append(self.get_mlp(block).register_forward_hook(
                self._store(self.mlp_out, i)))
        return self

    def __exit__(self, *exc):
        for h in self.handles:
            h.remove()
        self.handles = []

Three details in that code are load-bearing. .detach() stops the captured tensors keeping the autograd graph alive, which is the difference between a few megabytes and running out of memory. .cpu() moves them off the accelerator so a long capture does not compete with the model for device memory. And the hooks are removed in __exit__: a leaked hook stays attached to the module for the lifetime of the process and silently slows every later forward pass while filling a dictionary nobody reads.

The report

from transformers import AutoTokenizer

@torch.no_grad()
def report(model, tok, text, position=-1, topk=3):
    rec = Recorder(model)
    ids = tok(text, return_tensors="pt")

    with rec:
        out = model(**ids)

    print(f"prompt: {text!r}")
    print(f"final prediction: {tok.decode([out.logits[0, position].argmax()])!r}\n")
    print(f"{'layer':>5} {'|resid|':>9} {'attn+':>7} {'mlp+':>7} "
          f"{'entropy':>8}  top tokens (logit lens)")

    n = len(rec.blocks)
    prev_norm = None
    for i in range(n):
        h    = rec.residual[i][0, position]
        attn = rec.attn_out[i][0, position]
        mlp  = rec.mlp_out[i][0, position]

        logits = rec.unembed(rec.norm(h.to(next(model.parameters()).dtype)))
        probs  = logits.float().softmax(-1)
        top    = probs.topk(topk)
        ent    = -(probs * probs.clamp_min(1e-12).log()).sum().item()
        toks   = " ".join(f"{tok.decode([t]).strip()!r}" for t in top.indices)

        # each sublayer's contribution relative to the stream it was added to
        rel_attn = attn.norm().item() / h.norm().item()
        rel_mlp  = mlp.norm().item()  / h.norm().item()

        print(f"{i:>5} {h.norm().item():>9.1f} {rel_attn:>7.3f} {rel_mlp:>7.3f} "
              f"{ent:>8.2f}  {toks}")
        prev_norm = h.norm().item()

tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()
report(model, tok, "The Eiffel Tower is located in the city of")

Reading it

Four things in that table are worth your attention.

The residual norm grows with depth, usually substantially. This is expected: every layer adds to the stream and nothing subtracts on average. It is also why a fixed-size intervention means something different at layer 2 than at layer 20, and why steering coefficients do not transfer across layers.

The relative contribution columns show where the work happens. Layers whose attention or MLP contribution is a large fraction of the running stream are doing something; layers where both are small are mostly passing the state along. On most models this is not uniform, and the uneven distribution is the most immediately surprising thing the instrument shows you.

Entropy falls in steps, not smoothly. The layer where it drops sharply is where the model committed. Compare an easy prompt with a hard one of the same length and watch the commitment point move.

Candidates that rise and then fall mean something later actively suppressed them. That is the single best lead this instrument produces, and the follow-up is activation patching on the layers between the peak and the fall.

Instrumenting generation, not just one pass

Hooks fire during generate as well, and the behaviour surprises people. With the KV cache active, every step after the first passes only the newest token, so the captured tensors have sequence length 1 from step two onward. Your position=-1 indexing keeps working; any indexing into earlier positions does not.

class StepRecorder(Recorder):
    """Collects one row per generated token instead of overwriting."""
    def __init__(self, model):
        super().__init__(model)
        self.steps = []

    def _store(self, store, i):
        def hook(module, args, output):
            store[i] = first_tensor(output).detach().float().cpu()
        return hook

    def snapshot(self, layer, position=-1):
        self.steps.append(self.residual[layer][0, position].clone())

# usage: run generate inside the context, then inspect rec.residual after
# each call, or subclass further to append on every forward of the last block.
#
# NOTE: with use_cache=True, seq_len == 1 after the first step. If you want
# aligned positions across the whole sequence, either pass use_cache=False
# (much slower, full recomputation each step) or re-run one forward pass over
# the completed sequence afterwards. The second is almost always what you want.

That last comment is the practical advice: generate normally, then run a single instrumented forward pass over prompt-plus-completion. You get a clean, aligned grid and you avoid reasoning about cache behaviour. The only thing lost is the model’s state at the moment of each decision, which matters rarely.

From watching to intervening

The same hook mechanism modifies rather than records — return a value from the hook and it replaces the module’s output. This is the whole of ablation, patching and steering.

def ablate_attention(layer_idx, blocks, get_attn):
    """Zero one layer's attention contribution and see what breaks."""
    def hook(module, args, output):
        if isinstance(output, tuple):
            return (torch.zeros_like(output[0]),) + tuple(output[1:])
        return torch.zeros_like(output)
    return get_attn(blocks[layer_idx]).register_forward_hook(hook)

blocks, get_attn, get_mlp, norm, unembed = get_parts(model)
ids = tok("The Eiffel Tower is located in the city of", return_tensors="pt")

with torch.no_grad():
    base = model(**ids).logits[0, -1].softmax(-1)

for layer in range(len(blocks)):
    h = ablate_attention(layer, blocks, get_attn)
    with torch.no_grad():
        p = model(**ids).logits[0, -1].softmax(-1)
    h.remove()
    print(f"layer {layer:>2}  top token {tok.decode([p.argmax()])!r:>12}  "
          f"total variation from baseline {0.5 * (p - base).abs().sum():.3f}")

Zero ablation is the blunt instrument — zero is an off-distribution value, so a large effect may mean you broke the model rather than removed a function. Mean ablation over a set of prompts is the better default, and the patching page covers why, along with the self-repair result that makes a small ablation effect much harder to interpret than it looks.

The six things that go wrong

1. A hook returns a bare tensor to a module that returns a tuple. The next layer receives the wrong type and you get an unhelpful error deep in the model, or worse, silently wrong behaviour. Always reconstruct the tuple.
2. Hooks are never removed. They persist on the module. Use a context manager, as above, or expect a notebook whose results drift as leaked hooks accumulate.
3. Captured tensors keep the graph. Without .detach() inside torch.no_grad()-free code you retain the autograd graph for every layer and run out of memory on the second prompt.
4. Dtype mismatch on the lens. Captured tensors cast to float32 will not multiply against fp16 or bf16 weights. Cast back before calling the unembedding, as the report does.
5. The off-by-one on hidden states. If you use output_hidden_states=True instead of hooks, entry zero is the embedding output, so entry i+1 is layer i. Verify by checking the last entry reproduces the model’s real logits.
6. Attention weights come back empty. Fused kernels do not materialise the weight matrix. Load with attn_implementation="eager" if you need them.