Standardization

causalab_mini/address.py is the only file in the project that knows anything about a model's internals, and most of what it knows it asks nnterp for. The split is one sentence: which module or operation is "the mid-stream of layer 3" on this architecture is nnterp's question, and whether a document may name that place, and what happens there is mini's.

The mechanism on the nnterp side is one table of Address rows. A row is pure data: a module path in standardized names, which side of it, any .source operations to descend through, where the tensor sits in the value there, where it sits in the forward pass, and why a family has none. A family that differs differs by a row, not by a branch of code.

At a glance #

One name, two families that spell the block and its children nothing alike — tiny Llama's first block is model.model.layers.0, tiny GPT-2's is model.transformer.h.0:

from nnterp import StandardizedTransformer

for key in ["hf-internal-testing/tiny-random-LlamaForCausalLM",
            "hf-internal-testing/tiny-random-gpt2"]:
    model = StandardizedTransformer(key, device_map="cpu")
    row = model.internals["layers_mid"].address
    print(f"{key:48} {model.block_structure:10} {row.module!r} {row.io.value}")

prints:

hf-internal-testing/tiny-random-LlamaForCausalLM pre_norm   'post_attention_layernorm' input
hf-internal-testing/tiny-random-gpt2             pre_norm   'ln_2' input

Mini's block_mid is that row's name and nothing else, so the address mini compiles is equal on both checkpoints — same component, same layer, same resolved operation. That equality is asserted for every one of mini's components in tests/test_components.py::test_one_address_serves_both_families. Mini has no family axis and no family column; the axis is absorbed by the dependency.

Where the line is #

What nnterp answers, for every place it addresses:

• which child module this family spells the place with, and which side of it — or which operation inside whose forward;
• whether the place exists at all on this architecture, and on this layer, with the reason when it does not;
• where the tensor sits inside the value there;
• where the place sits in the forward pass;
• every width and head count — hidden_size, intermediate_size, head_dim, qk_head_dim, num_kv_heads.

What mini keeps beside the accessor's name is policy: which axis of the tensor the sequence runs along, whether it is per head and per which heads (q or kv, which differ under grouped-query attention), whether its last axis is the keys of a padded batch rather than a feature width, which model attribute publishes its width, whether a write there is refused, and the four interiors nnterp does not address.

What an Address row is #

Every row of the table becomes an attribute of the model — model.layers_output[i], model.lm_head_output() — and a member of model.internals, which is the rows by name in forward order. Adding a place is adding a row and nothing else. A row whose name the model already uses is refused at load with a message saying to rename it, rather than replacing the attribute.

model.logits and model.token_embeddings are the two exceptions, and they are spellings rather than exceptions to the table: each stays the tensor property it is and reads its row (logits and embeddings_output), so there is still one answer per place. Locations has the component table and why logits and lm_head are not the same tensor on every family.

A family that differs is one row #

FAMILY_ADDRESSES is a list of (model class or predicate, rows); the defaults are applied, then every entry that matches, then the user's RenameConfig. Later entries win. Each row is replace() on the default one, so an override states only what differs.

A place on a child nnterp does not rename #

nnterp standardizes the block, the mixer and the MLP. It does not rename their children, and families spell them differently: input_layernorm or ln_1, o_proj or c_proj or dense or out_proj, act_fn or act or gelu_impl. Which spelling a model uses is a question the checkpoint can answer, so STRUCTURAL_ADDRESSES holds the candidates and structural_addresses reads the module tree: the row's {name} becomes the one candidate that exists.

If none exists, or several do, the place is unavailable with a reason naming the candidates it looked for, the children the layers actually have, and RenameConfig(addresses={…}) as the way to say which. Guessing by shape is exactly the silent failure this removes.

Six rows are resolved this way. On the two fixtures:

layers_mid and mlps_norm_output land on the same module and differ by side: the mid-stream is what the pre-MLP norm reads, and the norm's output is what the MLP reads.

Block structure, and why the tree alone is not enough #

model.block_structure is how a block combines its two sublayers with the residual stream, read from the model's class and config: pre_norm, sandwich_norm (Gemma-2/3), post_norm (OLMo-2), parallel (GPT-NeoX, GPT-J, Falcon, Phi, StableLM-2) or residual_inside (BLOOM, MPT, DBRX). Each structural row's candidates are keyed by it, and a structure with no such place holds the reason instead of the candidates.

