Locations

A document names a place in a sites entry: a component, a layer where that component is one per layer, and optionally which heads or units of it. That is all the document says. Which child module spells the place on this checkpoint, which side of that module carries the tensor, how wide it is and where it sits in the forward pass are facts about the architecture, and the compiler asks nnterp for them before any weights are loaded.

There are 21 component names. Seventeen go through an nnterp accessor; four are mini's own interior rows — a tensor that never crosses a module boundary, addressed as one call inside a forward through nnsight's .source. (One accessor, the attention pattern, is an operation on nnterp's side too; mini does not have to know, because nnterp reaches it.) The table is a floor and not a fence: a name mini has never heard of but model.internals has is addressable as well.

At a glance #

explain compiles a document against the model, with no weights, and prints the place every read and every write resolved to. documents/v2/head_patching.json patches one attention head at a time and scores the result at the model's head:

$ uv run causalab-mini explain documents/v2/head_patching.json
<root>: Plan
  heads=0: Plan
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
          read  'v_cf' at attention_premix[0] pos={index:-1} features=[0]/4 via 'identity'
        forward 'patched' on 'base'  (4 rows x 11 tokens)
          write 'patch' at attention_premix[0] pos={index:-1} features=[0]/4 swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'
  …
  heads=3: Plan
    score: Observe  metrics=['iia/match', 'logit_diff/logit_diff']  saves=['iia.json', 'logit_diff.json']
        forward 'original' on 'counterfactual'  (4 rows x 11 tokens)
          read  'v_cf' at attention_premix[0] pos={index:-1} features=[3]/4 via 'identity'
        forward 'patched' on 'base'  (4 rows x 11 tokens)
          write 'patch' at attention_premix[0] pos={index:-1} features=[3]/4 swap(v_cf) via 'identity'
          read  'logits' at lm_head pos={index:-1} via 'identity'

attention_premix[0] is a component and a layer; features=[0]/4 is head 0 of four; lm_head carries no layer because there is only one of it. The whole document, and what it produces, is at the bottom of this page.

The site field #

Both authoring formats spell the layer as a one-element list, and both refuse a longer one: layers must be a one-element band; a band spanning several layers is one address and is not implemented. To sweep depth, sweep the site — see Layer sweep.

A layer outside the model's stack is refused when the document is compiled, by number: site 'target': layer 24 is outside the model's 16 layers.

The component list #

What each name is, and what the tensor at it is. Width is the model attribute nnterp publishes the feature axis as; a per-head tensor's width is that attribute times the head count, because every head is handed on side by side.

Four of these are the residual stream — embeddings, block_input, block_output, ln_final — and a read with "view": "logits" is refused anywhere else, because pushing a tensor through the final norm and the head only means something for a tensor those two expect.

Where each one lands on a checkpoint #

locate is what turns a name into an address, and what it stamps in comes from nnterp: the module this family spells the place with, which side of it, and the rank that puts reads in forward order. module is relative to the layer for a per-layer place and to the model for a whole-model one; an interior has no module, and carries the .source operation instead.

from causalab_mini import address
from causalab_mini.engine import NNterpEngine
from causalab_mini.plan.spec import Model

e = NNterpEngine.load(Model(key="hf-internal-testing/tiny-random-LlamaForCausalLM",
    revision="9fb191250dd56d0ba7ec9785a025ed29c03d5998", dtype="fp32",
    attn_implementation="eager"), dispatch=False)
g = NNterpEngine.load(Model(key="hf-internal-testing/tiny-random-gpt2",
    revision="71034c5d8bde858ff824298bdedc65515b97d2b9", dtype="fp32",
    attn_implementation="eager"), dispatch=False)
for name, row in address.describe().items():
    layer = 0 if row["layered"] else None
    a, b = e.locate(name, layer), g.locate(name, layer)
    print(f"{name:22s} rank={str(a.rank):8s} io={str(a.io):7s} llama={a.module!r:32s} gpt2={b.module!r}")

attention_query        rank=(0, 11)  io=None    llama=None                             gpt2=None
attention_key          rank=(0, 12)  io=None    llama=None                             gpt2=None
attention_scores       rank=(0, 15)  io=None    llama=None                             gpt2=None
attention_probs        rank=(0, 20)  io=output  llama='self_attn'                      gpt2='self_attn'
attention_z            rank=(0, 22)  io=None    llama=None                             gpt2=None
input_ids              rank=(-1, -2) io=input   llama='embed_tokens'                   gpt2='embed_tokens'
embeddings             rank=(-1, -1) io=output  llama='embed_tokens'                   gpt2='embed_tokens'
block_input            rank=(0, 0)   io=input   llama=''                               gpt2=''
attention_input_norm   rank=(0, 5)   io=output  llama='input_layernorm'                gpt2='ln_1'
attention_premix       rank=(0, 25)  io=input   llama='self_attn.o_proj'               gpt2='self_attn.c_proj'
attention_output       rank=(0, 30)  io=output  llama='self_attn'                      gpt2='self_attn'
block_mid              rank=(0, 33)  io=input   llama='post_attention_layernorm'       gpt2='ln_2'
mlp_input_norm         rank=(0, 36)  io=output  llama='post_attention_layernorm'       gpt2='ln_2'
mlp_input              rank=(0, 40)  io=input   llama='mlp'                            gpt2='mlp'
mlp_activation         rank=(0, 44)  io=output  llama='mlp.act_fn'                     gpt2='mlp.act'
mlp_neuron_output      rank=(0, 47)  io=input   llama='mlp.down_proj'                  gpt2='mlp.c_proj'
mlp_output             rank=(0, 50)  io=output  llama='mlp'                            gpt2='mlp'
block_output           rank=(0, 60)  io=output  llama=''                               gpt2=''
ln_final               rank=(2, 1000) io=output  llama='ln_final'                       gpt2='ln_final'
lm_head                rank=(2, 1001) io=output  llama='lm_head'                        gpt2='lm_head'
logits                 rank=(2, 1002) io=output  llama=''                               gpt2=''

Two families whose module trees share no path, one set of names. The module column is the whole of what standardization buys here: block_mid is the input of post_attention_layernorm on Llama and of ln_2 on GPT-2, and the document says block_mid either way.

rank is (layer, order), nnterp's internals.rank, and reads are issued in that order because nnsight cannot reach back to a value the model has already passed. The whole-model places sort outside the stack: the embeddings at layer −1, the final norm and the two head rows at layer N. Mini keeps its own numbering only for the four interiors nnterp does not address, interleaved into nnterp's.

The whole-model places #

Five components are one per model and take no layers: input_ids, embeddings, ln_final, lm_head and logits. The protocol format checks both directions — lm_head takes no layers, block_output is addressed at one layer.

input_ids is the only read-only place: it is the model's input, and swapping a float activation into a tensor of token ids means nothing, so a write naming it is refused. embeddings is the embedding table's output — the same tensor as block_input at layer 0 on a model with nothing between them, and not on GPT-2, which adds position embeddings.

lm_head and logits #

nnterp carries the head's output and the model's output as two rows, because on Gemma-2 they are two tensors: final_logit_softcapping bounds what the model predicts from and leaves the head's own output alone. Mini has both, lm_head and logits, and a token-space metric binds to either — the document says which.

Measured on trl-internal-testing/tiny-Gemma2ForCausalLM, whose cap is 30.0: a swap of 100.0 into the head reads back 100.0 at lm_head and 29.924 at logits, which is tanh(100/30)·30. On a family with no cap the two reads are the same tensor, torch.equal, which is why every shipped document goes on saying lm_head.

Both engines reach it. nnterp's logits row is the model itself — an empty module path — with a Selection saying the tensor is the value's logits field. The nnterp engine reads it through model.internals; the hooks engine hooks the root and applies the same Selection, which is the one rule it uses everywhere, unwrapping a tuple's first element being another Selection (FirstIfTuple) rather than a special case.

from causalab_mini.engine.engines.hooks import HooksEngine, engine as hooks
from causalab_mini.plan.document import ModelSpec

e = HooksEngine.load(ModelSpec("hf-internal-testing/tiny-random-LlamaForCausalLM",
    "9fb191250dd56d0ba7ec9785a025ed29c03d5998", "fp32"), device_map="cpu")
for name in ("lm_head", "logits"):
    a = e.locate(name)
    print(name, "module=", repr(a.module), "io=", a.io)
    print("   resolves to", type(hooks.resolve(a, e._names)).__name__,
          "| select", hooks.selection(a, e._names))

lm_head module= 'lm_head' io= output
   resolves to Linear | select None
logits module= '' io= output
   resolves to LlamaForCausalLM | select Index('logits',)

The attention pattern, and the scores under it #

Two places name the attention matrix, and they are on either side of the softmax.

• attention_scores is the softmax's argument: QKᵀ·scale with the causal and padding masks already added, so a masked key is −inf here and 0 after. It is mini's own interior row, reached through the .source of the function the attention call dispatches to.
• attention_probs is nnterp's attention_probabilities: the dropout call's output rather than the softmax's, which is the pattern the values are actually mixed with on every family — after the cast, and after an attention sink has been dropped — and the identity in eval mode. nnterp carries six per-family overrides, a sink tag and a load-time validator behind that one name. A write here is make this head attend there.

Both exist only under eager attention: sdpa and flash never materialize the matrix, so there is no tensor to address. The refusal says so and says what to write:

component 'attention_probs' exists only under 'eager' attention, and this model runs 'sdpa': say "attn_implementation": "eager" in the document's model block

attn_implementation is a field of the plan-shaped document's model block. The protocol format's model block takes key, revision and dtype and rejects anything else, so the pattern is reachable only from the plan-shaped format.

Saying eager also turns on nnterp's enable_attention_probs, which puts nnterp's own shape, sum and causality check behind the accessor. A document says eager for exactly one reason, and that is the one.

documents/data/weekdays_aligned exists for that reason: it is the subset of the weekdays pairs that tokenize to the same length under this tokenizer, and it is what documents/v2/attention_pattern_patching.json runs on. See Attention pattern.

Inside a forward #

Four components — attention_query, attention_key, attention_scores, attention_z — name a tensor that never crosses a module boundary. The address is a module plus one operation inside its forward, which nnsight's .source exposes. All four hang off one call: the one that hands the query, the keys and the values to whichever attention implementation the checkpoint runs. Three are its arguments, attention_z is its return, and attention_scores is an operation inside the function it dispatches to.

The operation is named by the call site — the needle "attention_interface(" — and resolved against the loaded checkpoint, not tabulated, because the occurrence suffix is a property of the transformers version. nnsight gives assignments the same namespace as calls, so the attention forward has both an attention_interface_0 (the assignment attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(...)) and an attention_interface_1 (the call). A name match would hit both; a call-shaped needle hits one, and find_op refuses unless exactly one matches, listing every operation the forward does have.

from causalab_mini.engine import NNterpEngine
from causalab_mini.plan.document import ModelSpec

engine = NNterpEngine.load(ModelSpec(
    "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "9fb191250dd56d0ba7ec9785a025ed29c03d5998", "fp32"), dispatch=False)
print(engine.locate("block_mid", 0))
print(engine.locate("attention_z", 0))
print([n for n in engine.model.attentions[0].source.names if "attention_interface" in n])

Address(component='block_mid', layer=0, op=None, module='post_attention_layernorm', io='input', rank=(0, 33), inside=False)
Address(component='attention_z', layer=0, op='attention_interface_1', module=None, io=None, rank=(0, 22), inside=False)
['attention_interface_0', 'attention_interface_1']

attention_scores goes one level deeper: it is nn_functional_softmax_0 inside the function attention_interface_1 dispatches to. That inner name is resolved in the trace, because which function the call dispatches to is a run-time fact; unlike the outer one it is a name and not a needle.

Heads and units #

heads and units are one mechanism at two grains. Both compile into a Selection: positions along the sequence, plus groups and take over the feature axis. The rule is one line — view the features as (groups, -1) and keep take of the groups.

What is gathered comes back flat: (…, len(take) · per_group), however the model holds it. A tensor the model keeps as two axes — the query, which is (batch, head, seq, head_dim) — and one it keeps flat — attention_premix, which is o_proj's input — look the same to a featurizer, a mechanism and a metric. A scatter puts the named groups back and leaves the others as they were.

Which axis the sequence runs along is a fact about the address too, not about the tensor: 1 at a module boundary, 2 inside the attention, where the tensor is (batch, head, seq, head_dim). It is passed to the gather rather than guessed from the shape. What the positions along that axis are is Positions.

Per-head width is head_dim (or qk_head_dim) times the head count, and which head count depends on the address: attention_key is per key/value head and everything else per query head. On a model with grouped-query attention those differ. On hf-internal-testing/tiny-random-GptOssForCausalLM — 4 query heads, 2 key/value heads, head_dim 16 — causalab-mini model reports attention_query at width 64 and attention_key at width 32.

A head or unit index past the end is refused when the document is compiled, against the real count: site 'target': head 3 of a 2-head tensor.

Asking a model what it has #

causalab-mini vocab prints the table itself — the names, which nnterp accessor each is, which are per head, which need eager, which are read-only. It loads nothing:

$ uv run causalab-mini vocab
components:
  attention_query        interior  [heads]
  attention_key          interior  [heads]
  attention_scores       interior  [heads]  needs eager attention
  attention_probs        nnterp attention_probabilities  [heads]  needs eager attention
  attention_z            interior  [heads]
  input_ids              nnterp embeddings_input  (read-only)
  embeddings             nnterp embeddings_output
  …
  ln_final               nnterp ln_final_output
  lm_head                nnterp lm_head_output
  logits                 nnterp logits
…

causalab-mini model <key> asks the checkpoint. It loads the config, the module tree and the tokenizer and no weights, and tries every component: the width where there is one, the resolved .source operation for an interior, and the reason where the place is not there.

$ uv run causalab-mini model hf-internal-testing/tiny-random-LlamaForCausalLM \
      --revision 9fb191250dd56d0ba7ec9785a025ed29c03d5998
hf-internal-testing/tiny-random-LlamaForCausalLM via nnterp: 2 layers, padding_side=left
  attention_query    ok   width=16  op=attention_interface_1
  attention_key      ok   width=16  op=attention_interface_1
  attention_scores   --   component 'attention_scores' exists only under 'eager' attention, and this model runs 'sdpa': say "attn_implementation": "eager" in the document's model block
  attention_probs    --   component 'attention_probs' exists only under 'eager' attention, and this model runs 'sdpa': say "attn_implementation": "eager" in the document's model block
  attention_z        ok   width=16  op=attention_interface_1
  input_ids          ok   no width (no featurizer here)
  embeddings         ok   width=16
  …
  ln_final           ok   width=16
  lm_head            ok   width=32000
  logits             ok   width=32000

When a place is not there #

Every refusal about existence is nnterp's unavailable_on, decided from the address and the module tree and answered before any trace. locate passes the reason through unchanged. Three kinds happen.

A mixture of experts has no single activation. A token goes through top-k of N experts, so there is no one activation function's output and no one down projection's input. The block's boundaries are unaffected:

$ uv run causalab-mini model hf-internal-testing/tiny-random-GptOssForCausalLM
hf-internal-testing/tiny-random-GptOssForCausalLM via nnterp: 2 layers, padding_side=left
  attention_query    ok   width=64  op=attention_interface_1
  attention_key      ok   width=32  op=attention_interface_1
  attention_scores   ok   no width (no featurizer here)  op=attention_interface_1
  attention_probs    --   component 'attention_probs': Attention probabilities are disabled for this model. Set enable_attention_probs=True when loading the model to enable them.
  …
  mlp_input          ok   width=64
  mlp_activation     --   component 'mlp_activation': mlps_activation does not exist on a mixture-of-experts layer, where a token goes through top-k of N experts and has no single activation. mlps_input and mlps_output are the block's boundaries and work as usual.
  mlp_neuron_output  --   component 'mlp_neuron_output': mlps_neurons does not exist on a mixture-of-experts layer, where a token goes through top-k of N experts and has no single activation. mlps_input and mlps_output are the block's boundaries and work as usual.
  mlp_output         ok   width=64
  …

That reason is a function of the block, not a fact about the model, because a model may be dense for its first blocks and a mixture of experts after — DeepSeek's first_k_dense_replace is exactly that. So mlp_activation can resolve at layer 0 and be refused at layer 12 of the same model, and the answer is per layer.

The attention_probs row above is refused for a different reason. This family runs eager by default, which is why attention_scores — mini's own interior — resolves; nnterp keeps its pattern row disabled until enable_attention_probs is asked for, and a document asks for it by saying "attn_implementation": "eager". Loaded that way the row resolves here too.

A hybrid's linear-attention blocks have no attention. On Qwen3-Next and Qwen3.5 hybrids some blocks mix with Gated DeltaNet rather than softmax attention. nnterp splits the stack into attention_layers and linear_attention_layers, and every attention accessor and the pattern raise on the latter: layer i is a linear-attention layer (Mixer); attentions/attention_probabilities are only defined on softmax-attention layers.

A block structure that has no such place. On a parallel block, attention and the MLP both read the block input, so there is no residual stream between them: block_mid and mlp_input_norm are refused. On a post-norm block each sublayer reads the residual stream itself and the norm comes after, so attention_input_norm and mlp_input_norm are refused. nnterp keys those rows by block structure and holds the reason where the structure has no such place, so they are refused with a sentence rather than resolving to a neighbouring tensor.

A name only nnterp knows #

RenameConfig(addresses={...}) is nnterp's extension point: a row of the address table, by accessor name, either replacing where an existing accessor reads or adding one of your own. An accessor added that way appears in model.internals, and a document may name it — mini's table is a floor, and a component it has never heard of is accepted as long as this model's internals have an accessor of that name.

from nnterp.rename_utils import Address as NnterpAddress, RenameConfig
from causalab_mini.address import AddressError
from causalab_mini.engine import NNterpEngine
from causalab_mini.plan.document import ModelSpec

engine = NNterpEngine.load(
    ModelSpec("hf-internal-testing/tiny-random-LlamaForCausalLM",
              "9fb191250dd56d0ba7ec9785a025ed29c03d5998", "fp32"),
    dispatch=False,
    rename_config=RenameConfig(addresses={
        "mlp_gate_output": NnterpAddress("mlp.gate_proj", order=42),
    }),
)
located = engine.locate("mlp_gate_output", 0)
print(located)
try:
    engine.width(located)
except AddressError as refusal:
    print("--", refusal)

Address(component='mlp_gate_output', layer=0, op=None, module='mlp.gate_proj', io='output', rank=(0, 42), inside=False)
-- the width of 'mlp_gate_output' is not a fact about the model: it depends on the batch (the attention pattern's key axis) or is not a tensor's

Mini claims nothing else about such a place. It is a module boundary, read at the accessor's own side, one per layer or not as nnterp says, not per head, and it has no width — so a featurizer, which needs one, is refused there rather than sized wrongly. The order in the row is what puts it in forward order beside everything else; give it one that reflects where the place actually is.

What is refused #

Every one of these happens before a forward runs — most at validation, the rest when the document is compiled against the model.

A read of four heads, and a read of the whole model #

documents/v2/head_patching.json is the interchange of patching.json done one attention head at a time. Two sites: target is attention_premix at layer 0 with a sweep over which head, and head is lm_head, which names no layer because there is only one of it.

{
  "header": {
    "description": "Head-by-head activation patching: the interchange of patching.json, but of ONE attention head's result at a time, swept over the four heads of layer 0. The site is `attention_premix` \u2014 the output projection's input, which is every head's result side by side \u2014 so this runs on the hooks engine too; `attention_z` is the same tensor held as (heads, head_dim) and gives the same numbers bit for bit on nnterp."
  },
  "model": {
    "key": "hf-internal-testing/tiny-random-LlamaForCausalLM",
    "revision": "9fb191250dd56d0ba7ec9785a025ed29c03d5998",
    "dtype": "fp32"
  },
  "roles": {
    "base": {
      "field": "input"
    },
    "counterfactual": {
      "field": "counterfactual_inputs[0]"
    }
  },
  "sites": {
    "target": {
      "component": "attention_premix",
      "layers": [
        0
      ],
      "heads": {
        "sweep": [
          [
            0
          ],
          [
            1
          ],
          [
            2
          ],
          [
            3
          ]
        ]
      }
    },
    "head": {
      "component": "lm_head"
    }
  },
  "interventions": {
    "patching": {
      "reads": {
        "v_cf": {
          "site": "target",
          "pos": -1,
          "model": "original",
          "input": "counterfactual"
        },
        "logits": {
          "site": "head",
          "pos": -1,
          "model": "patched",
          "input": "base"
        }
      },
      "writes": {
        "patch": {
          "site": "target",
          "pos": -1,
          "mechanism": "swap",
          "operand": "v_cf"
        }
      },
      "models": {
        "patched": {
          "input": "base",
          "writes": [
            "patch"
          ]
        }
      },
      "metrics": {
        "iia": {
          "kind": "match",
          "of": "logits",
          "expected": "cf_answer",
          "token_form": "space_prefixed"
        },
        "logit_diff": {
          "kind": "logit_diff",
          "of": "logits",
          "a": "cf_answer",
          "b": "base_answer",
          "token_form": "space_prefixed"
        }
      }
    }
  },
  "steps": {
    "score": {
      "kind": "observe",
      "interventions": "patching",
      "rows": {
        "base": "weekdays/train",
        "counterfactual": "weekdays/train"
      },
      "saves": [
        {
          "value": "iia",
          "file_path": "iia.json"
        },
        {
          "value": "logit_diff",
          "file_path": "logit_diff.json"
        }
      ]
    }
  }
}

The document's own note says why the site is attention_premix and not attention_z: the two are the same tensor held differently, attention_premix is a module boundary and so runs on the hooks engine too, and on nnterp the two give the same numbers bit for bit.

$ uv run causalab-mini run documents/v2/head_patching.json --device-map cpu --out out/heads
out/heads/document.json
out/heads/run.json
out/heads/heads=0/document.json
out/heads/heads=0/iia.json
out/heads/heads=0/logit_diff.json
…
out/heads/heads=3/document.json
out/heads/heads=3/iia.json
out/heads/heads=3/logit_diff.json

Four points, one per head, each a directory named by the swept axis. One row of out/heads/heads=1/iia.json:

{
 "example_id": "0",
 "metric": "iia",
 "value": 0.0,
 "eligible": true,
 "positions": null,
 "reason": "",
 "unit": "fraction",
 "estimand_version": "match/v1",
 "produced_by": "8334fa82ec1ce2f26e9a178eebb003c9898cdb6ee445448f25c5be4d2297865a"
}