Keying by structure rather than by name is the whole point, because the same child name means different things on different blocks:

• A sandwich-norm block spells the pre-MLP norm pre_feedforward_layernorm and also has a post_attention_layernorm — which there is the attention's post-norm. A candidate-existence rule would pick it and read a real tensor of the right shape from the wrong place.
• A post-norm block has the children and not the places: attention reads the residual stream itself, and the norm comes after the sublayer.
• A parallel block has no mid-stream at all, because attention and the MLP both read the block input — and on GPT-NeoX a second norm exists anyway, normalizing that input.

What holds the names to their meanings is that each place is defined by what it is of the block, as an identity over the tensors:

layers_mid       == layers_input + attentions_output
layers_output    == layers_mid   + mlps_output
mlps_norm_output == mlps_input

and, on a parallel block, layers_output == layers_input + attentions_output + mlps_output. A row that resolved to a plausible neighbour satisfies none of them.

Availability, before any trace #

model.internals.status() answers, for every row, None if it is available and otherwise the reason it is not — decided from the address and the module tree, without running the model. On the tiny Llama:

from nnterp import StandardizedTransformer

model = StandardizedTransformer("hf-internal-testing/tiny-random-LlamaForCausalLM",
                                device_map="cpu")
for name, reason in model.internals.status().items():
    print(f"{name:26} {reason or ''}")

embeddings_input
embeddings_output
layers_input
attentions_norm_output
attentions
attentions_input
attention_probabilities    Attention probabilities are disabled for this model. Set enable_attention_probs=True when loading the model to enable them.
attentions_premix
attentions_output
layers_mid
mlps_norm_output
mlps
mlps_input
mlps_activation
mlps_neurons
mlps_output
layers_output
ln_final_output
lm_head_output
logits

Twenty rows, in forward order, one reason. The tiny GPT-2 prints the same twenty names with the same one reason — the difference between the two families is entirely in what those names resolved to, which is the table above.

A family that lacks a place says so by name. The tiny GPT-NeoX, whose block is parallel:

…
attentions_output
layers_mid                 layers_mid does not exist on this model: its block is 'parallel'. Attention and the MLP both read the block input, so there is no residual stream between them.
mlps_norm_output           mlps_norm_output does not exist on this model: its block is 'parallel'. The MLP reads the block input, so there is no mid-stream for a norm to read (on GPT-NeoX the second norm exists, and normalizes the block input).
mlps
…

unavailable may be a function of the layer's module rather than a string, and status(layer=i) answers for one layer — which is what a model whose layers are not all alike needs. On a mixture of experts a token goes through top-k of N experts and has no single activation, so the two MLP interiors say so:

…
mlps_input
mlps_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.
mlps_neurons               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.
mlps_output
…

Every block of that checkpoint is a mixture, so its model-wide answer and each layer's are the same. The mixed case is yujiepan/deepseek-v3-tiny-random, whose layer 0 is dense and layer 1 a mixture: status(layer=0)["mlps_activation"] is None there and status(layer=1) gives the reason above. That checkpoint is not in this machine's cache, so the split is quoted from nnterp/tests/test_config.yaml and the assertion in test_block_invariants.py rather than run here.

Selection: where the tensor is inside the value #

A place does not always hold the tensor directly. A module may return its output beside a cache; the model's own output is an object with a logits field. Address.select is where the tensor is inside that value, and it is a pair — get to read and put to write — because a write has to hand the place back its whole value with one part replaced, rebuilding tuples, named tuples, lists and mappings on the way out.

Four rows of twenty carry one. Three are FirstIfTuple() — attentions_output, mlps_output, layers_output — and one is an Index: logits, which reads the logits field of the model's own output.

Forward order #

nnsight will not reach back to a value the model has already passed: reading layer 8 and then layer 2 in one trace raises, and a write has to go in above the read that observes it. So a place's rank in the forward pass is part of its row. internals.rank(name, layer) is (layer, order) for a per-layer place; a whole-model place gets (-1, order) where its order is negative and (num_layers, order) where it is not, which puts the embeddings below every layer and the final norm and head above.

Mini stamps that answer into the address rather than keeping a second numbering beside it — nnterp moves a place's order per family, and two numberings are how they drift apart with nothing saying so. The four interiors, which nnterp does not address, are mini's own numbers interleaved into nnterp's: 11, 12, 15 and 22, in the gaps nnterp leaves between its 10, 20, 25 and 30 inside the attention. Every component of mini's table on the tiny Llama, sorted by its rank:

rank        component              kind       nnterp row
(-1, -2)    input_ids              boundary   embeddings_input
(-1, -1)    embeddings             boundary   embeddings_output
(0, 0)      block_input            boundary   layers_input
(0, 5)      attention_input_norm   boundary   attentions_norm_output
(0, 11)     attention_query        interior   -
(0, 12)     attention_key          interior   -
(0, 15)     attention_scores       interior   -
(0, 20)     attention_probs        interior   attention_probabilities
(0, 22)     attention_z            interior   -
(0, 25)     attention_premix       boundary   attentions_premix
(0, 30)     attention_output       boundary   attentions_output
(0, 33)     block_mid              boundary   layers_mid
(0, 36)     mlp_input_norm         boundary   mlps_norm_output
(0, 40)     mlp_input              boundary   mlps_input
(0, 44)     mlp_activation         boundary   mlps_activation
(0, 47)     mlp_neuron_output      boundary   mlps_neurons
(0, 50)     mlp_output             boundary   mlps_output
(0, 60)     block_output           boundary   layers_output
(2, 1000)   ln_final               boundary   ln_final_output
(2, 1001)   lm_head                boundary   lm_head_output
(2, 1002)   logits                 boundary   logits

Twenty-one components. Seventeen are the name of an nnterp row, and four are interiors nnterp does not address. attention_probs is a row of nnterp's that is itself an operation — its address names two .source operations — so mini marks that address inside, which is what an engine that only knows module boundaries has to be told. The three nnterp rows mini's table does not name (attentions, attentions_input, mlps) are still addressable; see below.

How mini consumes it #

locate(model, component, layer) is the whole of it. It asks model.internals[accessor] for the row, asks unavailable_on(layer) whether the place exists here, asks internals.rank where it sits, and writes the resolved child, the side and the rank into an Address that stays pure data — so it pickles, sorts, prints, and travels in a plan that an engine without nnterp's accessors can still walk.

In the trace, a boundary is one line each way:

model.internals[address.accessor][address.layer]            # read
model.internals[address.accessor][address.layer] = tensor   # write

Two things are cross-checked rather than trusted. locate compares mini's per_layer against nnterp's and refuses if they disagree, naming the table as mini's to fix; and a place the family or the layer lacks is refused at compile time with nnterp's own reason, rather than inside someone else's process. causalab-mini model is that, for every component at once:

$ uv run causalab-mini model hf-internal-testing/tiny-random-GPTNeoXForCausalLM
hf-internal-testing/tiny-random-GPTNeoXForCausalLM via nnterp: 5 layers, padding_side=left
  attention_query    ok   width=32  op=attention_interface_1
  attention_key      ok   width=32  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_output   ok   width=32
  block_mid          --   component 'block_mid': layers_mid does not exist on this model: its block is 'parallel'. Attention and the MLP both read the block input, so there is no residual stream between them.
  mlp_input_norm     --   component 'mlp_input_norm': mlps_norm_output does not exist on this model: its block is 'parallel'. The MLP reads the block input, so there is no mid-stream for a norm to read (on GPT-NeoX the second norm exists, and normalizes the block input).
  mlp_input          ok   width=32
  mlp_activation     ok   width=37
  …
  logits             ok   width=1024

What is still mini's #

Four interiors: attention_query, attention_key, attention_scores, attention_z. The tensor never crosses a module boundary, so the address is a module plus one operation inside its forward, reached through nnsight's .source. The query and the key are positional arguments 1 and 2 of the attention interface call; attention_z is that call's return, element 0. nnterp's design document lists these among the things its internals layer does not address, along with per-head indexing and a tensor-layout descriptor — so the per-head slicing is mini's too, and a per-head tensor is handed on flat however the model holds it.

Mini names such an operation by its call site, not by its name, and find_op resolves the name against the loaded checkpoint. nnsight names an operation {callable}_{occurrence} and gives assignments the same namespace as calls, so this forward has both attention_interface_0 (the assignment that looks the implementation up) and attention_interface_1 (the call). A name match on attention_interface hits both and is refused with the forward's whole operation list; the needle "attention_interface(" is call-shaped and hits one.

attention_scores is a step further in: the softmax's argument inside the function the interface call dispatches to. nnsight can only open that nested source inside a trace — which function the call dispatches to is a run-time fact — so the table holds the inner operation's name and the engine resolves it where the run runs. It exists only under eager attention, and that is checkable on the client, so locate refuses at compile time and says what to write: "attn_implementation": "eager" in the document's model block.

Positions are mini's too, for the same reason the components are nnterp's: Llama's tokenizer prepends a BOS and GPT-2's prepends nothing, so without a content frame {"index": 0} is <s> on one family and the first word on the other — one document addressing two different things, silently. See Positions.

Adding an architecture, or a place #

An architecture nnterp has never seen is a RenameConfig: the names of its containers (attn_name, mlp_name, ln_final_name, lm_head_name, model_name, layers_name) and, where a place is somewhere else, an Address row for it. addresses={…} is the general form of the three narrower fields beside it: it replaces where an existing accessor reads, or adds one of your own.

from nnterp import StandardizedTransformer
from nnterp.rename_utils import Address, IOType, RenameConfig

config = RenameConfig(addresses={
    "attentions_gate": Address("self_attn.q_proj", IOType.OUTPUT, order=15),
})
model = StandardizedTransformer("hf-internal-testing/tiny-random-LlamaForCausalLM",
                                device_map="cpu", rename_config=config)
print("attentions_gate" in model.internals)
print(repr(model.internals["attentions_gate"]))

True
<attentions_gate: self_attn.q_proj.output, per layer>

Mini's table is a floor, not a fence. A component name mini has never heard of but model.internals has is addressable, because a document should be able to name what a user added. The same row, through mini, resolves at layer 1 to module='self_attn.q_proj', io='output', rank=(1, 15). Mini claims nothing else about such a place: no width, so a featurizer there is refused rather than sized wrongly, it is not per head, and it is a module boundary.

When transformers moves #

A .source row names operations the way nnsight names them, which is a fact about the installed transformers version and not about the architecture. An upgrade that moves a call, renames a local, or routes a family through the shared attention interface invalidates a row. Two suites stand against it:

A row that drifted onto the scores, the values or the post-mix output passes a shape check and fails these. And when a row does break, the error names the row, the module, the transformers version, the model class and every operation that does exist at that level of the source:

attention_probabilities: the forward of layers.0.self_attn.attention_interface_1 has no operation 'module_attn_dropout_0' with transformers 5.17.0 on GPT2LMHeadModel. Its operations are:
                             * def eager_attention_forward(module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs):
                             …
 nn_functional_softmax_0 ->  8     attn_weights = nn.functional.softmax(attn_weights, dim=-1)
 attn_weights_2          ->  +     ...
                             9
                            10     # Downcast (if necessary) back to V's dtype (if in mixed-precision) -- No-Op otherwise
 attn_weights_type_0     -> 11     attn_weights = attn_weights.type(value.dtype)
 attn_weights_3          ->  +     ...
 nn_functional_dropout_0 -> 12     attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
 attn_weights_4          ->  +     ...
                            13
                            14     return attn_output, attn_weights

The fix is a one-row edit, and the name this one needs — nn_functional_dropout_0 — is in the list the message printed.

Remote #

A request carries nothing of nnterp. The package is not registered for by-value serialization, so an NDIF server runs its own installed copy and must have it at the same version as the client. The same flag keeps the checkpoint off the client: remote=True sets allow_dispatch=False, every load-time check runs with scan() on the meta model, and no request is sent during construction.

Mini is on the same side of that line: nothing of it ships by value either, so the server needs causalab_mini installed at the client's version as well as nnterp, and a server without it fails with a ModuleNotFoundError rather than running a different copy. See Remote & NDIF.