# Getting Started nnterp gives every transformer family the same module names and the same standard values, on top of [nnsight](https://nnsight.net). This page installs it, runs a first trace, and sets that trace beside its nnsight form and its nnterp 1.x form. For the model you are about to load, the [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/) has a page per supported family: the block drawn with nnterp's names and every value on it. ## Installation nnterp is a PyPI package for Python 3.10 or later. It depends on nnsight 0.8 and transformers 5.17 or later, which pip installs with it. ```bash pip install nnterp ``` ### From source ```bash pip install git+https://github.com/ndif-team/nnterp.git ``` For a checkout that you intend to edit, run `pip install -e .` from the clone. ### Extras ```bash pip install "nnterp[models]" # sentencepiece, tiktoken, protobuf, pillow: tokenizers and processors some checkpoints need pip install "nnterp[vllm]" # vLLM, for StandardizedVLLM pip install "nnterp[test]" # pytest and pytest-xdist, for the test suite ``` ### For agents If an agent writes your nnterp code, give it the nnterp skills from [ndif-team/skills](https://github.com/ndif-team/skills): three skills, for the values, for the recipes that run across families, and for extending nnterp. In Claude Code: ``` /plugin marketplace add https://github.com/ndif-team/skills.git /plugin install nnterp@ndif-team ``` In Codex, `$skill-installer install https://github.com/ndif-team/skills`. Every page of this site is also served as Markdown, listed in [llms.txt](../llms.txt), with the whole site in [llms-full.txt](https://ndif-team.github.io/nnterp/llms-full.txt). ## Quick start A `StandardizedTransformer` is an nnsight `TransformersModel` with one set of module names on every family, and with values on its blocks that mean the same thing everywhere. Load GPT-2 and read block 5's values and the logits: ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): attn = model.layers[5].self_attn.attention_output.save() # what attention adds to the residual stream mlp = model.layers[5].mlp.mlp_output.save() # what the MLP adds resid = model.layers[5].layer_output.save() # the residual stream leaving block 5 logits = model.logits.save() print(attn.shape, mlp.shape, resid.shape, logits.shape) # torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 50257]) ``` The constructor reads the checkpoint's `model_type`, picks the matching family, and passes every other keyword to `TransformersModel`. ### The names `model.embed_tokens`, `model.layers[i]`, `model.layers[i].self_attn`, `model.layers[i].mlp`, `model.norm` and `model.lm_head` exist on every family. Because they are nnsight aliases, the native names keep working and both reach the same envoy: ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") print(model.layers[5] is model.transformer.h[5], model.layers[5].self_attn is model.transformer.h[5].attn) # True True ``` ### The values A name says which module; a value says which tensor. `layer_output` is a tensor whether the family's block returns a tensor or a tuple, and `attention_output` and `mlp_output` are what each sublayer adds to the residual stream, wherever the family adds it, so one identity holds on every family: ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): block = model.layers[5] x = block.input.save() a = block.self_attn.attention_output.save() m = block.mlp.mlp_output.save() out = block.layer_output.save() print(torch.allclose(x + a + m, out, atol=1e-4)) # True ``` A value is read, edited in place or assigned the way an nnsight `.output` is. ### Availability Not every checkpoint has every value. `model.support()` says, from the config and the module tree, what this one has and, for anything missing, why. A family's [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/) page shows the same answer for each of its public checkpoints. Reading a missing value inside a trace raises `nnterp.Unavailable` with the same reason, before the model runs: ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") print(model.support(layer=5)["self_attn.attention_probabilities"]) # read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' ``` Anything inside attention needs the eager implementation: ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") with model.trace("The Eiffel Tower is in"): pattern = model.layers[3].self_attn.attention_probabilities.save() # [batch, num_heads, query, key] print(pattern.shape) # torch.Size([1, 12, 7, 7]) ``` ### Methods The operations that an experiment repeats are methods on the model, written against the values, so they run on every family: ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") torch.manual_seed(0) vector = torch.randn(model.hidden_size) with model.trace("The Eiffel Tower is in"): lens = model.project_on_vocab(model.layers[6].layer_output).save() # logit lens at block 6 model.steer(8, vector, factor=3, token_positions=-1) # add to the stream leaving block 8 model.skip_layers(10, 11) # blocks 10 and 11 do not run logits = model.logits.save() print(model.tokenizer.decode(lens[0, -1].argmax()), logits.shape) # front torch.Size([1, 7, 50257]) ``` ## Coming from nnsight A `StandardizedTransformer` is a `TransformersModel`, so the trace is nnsight's. What changes is the block's body: the family's module paths become one set of names, and `.output` with its per-family tuple indexing becomes values that are tensors with one meaning. The nnsight trace, on GPT-2's own names: ```python from nnsight import TransformersModel model = TransformersModel("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): attn = model.transformer.h[5].attn.output[0].save() # (attn_output, attn_weights): index the tensor mlp = model.transformer.h[5].mlp.output.save() resid = model.transformer.h[5].output.save() # a tensor on GPT-2, a tuple on BLOOM or GPT-J logits = model.output.logits.save() print(attn.shape, mlp.shape, resid.shape, logits.shape) # torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 50257]) ``` The same trace with nnterp's names and values is the first block of the [quick start](#quick-start): the four tensors are the ones that the nnsight trace saved (the next block checks one), and that body runs unchanged on `EleutherAI/pythia-70m-deduped`, where the attention is `gpt_neox.layers[5].attention`, and on `bigscience/bloom-560m`, where the block returns a tuple and `layer_output` is still the tensor. ### What does not change `trace`, `generate`, `.save()`, `.input`, `.output`, invokes, `tracer.iter`, `.source`, `model.edit` and `remote=True` work as [nnsight's documentation](https://nnsight.net/features/1_getting/) describes them. The native names keep working on the same model, and a trace written for a `TransformersModel` runs on a `StandardizedTransformer` as it is: ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): native = model.transformer.h[5].attn.output[0].save() standard = model.layers[5].self_attn.attention_output.save() print(torch.equal(native, standard), model.layers[5].self_attn is model.transformer.h[5].attn, model.layers[5].self_attn.path) # True True model.transformer.h.5.attn ``` ## Coming from nnterp 1.x nnterp 1.x reached a block's tensors through accessors on the model, `model.layers_output[i]` and `model.attentions_output[i]`; nnterp 2 puts them on the block's envoy, `model.layers[i].layer_output` and `model.layers[i].self_attn.attention_output`, and the constructor validates nothing. | nnterp 1.x | nnterp 2 | status | | --- | --- | --- | | `model.embed_tokens`, `model.layers[i]`, `model.lm_head` | same | same | | `model.ln_final` | `model.norm` | renamed | | `model.layers_input[i]` | `model.layers[i].input` | renamed | | `model.layers_output[i]` | `model.layers[i].layer_output` | renamed | | `model.attentions[i]`, `model.attentions_input[i]` | `model.layers[i].self_attn`, `model.layers[i].self_attn.input` | renamed | | `model.attentions_output[i]` | `model.layers[i].self_attn.attention_output` | changed: the contribution, not the module's output; the same tensor on GPT-2 and Llama, different only where the family norms the module's output (Gemma-2) or adds the residual inside it (BLOOM) | | `model.mlps[i]`, `model.mlps_input[i]` | `model.layers[i].mlp`, `model.layers[i].mlp.input` | renamed | | `model.mlps_output[i]` | `model.layers[i].mlp.mlp_output` | changed: the contribution, not the module's output | | `model.attention_probabilities[i]` | `model.layers[i].self_attn.attention_probabilities` | renamed | | `model.logits`, `token_embeddings`, `next_token_probs`, `input_ids`, `input_size`, `attention_mask` | same; `logits`, `input_ids` and `attention_mask` are also assignable | same | | `model.num_layers`, `num_heads`, `hidden_size`, `vocab_size` | same, read-only | same | | `model.attn_probs_available` | `model.support()["self_attn.attention_probabilities"] is None` | changed | | `StandardizedTransformer(repo, enable_attention_probs=True)` | `StandardizedTransformer(repo, attn_implementation="eager")` | changed | | `check_renaming=`, `allow_dispatch=`, `check_attn_probs_with_trace=`, `allow_multimodal=` | nothing runs at construction | gone | | `rename_config=RenameConfig(...)` | `rename=`, `envoys=`, `family=`, `nnterp.families.register` | changed | | `StandardizedTransformer(repo, remote=True)` | `model.trace(prompt, remote=True)` | changed | | `device_map="auto"` unless given | nnsight's default; pass `device_map="auto"` | changed | | `trust_remote_code=`, `tokenizer_kwargs=`, `rename=` | same | same | | `model.skip_layer(i)` | `model.skip_layers(i, i)` | gone | | `model.skip_layers(start, end, skip_with, layer_returns_tuple)`, `detect_layer_output_type()` | `model.skip_layers(start, end, skip_with)` | changed | | `model.steer(layers, steering_vector, factor, positions, token_positions, batch_index)` | `model.steer(layers, vector, factor, token_positions, batch_index)` | changed | | `model.project_on_vocab(h)` | same name; applies the softcap on Gemma-2 | changed | | `model.probs_to_dict(tokens, probs)` | `model.probs_to_dict(probs, k=5)` | changed | | `model.get_topk_closest_tokens(h, k)`, a dict for a 1-D `h` | always a list of dicts | changed | | `nnterp.nnsight_utils.get_layer_output(model, i)` and the other `get_*`, `ModuleAccessor` | none: the names and values above | gone | | `get_token_activations`, `collect_token_activations_batched`, `collect_last_token_activations_session`, `compute_next_token_probs` | same | same | | `nnterp.prompt_utils`: `Prompt`, `get_first_tokens`, `run_prompts` | same; `run_prompts` shows a progress bar only when given `tqdm=` | same | | `nnterp.interventions`: `logit_lens`, `patchscope_lens`, `patchscope_generate`, `repeat_prompt`, `TargetPrompt` | none; the [Logit lens](../tutorials/logit-lens.md) and [Activation patching](../tutorials/activation-patching.md) tutorials | gone | | `nnterp.display.plot_topk_tokens`, `prompts_to_df` | none | gone | | `RenamingError`, `get_rename_dict` | `nnterp.Unavailable`, `nnterp.UnsupportedFamily` | changed | ### `layers_output` and the other accessors 1.x code read `model.layers_output[5]` and wrote `model.layers_output[5] = v`; nnterp 2 code reads and writes `model.layers[5].layer_output`. The module accessors `attentions[i]` and `mlps[i]` are the envoys `layers[i].self_attn` and `layers[i].mlp`, and `ln_final` is `norm`: ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): x = model.layers[5].input.save() # layers_input[5] attn = model.layers[5].self_attn.attention_output.save() # attentions_output[5] mlp = model.layers[5].mlp.mlp_output.save() # mlps_output[5] model.layers[5].layer_output[:, -1] += 0.5 * x[:, -1] # layers_output[5] edited in place normed = model.norm.output.save() # ln_final_output logits = model.logits.save() print(x.shape, attn.shape, mlp.shape, normed.shape, logits.shape) # torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 50257]) ``` `attentions_output[i]` and `mlps_output[i]` were the module's output, first element of a tuple; `attention_output` and `mlp_output` are what the sublayer adds to the residual stream, and `layers[i].input + attention_output + mlp_output == layer_output` on every family. On GPT-2 and Llama the two are the same tensor. On BLOOM the attention module returns the stream with the residual already added, and on Gemma-2 the block adds a post-norm's output, so 1.x code that read the module's output there now reads a different tensor: ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("bigscience/bloom-560m") with model.trace("The Eiffel Tower is in"): x = model.layers[2].input.save() attn = model.layers[2].self_attn.attention_output.save() # the contribution raw = model.layers[2].self_attn.output[0].save() # what 1.x attentions_output[2] gave print(torch.allclose(raw, x + attn)) # True ``` ### Read order nnterp 2 reads values in the model's forward order within a trace and raises `OutOfOrderError` otherwise. 1.x code that read block 5 before block 3 in one trace has to be reordered ([Attention: read order](../guide/values/attention.md#read-order)). ### The attention pattern 1.x code passed `enable_attention_probs=True`, which forced the eager attention and checked the pattern at load, and asked `model.attn_probs_available`. nnterp 2 forces nothing: pass `attn_implementation="eager"` yourself and ask `support()`, which answers from the config, before any trace. Reading the pattern on a model that runs `sdpa` raises `nnterp.Unavailable` at that line. !!! warning "The 1.x flags reach transformers, which rejects them" `enable_attention_probs=True`, `check_renaming=`, `allow_dispatch=` and the other 1.x keywords pass through `TransformersModel` to the checkpoint's constructor, which raises `TypeError: GPT2LMHeadModel.__init__() got an unexpected keyword argument 'enable_attention_probs'`, inside a `ValueError` from the loader. With `dispatch=True` that happens at load. Without it the load reads only the config, `support()` reports `sdpa`, and the error comes at the first trace. Remove the keywords. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") print(model.support()["self_attn.attention_probabilities"]) # None with model.trace("The Eiffel Tower is in"): pattern = model.layers[5].self_attn.attention_probabilities.save() # attention_probabilities[5] print(pattern.shape, pattern[0, 0].sum(-1)) # torch.Size([1, 12, 7, 7]) tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000], # grad_fn=) ``` `support()` returns `None` for a value that every block has, else the reason, per block. Ask it instead of `hasattr`: `hasattr(envoy, "attention_probabilities")` raises `Unavailable` when the value is missing rather than returning `False`. ### The constructor 1.x ran a renaming check and a dummy forward at load and raised `RenamingError` when a module was not where it expected. nnterp 2 reads the config's `model_type`, picks the family and builds the tree. A shipped family is used as it is. A `model_type` with no family loads under a best-effort default with a warning, and `UnsupportedFamily` is raised only when that default cannot name the modules either. A wrong assumption fails at the first trace instead, as `Unavailable` or nnsight's `SourceNotAvailable`, so run `model.support()` and one trace of the values that you need before a long experiment. `remote=True` belongs on `trace()`, where nnsight puts it, and the constructor no longer sets `device_map="auto"` for you. To keep a model on the CPU pass `device="cpu"`, since `device_map="cpu"` does not ([Loading](../guide/loading.md#device-versus-device_map)). ### `steer` and `skip_layers` The steering vector is `vector`, positions are `token_positions` only, and a list of `layers` has to be ascending because nothing sorts it. `skip_layer(i)` is `skip_layers(i, i)`, and `layer_returns_tuple` is gone, since each family declares whether its blocks return a tuple: ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") torch.manual_seed(0) vector = torch.randn(model.hidden_size) with model.trace("The Eiffel Tower is in"): model.skip_layers(3, 3) # skip_layer(3) model.steer([4, 6], vector, factor=2.0, token_positions=-1) # steer([4, 6], steering_vector=vector, factor=2.0, positions=-1) resid = model.layers[6].layer_output.save() print(model.get_topk_closest_tokens(resid[0, -1], k=3)) # [{' FIRE': 0.4005129635334015, ' exchange': 0.06274925172328949, ' smart': 0.0324697308242321}] ``` `get_topk_closest_tokens` returns a list of one dict per position whatever the rank of its input, where 1.x returned a bare dict for a single position, and its keys are decoded text (`' FIRE'`) rather than vocabulary tokens (`'ĠFIRE'`). `probs_to_dict` takes a distribution and `k`, not the token ids and their probabilities. ### `project_on_vocab` on Gemma-2 1.x computed `lm_head(ln_final(h))`. nnterp 2 adds what the model does after the head, which on Gemma-2 is the softcap `final_logit_softcapping * tanh(logits / final_logit_softcapping)`: a logit lens at the last block now equals `logits`, and the 1.x lens on Gemma-2 did not. On GPT-2, Llama and every family with nothing after the head the two agree. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("google/gemma-2-2b") with model.trace("The Eiffel Tower is in"): last = model.layers[-1].layer_output.save() logits = model.logits.save() lens = model.project_on_vocab(last) raw = model.lm_head(model.norm(last)) # what 1.x computed print(torch.equal(lens, logits), torch.equal(raw, logits), (raw - logits).abs().max().item()) # True False 51.002906799316406 ``` ### Helpers and errors `nnterp.nnsight_utils` keeps `get_token_activations`, the two `collect_*` functions and `compute_next_token_probs` with their 1.x signatures. The `get_layer_output` style accessors are gone, and `model.layers[i].layer_output` is what they wrapped. `nnterp.prompt_utils` is unchanged except that `run_prompts` shows no progress bar unless you pass `tqdm=`. `nnterp.interventions` and `nnterp.display` have no counterpart: `logit_lens` is a loop over `model.project_on_vocab(layer.layer_output)`, which the [Logit lens](../tutorials/logit-lens.md) tutorial writes out. `RenamingError` is replaced by `nnterp.Unavailable` (a value that this checkpoint lacks, with the reason) and `nnterp.UnsupportedFamily` (a `model_type` that nothing can standardize). The [Guide](../guide/index.md) has one page per name and value. # Guide - [Loading](loading.md): how a checkpoint becomes a `StandardizedTransformer`. - [Names](names.md): the module names every family answers to. - [Residual stream](values/residual-stream.md): the stream leaving a block and what each sublayer adds to it. - [Attention](values/attention.md): the tensors inside the attention module, from the queries to the head outputs. - [The whole model](values/whole-model.md): the values on the root, from the input ids to the logits. - [Layouts](values/layouts.md): the axis layout of every value. - [Availability](values/availability.md): which values a checkpoint has. - [Mixture of experts](architectures/mixture-of-experts.md): the router and the experts on every MoE family. - [Gated DeltaNet](architectures/delta-net.md): the linear-attention mixer on Qwen3-Next, Qwen3.5, OLMo-Hybrid and Kimi-Linear. - [Mamba](architectures/mamba.md): the selective-scan mixer on Mamba, Falcon-Mamba and Jamba. - [Mamba-2](architectures/mamba-2.md): the state-space mixer on Mamba-2, Nemotron-H, Bamba and Falcon-H1. - [Vision-language](architectures/vision-language.md): the image side of an image-text-to-text checkpoint. - [Methods](methods.md): `steer`, `skip_layers`, `project_on_vocab` and `get_topk_closest_tokens`. - [Generation](generation.md): the values under `model.generate`. - [Helpers](helpers.md): activations and next-token distributions over many prompts. - [Remote](remote.md): a `StandardizedTransformer` trace on NDIF. # Loading `StandardizedTransformer(repo_id, ...)` is nnsight's `TransformersModel` with one step before the load: it reads the checkpoint's config, takes the family from `model_type`, and hands that family's aliases and envoy classes to nnsight. Every `TransformersModel` keyword passes through unchanged. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu") print(model.config.model_type, model.family.__name__) print(model.num_layers, model.hidden_size, model.num_heads) with model.trace("The Eiffel Tower is in"): resid = model.layers[3].layer_output.save() logits = model.logits.save() print(resid.shape, logits.shape) # gpt2 nnterp.families.gpt2 # 12 768 12 # torch.Size([1, 7, 768]) torch.Size([1, 7, 50257]) ``` `model.family` is the module `nnterp.families.`, imported on first use, and `nnterp.families.known()` lists the shipped ones. A multimodal config nests the language model's config as `text_config`, and that inner `model_type` is the one looked up, so a Gemma 3 wrapper resolves to `gemma3_text` under either task. ## Arguments `family=`, `rename=`, `envoys=` and `tokenizer_kwargs=` are nnterp's own. Everything else, such as `dispatch`, `dtype`, `device_map`, `attn_implementation`, `revision` and `task`, goes to `TransformersModel` as documented at [nnsight.net](https://nnsight.net/features/5_loading/), and `revision` and `trust_remote_code` also reach the config read that picks the family. The full signature is under the Documentation tab ([standardized](../documentation/standardized.md)). `task` defaults to `"text-generation"`, and `task="image-text-to-text"` loads a vision-language checkpoint as its wrapper, with its processor ([Vision-language](architectures/vision-language.md)). Without `dtype` the weights load in the config's own dtype, float32 for GPT-2 and bfloat16 for Qwen2.5. Without `dispatch=True` the model is built on the meta device and the weights load at the first trace, while `model.support()` and the sizes, which read the config, work either way. ### `attn_implementation="eager"` The values inside attention (`attention_probabilities`, `attention_scores`, the queries, keys, values and head outputs) are read inside transformers' eager attention forward. The default implementation is the checkpoint's, usually `sdpa`, and those values are then unavailable: ```python model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu") print(model.config._attn_implementation) print(model.support(layer=0)["self_attn.attention_probabilities"]) # sdpa # read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' ``` Pass `attn_implementation="eager"` whenever a trace will touch the inside of attention. The boundary values (`layer_output`, `attention_output`, `mlp_output`, `logits`) do not depend on it ([Availability](values/availability.md)). ### `device=` versus `device_map=` `device="cpu"` keeps the whole model on the CPU. `device_map="cpu"` does not: nnsight's pipeline passes its own `device` as well, and on a machine with a GPU the model lands on `cuda:0`. Use `device=` for one device and `device_map="auto"` to spread a model over several. ```python on_cpu = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu") mapped = StandardizedTransformer("openai-community/gpt2", dispatch=True, device_map="cpu") print(on_cpu.device, mapped.device) # cpu cuda:0 ``` ### `rename=` and `envoys=` Both merge over the family's, and a key you give wins. The key forms are nnsight's ([Module Access](https://nnsight.net/features/6_modules/)): a dotted key resolves from the root and mounts there, and a single-component key binds on every envoy that has a child of that name. ```python model = StandardizedTransformer( "openai-community/gpt2", dispatch=True, device="cpu", rename={"transformer.wpe": "embed_positions", "mlp": "ffn"}, ) print(model.embed_positions is model.transformer.wpe) print(model.layers[0].ffn is model.layers[0].mlp) print(model.layers[0].self_attn is model.transformer.h[0].attn) # True # True # True ``` `envoys=` maps a module type or a dotted path to an `Envoy` subclass. nnsight tries the type keys first, and the families key their envoys by type, so only a type key displaces a family's envoy. With `envoys={GPT2MLP: MyMlp}`, a subclass of `nnterp.families.gpt2.Mlp`, `type(model.layers[0].mlp) is MyMlp` on every block. A path key reaches a module that no type key in the merged map matches, by its native path or any alias spelling of it (`"norm"` reaches GPT-2's `ln_f`). ### `tokenizer_kwargs=` `tokenizer_kwargs=` sets attributes on the loaded tokenizer, for the two that a batch of prompts usually needs. ```python model = StandardizedTransformer( "openai-community/gpt2", dispatch=True, device="cpu", tokenizer_kwargs={"padding_side": "left", "pad_token": "<|endoftext|>"}, ) with model.trace(["Hello", "The Eiffel Tower is in"]): mask = model.attention_mask.save() print(mask) # tensor([[0, 0, 0, 0, 0, 0, 1], # [1, 1, 1, 1, 1, 1, 1]]) ``` ### `family=` `family=` gives one load its family without reading the config or the registry: a module, or any object with `RENAME` and `ENVOYS` such as `types.SimpleNamespace(**vars(gpt2))`. To cover a `model_type` for every load, register the family instead ([Registering](../developing/registering.md)). ## A `model_type` with no family A `model_type` with no family module and nothing registered loads with `nnterp.families.default`, a best-effort family, and a warning. The default maps the spellings the shipped families use onto Llama's names and checks its guess at load on a shape-only scan. It serves `layer_output`, the root values, the sizes and the attention interior, and it refuses what a name cannot tell: `attention_output`, `mlp_output`, `layer_mid`, `mlp_norm_output`, a mixture's interior and `project_on_vocab`. GPT-2 has a family, so this block uses a tiny random checkpoint of a type that has none. ```python model = StandardizedTransformer("hf-tiny-v2/tiny-random-DiffLlamaForCausalLM", dispatch=True, device="cpu") # UserWarning: nnterp has no family for model_type 'diffllama'; the default family standardizes it as a # best-effort guess. Check model.support() for what it found, and add nnterp/families/diffllama.py (or # nnterp.families.register(family, 'diffllama')) for a standardization you can rely on. print(model.family.__name__) support = model.support() print(support["layer_output"]) print(support["mlp.mlp_output"][0]) # nnterp.families.default # None # the default family cannot tell what this sublayer adds to the stream; add a family module (docs/extending/adding-a-family.md) ``` When the guess finds no embedding, blocks, final norm or head, the load raises `UnsupportedFamily` and suggests a `rename=` read off the module tree; passing it completes the default's names. [Adding a family](../developing/adding-a-family.md) is the standardization you can rely on, and the families that have it are the pages of the [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/). ```python from nnterp import UnsupportedFamily try: StandardizedTransformer("hf-tiny-v2/tiny-random-CTRLLMHeadModel", dispatch=True, device="cpu") except UnsupportedFamily as error: print(error) # the default family cannot standardize CTRLLMHeadModel: found no embed_tokens, norm under any name it knows. # Name the native modules with rename={'': ''} at load; going by the module tree, # perhaps rename={'transformer.w': 'embed_tokens', 'transformer.layernorm': 'norm'}, or add a family module # named after the model_type (docs/extending/adding-a-family.md). model = StandardizedTransformer( "hf-tiny-v2/tiny-random-CTRLLMHeadModel", dispatch=True, device="cpu", rename={"transformer.w": "embed_tokens", "transformer.layernorm": "norm"}, ) print(model.family.__name__, model.support(layer=0)["layer_output"]) # nnterp.families.default None ``` ## An already-loaded module `repo_id` can be a `torch.nn.Module`. Its own `config` picks the family and the module is wrapped as it is, so a model you built or edited yourself gets the same tree. ```python from transformers import AutoModelForCausalLM raw = AutoModelForCausalLM.from_pretrained("openai-community/gpt2", attn_implementation="eager") model = StandardizedTransformer(raw) print(model.family.__name__) with model.trace("Hello world"): x = model.layers[0].layer_output.save() print(x.shape) # nnterp.families.gpt2 # torch.Size([1, 2, 768]) ``` The names the loaded tree answers to are on [Names](names.md). # Names Every family's tree also answers to Llama's names: `model.embed_tokens`, `model.layers[i]` with its `self_attn` and `mlp`, `model.norm` and `model.lm_head`. The names are nnsight aliases, a second attribute on the same envoy, so the native name keeps working and both reach one object. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu") with model.trace("The Eiffel Tower is in"): emb = model.embed_tokens.output.save() attn_in = model.layers[5].self_attn.input.save() mlp_in = model.layers[5].mlp.input.save() normed = model.norm.output.save() logits = model.lm_head.output.save() print(model.layers[5] is model.transformer.h[5]) print(model.layers[5].self_attn is model.transformer.h[5].attn) print(model.layers[5].self_attn.path, model.get("layers.5.self_attn").path) # True # True # model.transformer.h.5.attn model.transformer.h.5.attn ``` The same trace body runs on `HuggingFaceTB/SmolLM2-135M` (a Llama) and `EleutherAI/pythia-70m-deduped` (a GPT-NeoX). `envoy.path` is always the native path, and `model.get()` resolves a dotted string through the aliases. ## The vocabulary | name | GPT-2 | Llama | GPT-NeoX | | --- | --- | --- | --- | | `model.embed_tokens` | `transformer.wte` | `model.embed_tokens` | `gpt_neox.embed_in` | | `model.layers[i]` | `transformer.h[i]` | `model.layers[i]` | `gpt_neox.layers[i]` | | `model.layers[i].self_attn` | `...h[i].attn` | same | `...layers[i].attention` | | `model.layers[i].mlp` | same | same | same | | `model.norm` | `transformer.ln_f` | `model.norm` | `gpt_neox.final_layer_norm` | | `model.lm_head` | same | same | same | The containers are lifted out of the inner model, so it is `model.layers` on every family, never `model.model.layers`. Inside the sublayers the output projection, the activation and the down projection are `o_proj`, `act_fn` and `down_proj` (GPT-2's `c_proj`, `act` and `c_proj`). A module with no standard role keeps its native name only, such as GPT-2's `wpe`. Each family's own spellings are on its page under the Documentation tab, with [llama](../documentation/families/llama.md) as the reference. ## Key forms A family's `RENAME` is nnsight's `rename=` ([Module Access](https://nnsight.net/features/6_modules/)). A dotted key such as `transformer.h` resolves from the root and mounts there, which is what lifts `layers` to `model.layers`. A single-component key such as `attn` binds on every envoy that has a child of that name, so one entry renames the attention in every block. A key that resolves nowhere is skipped, which is how `embed_out` covers the older spelling of GPT-NeoX's head alongside `lm_head`. Because a single-component key binds everywhere, no family aliases a block-level name that also exists outside the blocks: OPT's block norm `final_layer_norm` keeps its native name, since an alias would also bind on the decoder's `final_layer_norm`, which is `model.norm`. ## The norms `input_layernorm` and `post_attention_layernorm` bind as aliases where a family spells them otherwise (GPT-2's `ln_1` and `ln_2`), but they are not part of the vocabulary, because their meaning differs between families. On a sequential block (Llama, GPT-2) the first feeds the attention and the second the MLP. On Gemma-2 and Gemma-3 `post_attention_layernorm` norms the attention's output, and the MLP's pre-norm is `pre_feedforward_layernorm`. On a parallel block (GPT-J, Phi, Falcon-7B) one `input_layernorm` feeds both sublayers, and OLMo-2 has only post-norms, so no `input_layernorm` exists to alias. What matters about a norm is its output, and that is `self_attn.input` and `mlp.input` on every family: ```python import torch with model.trace("The Eiffel Tower is in"): norm1 = model.layers[5].input_layernorm.output.save() attn_in = model.layers[5].self_attn.input.save() norm2 = model.layers[5].post_attention_layernorm.output.save() mlp_in = model.layers[5].mlp.input.save() print(torch.equal(norm1, attn_in), torch.equal(norm2, mlp_in)) # True True ``` On GPT-2 the two agree. The accessor form reads the right tensor on Gemma-2 and OLMo-2 as well, where the norm aliases would not. ## Blocks without a module A hybrid (Qwen3-Next, Qwen3.5, OLMo-Hybrid) has `linear_attn` on most blocks and `self_attn` on a few, never both, and OPT has no `mlp` module at all. Decide which blocks have which outside the trace, with `getattr(layer, "self_attn", None)`. [Availability](values/availability.md) shows how `support()` reports them. ## `print(model)` The repr prints a standard name beside its native one as `alias/native`, and each standard value with its layout and description, so `print(model.layers[i])` is the reference for what a block of this family has. ```python print(model.layers[0].self_attn) # GPT2Attention( # (c_attn): Conv1D() # (o_proj/c_proj): Conv1D() # (attn_dropout): Dropout(p=0.1, inplace=False) # (resid_dropout): Dropout(p=0.1, inplace=False) # (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention # (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention # (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention # (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked # (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it # (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream # (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with # (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection # ) ``` A value is listed whether or not this checkpoint has it; `support()` is what says. The [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/) draws each family's block under these names, with the native name beside each. The values themselves start with the [residual stream](values/residual-stream.md). # Residual stream `layer_output` is the residual stream leaving block `i`, `layer_input` the stream entering it, and `attention_output` and `mlp_output` what the two sublayers add to it. All four are `[batch, seq, hidden_size]` tensors on every family, and the two contributions are defined by one identity: `layer_input + attention_output + mlp_output == layer_output`. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): x = model.layers[5].layer_input.save() attn = model.layers[5].self_attn.attention_output.save() mlp = model.layers[5].mlp.mlp_output.save() out = model.layers[5].layer_output.save() print(x.shape, attn.shape, mlp.shape, out.shape) # torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) print(torch.equal(x + attn + mlp, out)) # True ``` The identity is exact on a sequential block. A parallel block, such as Pythia's or Falcon's, sums the three terms in another order and holds it to rounding (one float16 ulp on Pythia-70m), so compare with a tolerance in the block's dtype there. ## A tensor on every family A block's native `.output` is a tuple on GPT-J, GPT-Neo, BLOOM, MPT and Falcon. `layer_output` is the tensor the block returned, whatever wraps it: ```python bloom = StandardizedTransformer("bigscience/bloom-560m") with bloom.trace("The Eiffel Tower is in"): raw = bloom.layers[2].output.save() out = bloom.layers[2].layer_output.save() print(type(raw), type(out), torch.equal(raw[0], out)) # True ``` Assigning `layer_output` on such a block replaces the first element and leaves the others as they were, so the next block still receives a tuple. `attention_output` does the same where the attention module returns `(attn_output, attn_weights)`. A contribution is what the block adds to the stream, which is not always the module's return. On the families with a norm after the attention module, such as Gemma-2 and OLMo-2, `attention_output` is that norm's output and `self_attn.output[0]` is the tensor before that norm. On BLOOM and MPT the module adds the residual inside, and the contribution is the tensor entering that add. The identity is what to rely on, not `self_attn.output`. ## Between the sublayers Between `layer_input` and `layer_output` the block has four more module boundaries, and the sublayers three more values inside them: | value | where it is read | shape | | --- | --- | --- | | `layers[i].attention_norm_output` | `input_layernorm.output`, what the attention reads | `[batch, seq, hidden_size]` | | `layers[i].self_attn.attention_premix` | `o_proj.input`, every head's output before the output projection | `[batch, seq, num_heads, head_dim]` | | `layers[i].layer_mid` | `post_attention_layernorm.input`, the stream after the attention add | `[batch, seq, hidden_size]` | | `layers[i].mlp_norm_output` | `post_attention_layernorm.output` | `[batch, seq, hidden_size]` | | `layers[i].mlp_input` | `mlp.input`, what the MLP is called with | `[batch, seq, hidden_size]` | | `layers[i].mlp.mlp_activation` | `act_fn.output` | `[batch, seq, intermediate_size]` | | `layers[i].mlp.mlp_neurons` | `down_proj.input`, what the down projection sums | `[batch, seq, intermediate_size]` | ```python block = model.layers[5] with model.trace("The Eiffel Tower is in"): x = block.layer_input.save() normed = block.attention_norm_output.save() premix = block.self_attn.attention_premix.save() attn = block.self_attn.attention_output.save() mid = block.layer_mid.save() mlp_normed = block.mlp_norm_output.save() mlp_in = block.mlp_input.save() act = block.mlp.mlp_activation.save() neurons = block.mlp.mlp_neurons.save() print(premix.shape, act.shape, neurons.shape) # torch.Size([1, 7, 12, 64]) torch.Size([1, 7, 3072]) torch.Size([1, 7, 3072]) print(torch.equal(mid, x + attn), torch.equal(mlp_normed, mlp_in), torch.equal(act, neurons)) # True True True ``` `layer_mid` is `layer_input + attention_output`, and `mlp_norm_output` is `mlp_input` where nothing sits between the norm and the MLP. On GPT-2's plain MLP `mlp_neurons` is `mlp_activation` itself; on a gated MLP (Llama, Qwen) it is `act(gate_proj(x)) * up_proj(x)`, a different tensor at the same place. The keys above are those of a pre-norm block, under the standard names of [Names](../names.md). A block built otherwise points a value at its own module or marks it unavailable with the reason: a parallel block has no `layer_mid`, a Gemma-2 block norms the MLP's input in `pre_feedforward_layernorm`, and a routed mixture of experts has no single `mlp_activation` or `mlp_neurons` ([Availability](availability.md)). ## Reading, editing in place, assigning Every value on this page is the live tensor, so an in-place edit reaches the model and an assignment replaces it, as with [`.output`](https://nnsight.net/features/2_setting/): ```python with model.trace("The Eiffel Tower is in"): clean = model.logits.save() with model.trace("The Eiffel Tower is in"): model.layers[5].self_attn.attention_output[:, -1] = 0 model.layers[5].mlp.mlp_output[:] = 0 model.layers[7].layer_output = model.layers[7].layer_output * 2 edited = model.logits.save() print((edited - clean).abs().max().item()) # 31.438919067382812 ``` Within one trace, read a block's values in forward order. `model.order(i)` gives that order for block `i`, with equal ranks for values read at one location, and reading `layer_input` after `mlp_output` raises nnsight's `OutOfOrderError`: ```python print(model.order(5)) # {'layer_input': 0, 'attention_norm_output': 1, 'self_attn.attention_premix': 2, 'self_attn.attention_output': 3, 'layer_mid': 4, 'mlp_norm_output': 5, 'mlp_input': 6, 'mlp.mlp_activation': 7, 'mlp.mlp_neurons': 8, 'mlp.mlp_output': 9, 'layer_output': 10} ``` The identity takes another form on Gemma-4, which multiplies the sum by a per-block `layer_scalar`, on Doge and ZAYA, which gate the stream per channel before each add, and on DeepSeek-V4, where `layer_output` is `hc_mult` parallel streams, `[batch, seq, streams, hidden_size]`. Each family's page in the [Documentation](../../documentation/families/index.md) tab gives its formula, and [`Layer`](../../documentation/components/layer.md) the signatures. The values inside the attention sublayer are on the [Attention](attention.md) page. # Attention Six values on `model.layers[i].self_attn` are what happens inside the attention sublayer: the queries, keys and values that enter the softmax attention, the scores entering the softmax, the pattern the values are mixed with, and each head's output before the heads are concatenated and projected. They are read inside transformers' eager attention forward, so the model has to be loaded with `attn_implementation="eager"`. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") attn = model.layers[5].self_attn with model.trace("The Eiffel Tower is in"): q = attn.attention_queries.save() k = attn.attention_keys.save() v = attn.attention_values.save() scores = attn.attention_scores.save() pattern = attn.attention_probabilities.save() heads = attn.attention_head_outputs.save() print(q.shape, k.shape, v.shape) # torch.Size([1, 12, 7, 64]) torch.Size([1, 12, 7, 64]) torch.Size([1, 12, 7, 64]) print(scores.shape, pattern.shape, heads.shape) # torch.Size([1, 12, 7, 7]) torch.Size([1, 12, 7, 7]) torch.Size([1, 7, 12, 64]) print(torch.allclose(scores.softmax(-1), pattern), pattern[0, 0].sum(-1)) # True tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000], # grad_fn=) ``` | value | shape | what it is | | --- | --- | --- | | `attention_queries` | `[batch, num_heads, seq, qk_head_dim]` | the queries entering attention, after the query projection and, where the family has one, the rotary embedding | | `attention_keys` | `[batch, num_kv_heads, seq, qk_head_dim]` | the keys entering attention, before `repeat_kv` | | `attention_values` | `[batch, num_kv_heads, seq, head_dim]` | the values entering attention, before `repeat_kv` | | `attention_scores` | `[batch, num_heads, query, key]` | the scaled scores entering the softmax, with the causal mask applied | | `attention_probabilities` | `[batch, num_heads, query, key]` | the pattern the values are mixed with: after the softmax and its dropout, in the model's dtype | | `attention_head_outputs` | `[batch, seq, num_heads, head_dim]` | each head's mix of the values, before the output projection | The axis names are the sizes on the attention envoy, `attn.num_heads`, `attn.num_kv_heads`, `attn.head_dim` and `attn.qk_head_dim`, read off the module ([Layouts](layouts.md)). `qk_head_dim` equals `head_dim` outside latent attention (DeepSeek-V2 and V3). The pattern's rows sum to one, except on a model with an attention sink (GPT-OSS, MiMo-V2-Flash, DeepSeek-V4), where the sink's column has already been dropped and they sum to less. ## Eager attention A load without `attn_implementation` gets transformers' default, `sdpa` on most families, and the eager forward never runs. `model.support()` gives the reason per value, and reading one raises `nnterp.Unavailable` with the same reason before the model runs. `attention_output` and `attention_premix` (`o_proj.input`, the heads as the output projection receives them) stay available under any implementation. ```python default = StandardizedTransformer("openai-community/gpt2") print(default.support(layer=5)["self_attn.attention_probabilities"]) # read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' print(default.support(layer=5)["self_attn.attention_output"]) # None ``` ## Editing Each value is the tensor the forward holds, so an in-place edit lands and an assignment replaces that one value: assigning the keys leaves the queries and values as they were. GPT-2's and GPT-BigCode's queries are the exception: `c_attn` produces one tensor that `split` divides into queries, keys and values, and torch refuses an in-place edit of the query view (`RuntimeError: Output 0 of Select is a view and is being modified inplace`), so assign them. Their keys and values take in-place edits. ```python with model.trace("The Eiffel Tower is in"): clean = model.logits.save() with model.trace("The Eiffel Tower is in"): attn.attention_head_outputs[:, :, 7] = 0 ablated = model.logits.save() with model.trace("The Eiffel Tower is in"): attn.attention_queries = attn.attention_queries * 0 attn.attention_keys[:, :, 0] = 0 replaced = model.logits.save() print((ablated - clean).abs().max().item(), (replaced - clean).abs().max().item()) # 2.7936935424804688 13.05731201171875 ``` !!! warning "Overwriting or multiplying `attention_scores` lifts the causal mask" The masked entries hold the dtype's minimum (`-inf` on GPT-Neo). `scores[:] = 0` or `scores * 0` turns them into real scores, every query then attends to the tokens after it, and the pattern still looks like a valid distribution. Add to the scores instead, or keep the masked entries: `torch.where(scores <= torch.finfo(scores.dtype).min, scores, new)`. ```python with model.trace("The Eiffel Tower is in"): attn.attention_scores[:, :, :, 0] += 5.0 added = attn.attention_probabilities.save() with model.trace("The Eiffel Tower is in"): attn.attention_scores[:] = 0 overwritten = attn.attention_probabilities.save() print(added[0, 0, 2]) # tensor([9.9987e-01, 8.8083e-05, 4.0907e-05, 0.0000e+00, 0.0000e+00, 0.0000e+00, # 0.0000e+00], grad_fn=) print(overwritten[0, 0, 2]) # tensor([0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429, 0.1429], # grad_fn=) ``` ## Grouped-query attention The keys and values are read before `repeat_kv`, so on a checkpoint with fewer key/value heads than query heads their head axis is `num_kv_heads` wide while the queries, the scores, the pattern and the head outputs are `num_heads` wide. Query head `h` reads key/value head `h // (num_heads // num_kv_heads)`, so a head-wise edit of the keys touches every query head of that group. On Qwen2.5-0.5B, 14 query heads share 2 key/value heads: ```python qwen = StandardizedTransformer("Qwen/Qwen2.5-0.5B", attn_implementation="eager") attn = qwen.layers[5].self_attn print(attn.num_heads, attn.num_kv_heads, attn.head_dim) # 14 2 64 with qwen.trace("The Eiffel Tower is in"): k = attn.attention_keys.save() v = attn.attention_values.save() pattern = attn.attention_probabilities.save() heads = attn.attention_head_outputs.save() print(k.shape, pattern.shape, heads.shape) # torch.Size([1, 2, 7, 64]) torch.Size([1, 14, 7, 7]) torch.Size([1, 7, 14, 64]) groups = attn.num_heads // attn.num_kv_heads mixed = (pattern @ v.repeat_interleave(groups, dim=1)).transpose(1, 2) print(torch.allclose(mixed, heads, atol=1e-5)) # True ``` ## Read order The table lists the six in forward order, and all of them come before `attention_output`. Several blocks' values in one trace are fine in ascending block order; the rule is forward order across the whole model, not only within a block. `model.order(i)` gives the order for block `i`, with equal ranks for values read at one location: ```python print(qwen.order(5)) # {'layer_input': 0, 'attention_norm_output': 1, 'self_attn.attention_queries': 2, 'self_attn.attention_keys': 2, 'self_attn.attention_values': 2, 'self_attn.attention_scores': 3, 'self_attn.attention_probabilities': 4, 'self_attn.attention_head_outputs': 5, 'self_attn.attention_premix': 6, 'self_attn.attention_output': 7, 'layer_mid': 8, 'mlp_norm_output': 9, 'mlp_input': 10, 'mlp.mlp_activation': 11, 'mlp.mlp_neurons': 12, 'mlp.mlp_output': 13, 'layer_output': 14} ``` A read out of order raises nnsight's `OutOfOrderError`, which names the internal location rather than the value: ```python attn = model.layers[5].self_attn try: with model.trace("The Eiffel Tower is in"): out = attn.attention_output.save() pattern = attn.attention_probabilities.save() except Exception as e: print(type(e).__name__, e) # OutOfOrderError 'model.transformer.h.5.attn.source.attention_interface_1.fn.i0' was requested but the model already ran past it ``` Falcon's attention does its own arithmetic, and without alibi its values bind before the queries and keys leave the rotary embedding, so read `attention_values` first there; with alibi the order is queries, keys, values. Under `model.generate` the values are per forward call: on a decode step the queries, the scores, the pattern and the head outputs have one query position, and the keys, the values and the pattern's key axis are as long as the cache ([Generation](../generation.md)): ```python with model.generate("The Eiffel Tower is in", max_new_tokens=2) as tracer: for _ in tracer.iter[1:2]: pattern = attn.attention_probabilities.save() print(pattern.shape) # torch.Size([1, 12, 1, 8]) ``` The signatures are on [`Attention`](../../documentation/components/attention.md), and the values of the whole model, from the input to the logits, are on [The whole model](whole-model.md). # The whole model Inside a trace the root envoy answers for the whole run: `logits`, `lm_head_output`, `final_norm_output`, `token_embeddings`, `next_token_probs`, and the `input_ids`, `attention_mask` and `input_size` the model was called with. Outside a trace it answers for the sizes, `num_layers`, `hidden_size`, `num_heads` and the rest, read off the config. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") with model.trace("The Eiffel Tower is in"): ids = model.input_ids.save() mask = model.attention_mask.save() emb = model.token_embeddings.save() normed = model.final_norm_output.save() head = model.lm_head_output.save() logits = model.logits.save() probs = model.next_token_probs.save() print(ids.shape, emb.shape, normed.shape, logits.shape, probs.shape) # torch.Size([1, 7]) torch.Size([1, 7, 768]) torch.Size([1, 7, 768]) torch.Size([1, 7, 50257]) torch.Size([1, 50257]) print(torch.equal(head, logits), model.probs_to_dict(probs[0], k=3)) # True {' the': 0.2401435822248459, ' a': 0.08623422682285309, ' fact': 0.022629808634519577} ``` `probs_to_dict` gives the `k` most likely tokens of one distribution ([Methods](../methods.md)). | value | shape | what it is | writes | | --- | --- | --- | --- | | `input_ids` | `[batch, seq]` | the token ids the model was called with | assign | | `attention_mask` | `[batch, seq]` | the mask the model was called with; zeros are padding | assign | | `input_size` | `torch.Size([batch, seq])` | the ids' shape | read-only | | `token_embeddings` | `[batch, seq, hidden_size]` | the embedding module's output | assign | | `final_norm_output` | `[batch, seq, hidden_size]` | the final norm's output, what the head reads | assign, in place | | `lm_head_output` | `[batch, seq, vocab_size]` | the head's own projection | assign, in place | | `logits` | `[batch, seq, vocab_size]` | the logits the model returns, after anything it does past the head | assign | | `next_token_probs` | `[batch, vocab_size]` | `logits[:, -1].softmax(-1)` | read-only | The table is in forward order: the inputs first, `token_embeddings` before any block's value, and the last three after them all. ## logits and lm_head_output `logits` is what the model returns. On most families it is `lm_head_output`. Gemma-2 softcaps the head's output (`config.final_logit_softcapping`), Cohere multiplies it by `logit_scale` and Granite divides it by `logits_scaling`, and `logits` is the result. `model.project_on_vocab` ends with the same step, so the logit lens on the last block's `layer_output` equals `logits` on every family ([Methods](../methods.md)). ```python gemma = StandardizedTransformer("google/gemma-2-2b") with gemma.trace("The Eiffel Tower is in"): head = gemma.lm_head_output.save() logits = gemma.logits.save() cap = gemma.config.final_logit_softcapping print(cap, torch.equal(head, logits), torch.allclose(logits, cap * torch.tanh(head / cap))) # 30.0 False True print(head[0, -1].abs().max().item(), logits[0, -1].abs().max().item()) # 42.63672637939453 26.695619583129883 ``` Assigning `logits` replaces them in the model's output, so `tracer.result.logits` is what you set. ## token_embeddings and block 0 `token_embeddings` is `embed_tokens.output`, with whatever the embedding module does itself (Gemma's `sqrt(hidden_size)` scale) and nothing the model adds after it. GPT-2 adds its positional embeddings, BLOOM norms the output and Granite multiplies it by `embedding_multiplier`, so the stream entering block 0 is `layers[0].layer_input`: ```python with model.trace("The Eiffel Tower is in"): emb = model.token_embeddings.save() pos = model.transformer.wpe.output.save() first = model.layers[0].layer_input.save() print(torch.equal(emb, first), torch.equal(emb + pos, first)) # False True ``` ## next_token_probs `next_token_probs` is the softmax at the last position, derived from the output. That position is the last token of every row under left padding, which nnsight sets on a causal model's tokenizer: ```python with model.trace(["Hi", "The Eiffel Tower is in"]): ids = model.input_ids.save() probs = model.next_token_probs.save() print(ids) # tensor([[50256, 50256, 50256, 50256, 50256, 50256, 17250], # [ 464, 412, 733, 417, 8765, 318, 287]]) print(model.probs_to_dict(probs[0], k=3)) # {'.': 0.07607239484786987, ',': 0.06419526785612106, '\n': 0.035513173788785934} print(model.probs_to_dict(probs[1], k=3)) # {' the': 0.2401452511548996, ' a': 0.0862361341714859, ' fact': 0.02262996695935726} ``` The softmax is in the model's dtype, and in float16 thousands of entries underflow to exact zeros, so take a KL on `model.logits[:, -1].float().log_softmax(-1)` instead. Assigning `next_token_probs` raises `AttributeError: next_token_probs is derived from the logits and cannot be assigned; assign model.logits instead`. ## The input `input_ids` and `attention_mask` are what the model was called with, and the model runs on what you assign to them. Assigning `input_ids` does not resize the mask, so assign `attention_mask` too when the length changes: ```python with model.trace("Paris is the capital of"): other_ids = model.input_ids.save() other = model.logits.save() with model.trace("The Eiffel Tower is in"): model.input_ids = other_ids model.attention_mask = torch.ones_like(other_ids) swapped = model.logits.save() print(torch.equal(swapped, other)) # True ``` `input_size` is the ids' `torch.Size`, read-only. A `torch.Size` bound inside the trace does not survive it, so save `torch.tensor(model.input_size)` to keep it. ## Sizes The sizes are read-only attributes of the root, readable before any trace and without `dispatch`: ```python print(model.num_layers, model.hidden_size, model.vocab_size) # 12 768 50257 print(model.num_heads, model.num_kv_heads, model.head_dim, model.qk_head_dim, model.intermediate_size) # 12 12 64 64 3072 print(model.layers[5].self_attn.num_heads, model.layers[5].self_attn.head_dim, model.layers[5].mlp.intermediate_size) # 12 64 3072 ``` | size | the plain rule | | --- | --- | | `num_layers` | `len(model.layers)` | | `hidden_size` | `config.hidden_size` | | `vocab_size` | `config.vocab_size` | | `num_heads` | `config.num_attention_heads` | | `num_kv_heads` | `config.num_key_value_heads`, else `num_heads` | | `head_dim` | `config.head_dim` when the config has it (Qwen3, Gemma), else `hidden_size // num_heads` | | `qk_head_dim` | `head_dim`, unless the family separates them (DeepSeek's latent attention) | | `intermediate_size` | `config.intermediate_size` | The rule reads the text config (a multimodal checkpoint's `text_config`), unless the family's module defines a function of the same name, which it does where its config spells a size its own way: Falcon's `num_kv_heads` is `config.num_kv_heads` on the 40B layout and 1 under `multi_query`, GPT-2's `intermediate_size` is `config.n_inner`, four times `hidden_size` when unset. `model.hidden_size = 5` raises an `AttributeError` saying that the size is read off the config and that a family defines `def hidden_size(model)` to say it otherwise. A root size is the config's value, which equals every block's where the blocks agree. A block's own sizes are on its attention (`num_heads`, `num_kv_heads`, `head_dim`, `qk_head_dim`) and its MLP (`intermediate_size`, one routed expert's width on a mixture of experts), read off the module, and they differ from the root's on Gemma-4 (`head_dim` on every checkpoint, `num_kv_heads` on 26B-A4B and 31B, the MLP width on E2B), MiMo-V2-Flash (`num_kv_heads`) and Laguna (`num_heads`). The signatures are on [`StandardizedTransformer`](../../documentation/standardized.md), and the axes behind every shape on this page are on [Layouts](layouts.md). # Layouts Every value has one axis layout, the same on every family: `attention_probabilities` is `[batch, num_heads, query, key]` on GPT-2, Llama and BLOOM alike. The layout is a named `jaxtyping` type from `nnterp.components` (`Residual`, `Pattern`, `Queries`, ...), read off the value as `value.layout`, with its axes as `value.dims`. ```python from nnterp import Attention, Layer, StandardizedTransformer from nnterp.components import Pattern, Queries, Residual print(Layer.layer_output.dims, Layer.layer_output.layout is Residual) print(Attention.attention_queries.dims) model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu", attn_implementation="eager") with model.trace("The Eiffel Tower is in"): q = model.layers[0].self_attn.attention_queries.save() pattern = model.layers[0].self_attn.attention_probabilities.save() out = model.layers[0].layer_output.save() print(tuple(q.shape), tuple(pattern.shape), tuple(out.shape)) print(isinstance(q, Queries), isinstance(pattern, Pattern), isinstance(out, Pattern)) print(isinstance(q[0], Queries), isinstance(q.long(), Queries)) # ('batch', 'seq', 'hidden_size') True # ('batch', 'num_heads', 'seq', 'qk_head_dim') # (1, 12, 7, 64) (1, 12, 7, 7) (1, 7, 768) # True True False # False False ``` `isinstance(tensor, value.layout)` checks rank and dtype, not axis sizes: `q[0]` fails on rank and `q.long()` on dtype. Read the layout off the class (`Attention.attention_queries`) or off the instance's type (`type(model.layers[0].self_attn).attention_queries`). A family that redefines a value annotates it with the same name, so the layout cannot drift from the base. ## The names The values GPT-2 serves carry these layouts, and `print(model)` shows each value's beside its description ([Names](../names.md#printmodel)). | layout | axes | values | | --- | --- | --- | | `Residual` | `batch seq hidden_size` | `layer_input`, `layer_output`, `attention_output`, `mlp_output`, `token_embeddings`, the block's norm boundaries | | `Logits` | `batch seq vocab_size` | `logits`, `lm_head_output` | | `NextTokenProbs` | `batch vocab_size` | `next_token_probs` | | `Tokens` | `batch seq`, integer | `input_ids`, `attention_mask` | | `Queries` | `batch num_heads seq qk_head_dim` | `attention_queries` | | `Keys` | `batch num_kv_heads seq qk_head_dim` | `attention_keys` | | `Values` | `batch num_kv_heads seq head_dim` | `attention_values` | | `Pattern` | `batch num_heads query key` | `attention_scores`, `attention_probabilities` | | `HeadOutputs` | `batch seq num_heads head_dim` | `attention_head_outputs`, `attention_premix` | | `Neurons` | `batch seq intermediate_size` | `mlp_activation`, `mlp_neurons` | There are 37 names in all. The mixture's (`RouterLogits`, `ExpertWeights`, ...), the recurrent mixers' (`State`, `ScanQK`, `SSDValues`, ...) and the vision tower's (`Patches`, `ImageFeatures`) are on their Guide pages and defined beside the envoy that serves them under the Documentation tab ([components](../../documentation/components/index.md)). `nnterp.components` exports 34 of them; the root's three, `Logits`, `NextTokenProbs` and `Tokens`, are in `nnterp.standardized`. A value with no tensor annotation, such as `input_size`, has `layout` and `dims` of `None`. ## Axes named as sizes An axis that is a fact about the model is named exactly as the size attribute that holds it (`hidden_size`, `vocab_size`, `num_heads`, `num_kv_heads`, `head_dim`, `qk_head_dim`, `intermediate_size`, `num_experts`, `top_k`), and `nnterp.axis_size(envoy, axis)` resolves it: the attribute on the envoy's class when it declares one, else on the first envoy above it that does, up to the root. The names are constants in `nnterp.components` (`HEADS` is `"num_heads"`). ```python from nnterp import axis_size attn = model.layers[0].self_attn print(axis_size(attn, "num_heads"), axis_size(attn, "num_kv_heads"), axis_size(attn, "head_dim")) print(axis_size(model.layers[0], "hidden_size"), axis_size(model, "vocab_size")) print({axis: axis_size(attn, axis) for axis in Attention.attention_keys.dims if axis not in ("batch", "seq")}) try: axis_size(attn, "key") except ValueError as error: print(error) # 12 12 64 # 768 50257 # {'num_kv_heads': 12, 'qk_head_dim': 64} # axis 'key' of a value on model.transformer.h.0.attn is not a fact about the model: no envoy from there up to the root holds a size by that name ``` `num_heads` on an attention value is that block's own count, which differs from the root's on Gemma-4, MiMo-V2-Flash and Laguna. The other axes (`batch`, `seq`, `query`, `key`, a mixer's `heads`, `key_dim`, `state_dim`, ...) are not facts about the model, so `axis_size` refuses them and their size is the tensor's. ## The sequence axis Batch is axis 0 everywhere. The sequence axis is 1 on every value that has one, except the queries, keys and values of softmax attention, where it is 2, the `[batch, num_heads, seq, head_dim]` layout transformers hands its attention interface. `attention_head_outputs` is back to sequence first, and a pattern has two token axes, `query` then `key`. The last token is therefore `out[:, -1]` for a residual-stream value, `q[:, :, -1]` for the queries, and `pattern[:, :, -1, :]` for the last query row. ```python print(tuple(out[:, -1].shape), tuple(q[:, :, -1].shape), tuple(pattern[:, :, -1, :].shape)) # (1, 768) (1, 12) (1, 12, 7) ``` Under grouped-query attention the keys and values are read before `repeat_kv`, so their head axis is `num_kv_heads` wide: on Qwen3-0.6B the queries come out `(1, 16, seq, 128)` and the keys `(1, 8, seq, 128)`. GPT-2 has no grouping, so the two counts coincide above. ## The exceptions On three families one value has a layout of its own. On DeepSeek-V4 `layer_output` (and `layer_input`) is `Streams`, `[batch, seq, streams, hidden_size]`, because the block carries parallel copies of the residual stream; the contributions stay `Residual`, and code written for `Residual` runs on a `Streams` value without an error and answers per stream. On Kimi-Linear `linear_attn.decays` is `ChannelGates`, `[batch, seq, heads, key_dim]`, one decay per key channel instead of the `Gates` `[batch, seq, heads]` of the other DeltaNet families. On OPT `layer_mid` is `FlatResidual`, `[tokens, hidden_size]`, the stream with batch and sequence flattened as its block holds it around the feed-forward path. ```python from nnterp.components import ChannelGates, FlatResidual, Streams from nnterp.families import deepseek_v4, kimi_linear, opt print(deepseek_v4.Layer.layer_output.layout is Streams, deepseek_v4.Layer.layer_output.dims) print(kimi_linear.LinearAttention.decays.layout is ChannelGates, kimi_linear.LinearAttention.decays.dims) print(opt.Layer.layer_mid.layout is FlatResidual, opt.Layer.layer_mid.dims) # True ('batch', 'seq', 'streams', 'hidden_size') # True ('batch', 'seq', 'heads', 'key_dim') # True ('tokens', 'hidden_size') ``` Where a recipe must run across families, check `value.layout` (or the tensor's rank) rather than assuming `Residual`. A value of your own takes a name the same way, as its return annotation ([Custom values](../../developing/custom-values.md)). # Availability Not every checkpoint has every value. A model loaded with `sdpa` never builds an attention pattern, and OPT has no `mlp` module. `model.support()` says which values this checkpoint has before any trace, from the config and the module tree alone. `None` means available, and anything else is the reason, a string. A family's page in the [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/) shows the same report for each of its public checkpoints, with the reason beside every value the checkpoint lacks. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu") support = model.support() print(support["layer_output"]) print(support["self_attn.attention_output"]) print(support["self_attn.attention_probabilities"][0]) print(len(support["self_attn.attention_probabilities"])) # None # None # read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' # 12 ``` Loaded with `attn_implementation="eager"`, every entry on GPT-2 is `None`. ## Three forms `model.support()` is the root's values plus every block value, by dotted name. A block value is `None` when every block has it, else `{layer: reason}` for the blocks that lack it, which is why the entry above is a dict with 12 keys. `model.support(layer=i)` is one block, flat, and `envoy.support()` is one envoy's own values without the prefix. ```python print(model.support(layer=0)["self_attn.attention_scores"]) print(model.layers[0].mlp.support()) # read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' # {'mlp_activation': None, 'mlp_neurons': None, 'mlp_output': None} ``` The keys come from the tree. Every child of a block that carries standard values is listed under its standard name, so a value you add through `envoys=` appears as `self_attn.` or `mlp.`. A module that some blocks lack is reported on those blocks, and a module no block has (OPT's `mlp`) has no key at all, so guard with `support.get("mlp.mlp_output")` on a checkpoint that may lack it. A hybrid reads as a per-block dict. GPT-2 has every module, which is why this block uses Qwen3.5-0.8B, where each block carries either `linear_attn` or `self_attn`. ```python hybrid = StandardizedTransformer("Qwen/Qwen3.5-0.8B", attn_implementation="eager") support = hybrid.support() print(sorted(support["self_attn.attention_output"])) print(support["self_attn.attention_output"][0]) print(sorted(support["linear_attn.attention_output"])) print([i for i, layer in enumerate(hybrid.layers) if getattr(layer, "self_attn", None) is not None]) # [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20, 21, 22] # no self_attn module on this block # [3, 7, 11, 15, 19, 23] # [3, 7, 11, 15, 19, 23] ``` No `dispatch=True` here: `support()` reads the config and the tree, which a meta-device model has. ## The reasons | reason, as it starts | values | when | | --- | --- | --- | | `read inside the eager attention forward, but this model runs 'sdpa'` | the attention interior | loaded with any `attn_implementation` but `"eager"` | | `this checkpoint sets reorder_and_upcast_attn` | the GPT-2 interior | a GPT-2 config with that flag, which takes GPT-2's own attention path | | `no self_attn module on this block`, `no linear_attn module on this block` | every value of that module | a hybrid's blocks: some other block has the module and this one does not | | `a parallel block: the attention and the MLP read the same input, so there is no stream between them` | `layer_mid` | a parallel block (Cohere, GPT-NeoX with `use_parallel_residual`): the [contribution identity](residual-stream.md#between-the-sublayers) still holds, and only `layer_mid` has no tensor | | `no act_fn module on this mlp`, `no down_proj module on this mlp` | `mlp_activation`, `mlp_neurons` | the family aliases no module to that name, such as an activation applied as a function | | `routed experts; read per-expert via the Moe component` | `mlp_activation`, `mlp_neurons` | the block's MLP is a mixture of experts ([Mixture of experts](../architectures/mixture-of-experts.md)) | | `the state after each token is materialized only by the token-by-token kernel` | `linear_attn.state`, `linear_attn.states` | a gated DeltaNet mixer before `nnterp.route_kernels(model.family, "torch")` ([Gated DeltaNet](../architectures/delta-net.md)) | | `the default family cannot tell what this sublayer adds to the stream` | `attention_output`, `mlp_output` | a `model_type` with no family ([Loading](../loading.md#a-model_type-with-no-family)) | The reasons are evaluated on the instance, so a change after the load changes the answer: `route_kernels` turns the `state` entries to `None`. `support()` predicts from the config; if a run's forward then takes a path the family does not expect, the read raises nnsight's `SourceNotAvailable` naming the missing operation ([Intermediate Operations](https://nnsight.net/features/11_source/)). ## `nnterp.Unavailable` at the read Reading or writing an unavailable value raises `nnterp.Unavailable` with the same reason, at that line, before the model runs. ```python import nnterp try: with model.trace("The Eiffel Tower is in"): pattern = model.layers[0].self_attn.attention_probabilities.save() except nnterp.Unavailable as error: print(error) # model.transformer.h.0.attn.attention_probabilities is not available: read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' ``` A module that does not exist on a block is an ordinary missing attribute: `model.layers[0].mlp` on OPT raises `AttributeError`, and `getattr(layer, "mlp", None)` is `None`, which is the form to pick blocks with outside the trace. ## `hasattr` `hasattr` treats only `AttributeError` as absence. An unavailable value raises `Unavailable` through it, and an available value read outside a trace raises nnsight's `ValueError`, so `hasattr` never answers `False`. Check `support()` outside the trace instead. ```python try: hasattr(model.layers[0].self_attn, "attention_probabilities") except nnterp.Unavailable as error: print(type(error).__name__) try: hasattr(model.layers[0], "layer_output") except ValueError as error: print(error) # Unavailable # Cannot access `model.transformer.h.0.output` outside of interleaving ``` ## Guarding code Branch on `support()` outside the trace, and the body reads only what exists. ```python model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu", attn_implementation="eager") support = model.support() has_pattern = support["self_attn.attention_probabilities"] is None has_mlp = support.get("mlp.mlp_output", "absent") is None with model.trace("The Eiffel Tower is in"): if has_pattern: pattern = model.layers[0].self_attn.attention_probabilities.save() if has_mlp: mlp = model.layers[0].mlp.mlp_output.save() print(pattern.shape, mlp.shape) # torch.Size([1, 12, 7, 7]) torch.Size([1, 7, 768]) ``` The repr lists a value whether or not this checkpoint has it. Only a value a family marks unavailable in its class body prints as `Unavailable: `, and a config-dependent reason shows in `support()` alone. Once a value is available, its shape is on [Layouts](layouts.md). # Mixture of experts On a mixture-of-experts checkpoint `layers[i].mlp` is a `Moe`, an `Mlp` that also serves the routing and the experts: `router_logits`, `expert_weights`, `expert_indices`, `expert_outputs`, `routed_output` and `shared_expert_output`. `mlp_output` stays what the block adds to the residual stream, as on every MLP. The smallest public checkpoint with a mixture that nnterp can read is ibm-granite/granite-3.0-1b-a400m-base, 32 experts and 8 slots per token, where a slot is one of a token's `top_k` routing choices. The names are the same on all 39 MoE families. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("ibm-granite/granite-3.0-1b-a400m-base", device="cuda", dtype=torch.float32) moe = model.layers[3].mlp print(moe.num_experts, moe.top_k) with model.trace("The Eiffel Tower is in the city of"): logits = moe.router_logits.save() weights = moe.expert_weights.save() indices = moe.expert_indices.save() slots = moe.expert_outputs.save() routed = moe.routed_output.save() print(tuple(logits.shape), tuple(weights.shape), tuple(indices.shape), tuple(slots.shape), tuple(routed.shape)) print(indices[0, -1].tolist()) print(torch.allclose(slots.sum(2), routed, atol=1e-5)) # 32 8 # (1, 11, 32) (1, 11, 8) (1, 11, 8) (1, 11, 8, 1024) (1, 11, 1024) # [2, 0, 10, 30, 20, 16, 18, 22] # True ``` ## The values | value | shape | what it is | | --- | --- | --- | | `router_logits` | `[batch, seq, num_experts]` | the router's logits before the scoring (a softmax, a top-k softmax or a sigmoid, named by `moe.SCORING`); writable, and the router scores what is written | | `expert_weights` | `[batch, seq, top_k]` | what multiplies each slot's expert output, after the router's normalization and scaling | | `expert_indices` | `[batch, seq, top_k]`, int64 | the expert each slot sends the token to | | `expert_outputs` | `[batch, seq, top_k, hidden_size]` | each slot's expert output times its weight; `expert_outputs.sum(2) == routed_output` | | `routed_output` | `[batch, seq, hidden_size]` | the routed experts' sum, without the shared expert | | `shared_expert_output` | `[batch, seq, hidden_size]` | the shared expert's contribution, where the mixture has one | The weights and indices are read as the experts module receives them, so a write there is what the experts compute with. `moe.num_experts` and `moe.top_k` are read off the module. On Granite the block multiplies each sublayer's output by `residual_multiplier`, 0.22 here, and `mlp_output` is `routed_output * 0.22`. The experts' values are unscaled on every family. Signatures are on the [`moe` page](../../documentation/components/moe.md) of the Documentation tab. ## Ablating and rerouting ```python prompt = "The Eiffel Tower is in the city of" with model.trace(prompt): clean = model.logits[0, -1].save() with model.trace(prompt): # expert 2 out of every slot that chose it moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == 2, 0) ablated = model.logits[0, -1].save() with model.trace(prompt): # the last token's first slot goes to expert 3 instead rerouted = moe.expert_indices.clone() rerouted[:, -1, 0] = 3 moe.expert_indices = rerouted rerouted_weights = moe.expert_weights.save() moved = model.logits[0, -1].save() with model.trace(prompt): # written before the scoring, the logits decide the routing forced = torch.full((moe.num_experts,), -10.0, device="cuda") forced[[2, 5]] = 10.0 moe.router_logits[:, -1] = forced chosen = moe.expert_indices.save() print((ablated - clean).abs().max().item(), (moved - clean).abs().max().item()) print(torch.equal(rerouted_weights, weights), chosen[0, -1, :2].tolist()) # 0.7101480960845947 1.5887203216552734 # True [5, 2] ``` Zeroing a weight removes exactly that slot's term from `routed_output`. Assign `expert_weights` in a single-prompt trace and write in place (`[:] =`) when the trace has several invokes: each invoke is served its own rows of the model's flat routing tensor as a view, so an in-place write lands on those rows, and an assignment is spliced back into the flat tensor. A rerouted slot computes `weight * expert_3(x)` with the weight the router gave the expert it replaced, and nothing is renormalized. A write to `router_logits` lands before the scoring, and the weights and indices follow it. The two forced experts share the weight, 0.5 each under Granite's top-k softmax. !!! warning "Take the baseline from the same run" An edit to one block's routing changes the stream. Later blocks can then send the token to other experts, and the rest of the forward rounds differently, which is why the clean run above is in float32, where the unedited part rounds the same way. Under two or more invokes the batch is padded and rounds differently from the single prompt (`base` below is not `clean`). Edit the routing in place there and take the baseline from an unedited invoke of the same batch: ```python with model.trace() as tracer: with tracer.invoke(prompt): base = model.logits[0, -1].save() with tracer.invoke(prompt): moe.expert_weights.masked_fill_(moe.expert_indices == 2, 0) edited = model.logits[0, -1].save() print(torch.equal(base, clean), (edited - base).abs().max().item()) # False 0.7101457118988037 ``` ## Read order Within a trace, read a mixture's values in forward order: `router_logits`, then `expert_weights` and `expert_indices`, then `expert_outputs`, `routed_output` and `mlp_output`. A read out of order raises nnsight's `OutOfOrderError`. Where `shared_expert_output` falls differs by family: after `routed_output` on DeepSeek-V3 and on Qwen2-MoE, Qwen3-Next and Qwen3.5-MoE (where it is the shared expert's output times its sigmoid gate), before the router on Hunyuan, ERNIE and Gemma-4. JetMoE's router computes the indices before the weights, and `expert_indices` comes first there. Granite has no shared expert, and `support()` says so: ```python print(model.support(layer=3)["mlp.shared_expert_output"]) # this mixture has no shared expert ``` ## Pad tokens The router routes pad tokens like any other. When counting usage over a batch, keep the real positions with `model.attention_mask`: ```python prompts = ["The Eiffel Tower is in the city of", "def add(a, b):\n return a + b"] with model.trace(prompts): mask = model.attention_mask.save() # [batch, seq], 0 at the pad positions indices = moe.expert_indices.save() real = mask.bool() usage = torch.bincount(indices[real].flatten(), minlength=moe.num_experts) print(mask[0].tolist()) print(usage.tolist()) # [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1] # [9, 8, 14, 0, 9, 1, 5, 6, 1, 4, 2, 1, 4, 13, 0, 0, 12, 0, 7, 1, 3, 3, 11, 5, 12, 2, 15, 10, 3, 3, 12, 8] ``` ## `expert_outputs` needs the grouped experts transformers runs the routed experts through one of several implementations, chosen at load like `attn_implementation`: `grouped_mm` (the default), `batched_mm`, `eager` (a Python loop over the experts) or a hub kernel. The per-slot outputs exist as one tensor only under the first two. Under the others `expert_outputs` is unavailable, and `support()` names the kwarg. The routing pair, `routed_output`, and `mlp_output` are module boundaries and read the same under every implementation. ```python eager = StandardizedTransformer("ibm-granite/granite-3.0-1b-a400m-base", device="cuda", experts_implementation="eager") print(eager.support(layer=3)["mlp.expert_outputs"]) # read inside transformers' grouped_mm / batched_mm experts forward, but this model runs 'eager'; load with experts_implementation='grouped_mm' (the default) or 'batched_mm' ``` [Expert ablation](../../tutorials/expert-ablation.md) sweeps every expert's effect on a prediction. # Gated DeltaNet On Qwen3-Next, Qwen3.5, OLMo-Hybrid and Kimi-Linear three blocks in four replace attention with a gated DeltaNet mixer, `layers[i].linear_attn`, a `LinearAttention`. It projects queries, keys and values like attention but mixes them through a recurrent state, one `[key_dim, value_dim]` matrix per head: each token decays the state by a gate, writes its key/value pair in with a strength, and the query reads against it. There is no pattern and no scores. Qwen/Qwen3.5-0.8B is the smallest public checkpoint with one. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("Qwen/Qwen3.5-0.8B", device="cpu") linear = [i for i, layer in enumerate(model.layers) if getattr(layer, "linear_attn", None) is not None] softmax = [i for i, layer in enumerate(model.layers) if getattr(layer, "self_attn", None) is not None] print(linear[:4], softmax[:4], model.config.layer_types[:4]) mix = model.layers[0].linear_attn with model.trace("The Eiffel Tower is in the city of"): q = mix.attention_queries.save() g = mix.decays.save() b = mix.betas.save() y = mix.attention_head_outputs.save() state = mix.state_output.save() out = mix.attention_output.save() print(tuple(q.shape), tuple(g.shape), tuple(y.shape), tuple(state.shape), tuple(out.shape)) print(g.dtype, g.max().item() <= 0, round(b.min().item(), 3), round(b.max().item(), 3)) # [0, 1, 2, 4] [3, 7, 11, 15] ['linear_attention', 'linear_attention', 'linear_attention', 'full_attention'] # (1, 9, 16, 128) (1, 9, 16) (1, 9, 16, 128) (1, 16, 128, 128) (1, 9, 1024) # torch.float32 True 0.044 0.996 ``` A block has `self_attn` or `linear_attn`, never both, following `config.layer_types`. Decide which blocks have which outside the trace, as above, since `getattr(envoy, name, None)` inside a trace can trip a served value. `model.support(layer=i)` reports every `self_attn` value as `no self_attn module on this block` on a linear block, and every `linear_attn` value the same way on an attention block. ## The values | value | shape | what it is | | --- | --- | --- | | `attention_output` | `[batch, seq, hidden_size]` | what the mixer adds to the residual stream | | `attention_queries`, `attention_keys` | `[batch, seq, heads, key_dim]` | what the delta rule receives, before the kernel's l2-norm and the queries' `1/sqrt(key_dim)` scale | | `attention_values` | `[batch, seq, heads, value_dim]` | the values the delta rule receives | | `decays` | `[batch, seq, heads]`, float32 | the gate: the log of how much of the state each token keeps, non-positive | | `betas` | `[batch, seq, heads]` | how strongly each token's pair is written, in (0, 1); in (0, 2) on OLMo-Hybrid | | `state_input` | `[batch, heads, key_dim, value_dim]` | the state the call starts from: `None` on a fresh prompt, a copy of the cached state on a decode step | | `state_output` | `[batch, heads, key_dim, value_dim]` | the state after the call's last token, what the next decode step starts from | | `attention_head_outputs` | `[batch, seq, heads, value_dim]` | each head's read of the state, before the gated norm and the output projection | `heads` is the mixer's `num_v_heads`, and the queries and keys are repeated up to it. On Kimi-Linear the decay is one per key channel, so `decays` is `[batch, seq, heads, key_dim]`. The sequence axis is 1 on every recurrent mixer, where softmax attention's queries put the heads first. Signatures are on the [`linear_attention`](../../documentation/components/linear_attention.md) and [`recurrent`](../../documentation/components/recurrent.md) pages of the Documentation tab. ## Kernels and `route_kernels` A recurrent mixer (a gated DeltaNet, a Mamba-1 or a Mamba-2 mixer) runs its sequence through a kernel function that the modeling module calls: here `torch_chunk_gated_delta_rule` on a prompt and `torch_recurrent_gated_delta_rule` on each decode step of `generate`, chosen by the forward's own test for one token over a cached state. Everything but `attention_output` is read at that call. nnterp reads the call that fires on this step, so the same value works in a `trace` and at every step of [`tracer.iter`](https://nnsight.net/features/4_multiple_token/), and the state hands off from one step to the next: ```python prompt = "The Eiffel Tower is in the city of" entering, leaving = [], [] # bound outside: a list first bound inside the body does not survive it with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[:3]: s = mix.state_input entering.append(s.save() if s is not None else None) leaving.append(mix.state_output.save()) print(entering[0] is None, torch.equal(leaving[0], entering[1]), torch.equal(leaving[1], entering[2])) # True True True ``` The kernels must be transformers' pure-torch ones: with `flash-linear-attention` installed the names dispatch to compiled kernels with no Python source, and every kernel value reports that, naming `route_kernels` as the fix. A submodule's value (`linear_attn.in_proj_qkv.output`) and a kernel value cannot share a trace in either order, because the kernel values look up an op at the top of the mixer's forward, and the read raises `OutOfOrderError` naming it. The queries, keys and values are views that torch refuses to edit in place under autograd (`Output 0 of Select is a view and is being modified inplace`): assign them, or edit in place under `torch.no_grad()`. `attention_head_outputs`, `decays` and `betas` take in-place edits. The chunked kernel carries the state between 64-token chunks and never materializes it per token. Until the family is routed, `state` and `states` report: ```python print(model.support(layer=0)["linear_attn.states"]) # the state after each token is materialized only by the token-by-token kernel; the chunked kernel a prompt runs through carries it between chunks. Call nnterp.route_kernels(model.family, 'torch') before tracing this layer (slower, like attn_implementation='eager') ``` `nnterp.route_kernels(family, "torch")` rebinds the family's two kernel names, in its modeling module, to transformers' pure-torch kernels. On a gated DeltaNet both go to the token-by-token loop, and a prompt then runs through it too, slower, as eager attention is. The call is process-wide, like installing a kernel, and applies to every model of the family. `"default"` restores the import-time binding. Call it before the first trace of the layer, because a forward that nnsight has already instrumented keeps the kernel it was compiled with. `family` is the family module, or `model.family` on a loaded model. ## The state after every token Routed, `states` is the state after every token of the call, `state` is one occurrence per token that `tracer.iter` reads or writes, and `state_after(t)` and `set_state_after(t, value)` are the same two operations as calls: ```python import nnterp from nnterp import StandardizedTransformer, route_kernels route_kernels(nnterp.families.qwen3_5_text, "torch") # before the first trace of the layer; process-wide model = StandardizedTransformer("Qwen/Qwen3.5-0.8B", device="cpu") mix = model.layers[0].linear_attn prompt = "The Eiffel Tower is in the city of" with model.trace(prompt): states = mix.states.save() # [batch, seq, heads, key_dim, value_dim] final = mix.state_output.save() clean = model.logits[0, -1].save() with model.trace(prompt) as tracer: for t in tracer.iter[4]: mix.state = torch.zeros_like(final) # tokens 5.. continue from an empty state for t in tracer.iter[5]: s5 = mix.state.save() edited = model.logits[0, -1].save() print(tuple(states.shape), torch.equal(states[:, -1], final)) print(round(s5.norm().item(), 2), round(states[:, 5].norm().item(), 2), round((edited - clean).abs().max().item(), 2)) # (1, 9, 16, 128, 128) True # 2.85 6.65 4.41 ``` A write at token `t` changes token `t`'s own output as well, since its query reads the state it has just written. Reads in one trace follow the forward: positions before a write are read before it, positions after it afterwards, and `states`, which reads every position, goes before any write. Among the state values the order is `state_input`, then `states`, then `state_output`; a read out of that order raises `OutOfOrderError` naming the kernel call, not the value. `states` is read-only. Under `generate` the prompt is one kernel call and each decode step another over one token, so the per-token view is an inner `tracer.iter` on step 0 and `state_output` on the later steps. !!! warning "The state is not the whole memory" A short convolution (width 4) in front of the kernel mixes each token's queries, keys and values with the three tokens before it. `set_state_after(t, zeros)` is not "forget everything before `t`": tokens `t + 1` to `t + 3` still see tokens up to `t` through the convolution. [DeltaNet state](../../tutorials/delta-net-state.md) patches and tracks the state across prompts, and [Mamba](mamba.md) and [Mamba-2](mamba-2.md) serve their mixers under the same names. # Mamba On Mamba and Falcon-Mamba every block is a norm and a Mamba-1 mixer, `layers[i].linear_attn`, a `SelectiveScan`. Jamba puts one on every block except the attention blocks that `attn_layer_period` places. Per channel of its inner width the mixer keeps a state of `state_dim` numbers and runs, for each token, `h = exp(dt * A) * h + dt * B * x` and `y = C . h + D * x`. nnterp serves the scan's tensors under the recurrent mixer's shared names: `C` reads the state as a query does, `B` says where a token writes as a key does, `x` is what it writes, `dt` is the write strength (`betas`) and `dt * A` the log decay (`decays`). state-spaces/mamba-130m-hf is the smallest public checkpoint. ```python import torch import nnterp from nnterp import StandardizedTransformer, route_kernels route_kernels(nnterp.families.mamba, "torch") # before any trace: mamba_ssm's kernels have no source and need CUDA model = StandardizedTransformer("state-spaces/mamba-130m-hf", device="cpu") mix = model.layers[0].linear_attn with model.trace("The Eiffel Tower is in the city of"): C = mix.attention_queries.save() B = mix.attention_keys.save() x = mix.attention_values.save() dt = mix.betas.save() decays = mix.decays.save() y = mix.attention_head_outputs.save() state = mix.state_output.save() out = mix.attention_output.save() print(tuple(C.shape), tuple(B.shape), tuple(x.shape), tuple(dt.shape), tuple(decays.shape)) print(tuple(y.shape), tuple(state.shape), tuple(out.shape)) # (1, 10, 1, 16) (1, 10, 1, 16) (1, 10, 1536) (1, 10, 1536) (1, 10, 1536, 16) # (1, 10, 1536) (1, 1536, 16) (1, 10, 768) ``` `route_kernels` comes first because transformers calls `mamba_ssm`'s compiled kernels when that package is installed, and they have no Python source to read inside and do not run on CPU: unrouted, a CPU trace fails inside the kernel with `Expected u.is_cuda() to be true`, even when it reads only `layer_output`. What the call does, and why a recurrent mixer's values need it, is on the [DeltaNet page](delta-net.md#kernels-and-route_kernels). On Mamba-1 the pure-torch scan is itself a token loop, and each kernel name is bound to its own pure-torch function. A Mamba block has no `self_attn` and no `mlp`: `support()` lists neither, and `layers[i].input + linear_attn.attention_output == layer_output`. Jamba's blocks are Llama's, with one mixer each. ## The values | value | shape | what it is | | --- | --- | --- | | `attention_output` | `[batch, seq, hidden_size]` | what the mixer adds to the residual stream | | `attention_queries` | `[batch, seq, 1, state_dim]` | `C`, what each token reads the state with; one group, shared by every channel | | `attention_keys` | `[batch, seq, 1, state_dim]` | `B`, what each token writes into the state along | | `attention_values` | `[batch, seq, channels]` | `x`, the input each channel writes, after the convolution and the activation | | `betas` | `[batch, seq, channels]` | `softplus(dt + dt_bias)`, the step size; derived, read-only | | `decays` | `[batch, seq, channels, state_dim]` | `dt * A`, the log decay per channel and state dimension; derived, read-only | | `state_input` | `[batch, channels, state_dim]` | a decode step's cached state, a copy; `None` on a prompt, whose scan starts from zeros | | `attention_head_outputs` | `[batch, seq, channels]` | `y = C . h + D * x`, before the `silu(z)` gate and `out_proj` | | `state_output` | `[batch, channels, state_dim]` | the state after the call's last token, as the cache receives it | `channels` is the mixer's `intermediate_size` (1536 here) and `state_dim` its `ssm_state_size` (16). The kernel takes these tensors channel-first; each value is a tokens-first view of its argument, laid back out on assignment. Unlike a DeltaNet's, `decays` has a state axis, because `A` is per channel and state dimension. `betas` and `decays` are derived from the kernel's arguments and read-only. On Falcon-Mamba and Jamba `B`, `C` and the step size pass through RMS norms before the scan, and the values are the normed ones. Signatures are on the [`selective_scan` page](../../documentation/components/selective_scan.md) of the Documentation tab. ## The state The pure-torch scan loops over the tokens, and once routed the state after every token of a prompt exists with no second kernel: `states` stacks them, `state` is one occurrence per token for `tracer.iter` to read or write, and `state_after(t)` and `set_state_after(t, value)` are the same two operations as calls: ```python prompt = "The Eiffel Tower is in the city of" with model.trace(prompt): states = mix.states.save() # [batch, seq, channels, state_dim] final = mix.state_output.save() with model.trace(prompt) as tracer: for t in tracer.iter[2]: mix.state = torch.zeros_like(final) # tokens 3.. continue from an empty state for t in tracer.iter[3]: s3 = mix.state.save() print(tuple(states.shape), torch.equal(states[:, -1], final)) print(round(s3.norm().item(), 2), round(states[:, 3].norm().item(), 2)) # (1, 10, 1536, 16) True # 3.66 10.6 ``` A write at token `t` changes token `t`'s own output too, since `y_t` reads `h_t`, and the convolution in front of the scan (width 4) still carries the three tokens before `t` into the three after it, as on a [DeltaNet](delta-net.md#the-state-after-every-token). Under `generate` the prompt runs the scan and each decode step runs `mamba_selective_state_update` over one token, the state handed from `state_output` to the next step's `state_input`: ```python with model.generate(prompt, max_new_tokens=2, do_sample=False) as tracer: for step in tracer.iter[1]: # a decode step: the state is updated first, then read entering = mix.state_input.save() leaving = mix.state_output.save() y1 = mix.attention_head_outputs.save() print(tuple(entering.shape), tuple(y1.shape)) # (1, 1536, 16) (1, 1, 1536) ``` !!! warning "Read order on a decode step" The scan computes `y` before it returns the state, but the decode kernel updates the state first and reads `y` from it. On a decode step read `state_output` before `attention_head_outputs`, as above. The other order does not raise: `state_output` binds the next step's state, and only on the last step is the loop cut short, with a `was never reached` warning that blames the loop. ## Assign, do not edit in place `C` and `B` are views that torch refuses to edit in place while autograd is on. Assign them, or edit in place under `torch.no_grad()`. `attention_values` and `attention_head_outputs` take in-place edits either way, and a write to `state_output` changes what the next decode step starts from, not this call's output. ```python with model.trace(prompt): mix.attention_queries = mix.attention_queries * 0 # assign: C is a view of the kernel's argument y_noC = mix.attention_head_outputs.save() with model.trace(prompt): with torch.no_grad(): mix.attention_queries[:, -1] = 0 # or edit in place without autograd y_noC_last = mix.attention_head_outputs.save() print(torch.allclose(y_noC[:, -1], y_noC_last[:, -1])) try: with model.trace(prompt): mix.attention_queries[:, -1] = 0 except RuntimeError as e: print(str(e)[:60]) try: with model.trace(prompt): mix.betas = mix.betas * 0 except AttributeError as e: print(e) # True # Output 0 of Select is a view and is being modified inplace. # betas is derived and read-only ``` Mamba-2's mixer has the same names and a different state: [Mamba-2](mamba-2.md). # Mamba-2 On Mamba-2, Nemotron-H, Bamba and Falcon-H1 the recurrent mixer `layers[i].linear_attn` is a `StateSpace`, a Mamba-2 (SSD) mixer: linear attention with one scalar decay per head. Per head, with the state a `head_dim` by `state_dim` matrix, each token does `h = exp(dt * A) * h + dt * x B^T` and `y = h C + D * x`. The shared names mean what they mean on a [gated DeltaNet](delta-net.md): `C` is `attention_queries`, `B` is `attention_keys`, `x` is `attention_values`, `dt` is `betas`, `A * dt` is `decays` and `y` is `attention_head_outputs`. AntonV/mamba2-130m-hf is the smallest public checkpoint. On it every block is the mixer alone. Nemotron-H and Bamba mix such blocks with attention and MLP blocks, and Falcon-H1 runs a mixer and attention side by side on every block. ```python import torch import nnterp from nnterp import StandardizedTransformer, route_kernels route_kernels(nnterp.families.mamba2, "torch") # before any trace: mamba_ssm's kernels have no source and need CUDA model = StandardizedTransformer("AntonV/mamba2-130m-hf", device="cpu") mix = model.layers[0].linear_attn with model.trace("The Eiffel Tower is in the city of"): C = mix.attention_queries.save() B = mix.attention_keys.save() x = mix.attention_values.save() dt = mix.betas.save() decays = mix.decays.save() y = mix.attention_head_outputs.save() state = mix.state_output.save() out = mix.attention_output.save() print(tuple(C.shape), tuple(B.shape), tuple(x.shape), tuple(dt.shape), tuple(decays.shape)) print(tuple(y.shape), tuple(state.shape), tuple(out.shape)) print(decays.dtype, decays.max().item() <= 0, dt.min().item() > 0) # (1, 10, 1, 128) (1, 10, 1, 128) (1, 10, 24, 64) (1, 10, 24) (1, 10, 24) # (1, 10, 24, 64) (1, 24, 128, 64) (1, 10, 768) # torch.float32 True True ``` `route_kernels` comes first, as on the [DeltaNet page](delta-net.md#kernels-and-route_kernels): with `mamba_ssm` installed transformers calls its compiled kernels, which have no source to read and do not run on CPU. Unrouted, a CPU trace fails inside the kernel (`Pointer argument cannot be accessed from Triton` when a GPU is visible). On Mamba-2 only the two scan kernels are rerouted, and a prompt keeps the chunked scan. ## The values | value | SSD | shape | what it is | | --- | --- | --- | --- | | `attention_output` | | `[batch, seq, hidden_size]` | what the mixer adds to the residual stream | | `attention_queries` | `C` | `[batch, seq, groups, state_dim]` | what reads the state; one per group of heads | | `attention_keys` | `B` | `[batch, seq, groups, state_dim]` | where each token writes into the state; one per group | | `attention_values` | `x` | `[batch, seq, heads, head_dim]` | what each token writes | | `betas` | `dt` | `[batch, seq, heads]` | `softplus(dt + dt_bias)`, clamped to `time_step_limit` on a prompt: the write strength | | `decays` | `A * dt` | `[batch, seq, heads]`, float32 | the log of how much of the state each token keeps, non-positive | | `state_input` | `h` in | `[batch, heads, state_dim, head_dim]` | the state the call starts from: `None` on a fresh prompt, a copy of the cached state on a decode step | | `state_output` | `h` out | `[batch, heads, state_dim, head_dim]` | the state after the call's last token | | `attention_head_outputs` | `y` | `[batch, seq, heads, head_dim]` | each head's read of the state plus the `D` skip, before the gated norm and `out_proj` | `groups` is the mixer's `n_groups`, 1 here (the heads of a group share `B` and `C`), `heads` its `num_heads` (24), `head_dim` 64 and `state_dim` its `ssm_state_size` (128). The state is served key side first, the layout every recurrent mixer's state shares and the transpose of the cache's buffer. An assignment is transposed back. Assign `attention_queries`, `attention_keys` and `attention_values` rather than editing them in place, as on the other recurrent mixers. Signatures are on the [`state_space` page](../../documentation/components/state_space.md) of the Documentation tab. ## `betas` and `decays` are one argument Both are the kernel's `dt` seen through `dt_bias`, the softplus and `A`. An assignment is carried back into `dt`, and the kernel computes the gate that was written. An in-place edit (`mix.betas[:, t] = 0`) does not reach `dt`. ```python prompt = "The Eiffel Tower is in the city of" with model.trace(prompt): mix.decays = mix.decays * 0.5 # half the log decay is half of dt: half-strength writes too halved = mix.betas.save() with model.trace(prompt): betas = mix.betas.clone() betas[:, 3] = 0 # token 3 writes nothing and decays nothing mix.betas = betas skipped = model.logits[0, -1].save() with model.trace(prompt): mix.betas = mix.betas # the round trip through the softplus's inverse rounds dt written_back = model.logits[0, -1].save() with model.trace(prompt): clean = model.logits[0, -1].save() print(torch.allclose(halved, dt * 0.5, atol=1e-6)) print((skipped - clean).abs().max().item() > 0, (written_back - clean).abs().max().item()) # True # True 0.000125885009765625 ``` !!! warning "Writing `decays` rewrites `betas`" There is no way to change the decay alone. `decays * 0.5` halves `dt`: each token keeps more of the state and writes at half strength. `decays = 0` ("keep everything") sets `dt` to zero, which writes nothing, and the state stays at zero through the whole prompt. `betas[:, t] = 0` is the exact "skip this token": no write and `exp(0) = 1`, no decay, and the state after token `t` equals the state after `t - 1`. Nemotron-H's scan clamps `dt` to `time_step_min` first, and there the state still moves. A written `betas` must be positive or zero, and on a prompt the kernel clamps it to `time_step_limit` again. Writing the same values back is not a no-op: the round trip rounds `dt`, by 1.3e-4 in the last logit here. Compare an edit against an unedited run, not against a write-back. ## The state after every token The chunk scan computes the state at every chunk boundary in one tensor, and between boundaries there is nothing to read: ```python from nnterp import chunk_per_token print(model.support(layer=0)["linear_attn.states"]) chunk_per_token(model) # this model only; undo with chunk_per_token(model, False) with model.trace(prompt): states = mix.states.save() # [batch, seq, heads, state_dim, head_dim] final = mix.state_output.save() with model.trace(prompt): betas = mix.betas.clone() betas[:, 3] = 0 mix.betas = betas held = mix.states.save() print(tuple(states.shape), torch.equal(states[:, -1], final), torch.equal(held[:, 3], held[:, 2])) print(model.support(layer=0)["linear_attn.set_state_after"]) chunk_per_token(model, False) # the chunk scan materializes the state only at chunk boundaries, every 256 tokens (chunk_size=256); call nnterp.chunk_per_token(model) to set every mixer's chunk_size to 1, so every token is a boundary (slower on long prompts) # (1, 10, 24, 128, 64) True True # the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts ``` `nnterp.chunk_per_token(model)` sets every Mamba-2 mixer's `chunk_size` to 1. Every token is then a boundary, and `states` and `state_after(t)` read the state after every token of the call. The call touches this model only, because the forward reads the attribute on every call, where `route_kernels` rebinds a family's kernels for the process. `chunk_per_token(model, False)` restores the chunk size each mixer was built with. The recurrence between chunks is quadratic in their number, and a long prompt runs markedly slower. The same scan in another order also moves the logits by rounding. `states` is a read-only copy. `state` and `set_state_after` stay unavailable either way, with the reasons `support()` prints: the scan has no per-token occurrence for `tracer.iter` to walk, as a DeltaNet has, and it computes every boundary state in one cumulative step from the initial state, in which a state written at one token does not flow into the next. To change where a call starts, assign `state_input`. ## The state between steps A prompt runs `mamba2_chunk_scan` and each decode step runs `mamba2_selective_state_update`, which takes its arguments in other places and has no sequence axis. The values select from the kernel that fires, and a decode step's tensors are served with a sequence axis of 1. ```python entering, leaving = [], [] with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[:3]: s = mix.state_input entering.append(s.save() if s is not None else None) leaving.append(mix.state_output.save()) print(entering[0] is None, torch.equal(leaving[0], entering[1]), torch.equal(leaving[1], entering[2])) # True True True ``` The update writes the new state into the cache and returns only `y`, which is where a decode step's `state_output` is read. Read it before `attention_head_outputs` on a decode step, the trap the [Mamba page](mamba.md#the-state) describes. On Falcon-H1-0.5B and the Falcon-H1-Tiny checkpoints, which set `mamba_rms_norm` off, a decode step's `attention_head_outputs` is gated by `silu(z)` where a prompt's is not. # Vision-language An image-text-to-text checkpoint is a text model plus a vision tower, a projector, and a step that scatters the projected image features into the text stream at the image tokens. Load it with `task="image-text-to-text"` and pass the image to the trace. The text model keeps its standard names and values, `model.vision` and `model.projector` name the vision side, and two values of the tower say where the image enters the text model: `vision.image_token_mask` and `vision.image_features`. GPT-2 has no tower, so this page runs `llava-hf/llava-interleave-qwen-0.5b-hf`, a SigLIP tower on a Qwen2 text model. ```python import torch from PIL import Image from nnterp import StandardizedTransformer model = StandardizedTransformer("llava-hf/llava-interleave-qwen-0.5b-hf", task="image-text-to-text", device="cpu", attn_implementation="eager") image = Image.new("RGB", (64, 64), "red") messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Answer in one word: what color is this?"}]}] prompt = model.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) with model.trace(prompt, images=[image]): mask = model.vision.image_token_mask.save() # first: it comes off the inputs pattern = model.vision.layers[0].self_attn.attention_probabilities.save() patches = model.vision.layers[0].layer_output.save() projected = model.projector.output.save() features = model.vision.image_features.save() first = model.layers[0].input.save() logits = model.logits.save() print(tuple(mask.shape), mask.sum().item()) # (1, 748) 729 print(tuple(pattern.shape), tuple(patches.shape), tuple(projected.shape), tuple(features.shape)) # (1, 16, 729, 729) (1, 729, 1152) (1, 729, 1024) (729, 1024) print(torch.equal(first[mask], features)) # True ``` The processor's chat template puts the image placeholder in the prompt, and the processor expands it to one token per projected patch, 729 here. A tower block's values have the `Patches` layout, `[images, patches, vision_hidden_size]`, and its attention interior has the same names as a text block's ([Attention](../values/attention.md)). The tower's sizes are on `model.vision` (`num_layers` 26, `vision_hidden_size` 1152), separate from the text model's (`model.num_layers` 24, `model.hidden_size` 1024). `image_token_mask` is `[batch, seq]` bool. `image_features` is `[image_tokens, hidden_size]`, the tensor the wrapper writes into the token embeddings at the image tokens, in row-major order, so the equality above is exact and holds on every wrapper. It is where to ablate, patch or steer the image as the text model sees it. ```python with model.trace(prompt, images=[image]): model.vision.image_features[:] = 0 ablated = model.logits.save() print(model.probs_to_dict(logits[0, -1].softmax(-1), k=2), model.probs_to_dict(ablated[0, -1].softmax(-1), k=2)) # {'Red': 0.984375, 'Orange': 0.004852294921875} {'Answer': 0.56640625, 'The': 0.2080078125} ``` What a row of a tower value is differs by encoder. SigLIP has one row per image, in raster order. CLIP, Llava 1.5's tower, puts the CLS token first on the patch axis, and Llama 4's encoder puts it last. Pixtral, the Qwen ViT and RADIO pack every image's patches into one row, `[1, patches, vision_hidden_size]`, and Gemma 4 pads the patch axis to a fixed row count. On the Qwen ViT and RADIO the attention runs once per image, so their `attention_probabilities` and `attention_scores` are `Unavailable` and the queries and keys are served whole. ## Image rows and text rows Every value of the text model is over the whole sequence, image tokens and text tokens alike. The mask splits them, and boolean indexing flattens the batch: `out[mask]` is `[image_tokens, hidden_size]` in the order of `image_features`, and `out[~mask]` the text rows. To keep the batch shape, mask instead of indexing (`out.masked_fill(~mask[..., None], 0)`). ```python with model.trace(prompt, images=[image]): mask = model.vision.image_token_mask.save() out = model.layers[6].layer_output.save() print(tuple(out[mask].shape), tuple(out[~mask].shape)) # (729, 1024) (19, 1024) with model.trace(prompt, images=[image]): positions = model.vision.image_token_mask model.layers[6].layer_output[positions] = 0 # the image positions only edited = model.logits.save() print(model.probs_to_dict(edited[0, -1].softmax(-1), k=2)) # {'Black': 0.169921875, 'White': 0.0966796875} ``` The in-place write through the mask is a one-sided edit, which leaves the text rows untouched. [Image pathway](../../tutorials/image-pathway.md) builds the ablation and patching recipes on it. ## The projector `model.projector` is the last module before the scatter, and `projector.input` is what it receives, which is not `tower_output` on every host. Here the projector reads the last block's `layer_output` before the tower's final norm (`vision.norm`), so a write to `tower_output` never reaches the text model. Llava 1.5 reads block `-2` without its CLS token, and Gemma 3 pools `tower_output`. On the Qwen ViT the projector is the tower's merger and runs inside the tower, so read `projector.input` and `projector.output` before `tower_output`. ```python with model.trace(prompt, images=[image]): last_block = model.vision.layers[-1].layer_output.save() tower = model.vision.tower_output.save() received = model.projector.input.save() print(torch.equal(received, last_block), torch.equal(received, tower)) # True False ``` ## Read order and one image per trace Reads follow the forward: `image_token_mask` first, then the tower's values in order (a block's attention interior before its `layer_output`), the projector's, `image_features`, and then the text model's. A read out of that order raises `OutOfOrderError`, as in any trace. A trace carries one image-carrying invoke, and a second raises `NotImplementedError: Can't batch these inputs`. Several images go in that one invoke as a list, and each becomes a row of the tower's values: ```python blue = Image.new("RGB", (64, 64), "blue") messages = [{"role": "user", "content": [{"type": "image"}, {"type": "image"}, {"type": "text", "text": "Compare them."}]}] two = model.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) with model.trace(two, images=[image, blue]): mask = model.vision.image_token_mask.save() patches = model.vision.layers[0].layer_output.save() features = model.vision.image_features.save() print(mask.sum().item(), tuple(patches.shape), tuple(features.shape)) # 1458 (2, 729, 1152) (1458, 1024) ``` Under `generate` the tower runs once, for the prompt: `image_token_mask` is `[batch, prompt_len]` on step 0 and `[batch, 1]`, all false, on every decode step, and the tower's values and `image_features` occur on step 0 only, so read them in `tracer.iter[0]` ([Generation](../generation.md)). `model.support()` lists the tower's values under `vision.`, such as `"vision.image_features"` and `"vision.self_attn.attention_probabilities"`, and `model.vision.support()` lists them alone. A load without `task="image-text-to-text"` builds the wrapper with no processor, so no image can reach it: `model.vision.support()` is then `{}` and reading an image value raises `Unavailable` saying so. # Methods `StandardizedTransformer`'s methods are written against the standard values, so each runs unchanged on every family. `skip_layers` and `steer` edit the residual stream inside a trace, `project_on_vocab` decodes a hidden state through the model's own head, and `get_topk_closest_tokens` and `probs_to_dict` turn the result into tokens. The signatures are in the [Documentation tab](../documentation/standardized.md#nnterp.standardized.Standardized). ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu") prompt = "The Eiffel Tower is in the city of" with model.trace(prompt): before = model.layers[3].layer_output.save() model.skip_layers(4, 7) # blocks 4 to 7 do not run after = model.layers[7].layer_output.save() logits = model.logits.save() print(torch.equal(before, after)) # True print(model.probs_to_dict(logits[0, -1].softmax(-1), k=3)) # {' the': 0.057030126452445984, ' London': 0.03870350494980812, ' New': 0.024094339460134506} ``` Without the skip the model puts `' Paris'` first, at 0.070. ## skip_layers `skip_layers` hands the stream entering block `start` to block `end + 1`. Both indices are inclusive, and negative ones count from the end. `skip_with`, when given, replaces that stream with a tensor of your own, `[batch, seq, hidden_size]`. A skipped block's submodules never run, so nothing under its `self_attn` or `mlp` can be read ([nnsight: skip execution](https://nnsight.net/features/13_skip/)). On a family with tuple-returning blocks, such as GPT-J, BLOOM or Falcon, the stream is packed the way the block would have returned it. The call reads `layers[start].input` itself, so a read of that location after it is out of order: ```python try: with model.trace(prompt): model.skip_layers(4, 7) entering = model.layers[4].layer_input.save() except Exception as error: print(type(error).__name__, error) # OutOfOrderError 'model.transformer.h.4.input.i0' was requested but the model already ran past it ``` Read `layers[start - 1].layer_output` before the call instead, as the first block does. ## steer `steer` adds `factor * vector`, with `factor` 1.0 by default, to `layer_output` of each block in `layers` (an int, or an ascending list), in place, so the add reaches the model. `vector` is `[hidden_size]` or anything broadcastable to the selected slice, and is cast to the stream's dtype and device. `token_positions` (an int, a list or a slice) and `batch_index` restrict the add; both default to everything. ```python torch.manual_seed(0) direction = torch.randn(model.hidden_size) with model.trace(prompt): clean = model.layers[6].layer_output.save() with model.trace(prompt): model.steer(6, direction, factor=3.0, token_positions=-1) steered = model.layers[6].layer_output.save() logits = model.logits.save() print(torch.allclose(steered[:, -1] - clean[:, -1], 3 * direction, atol=1e-5), torch.equal(steered[:, :-1], clean[:, :-1])) # True True print(model.probs_to_dict(logits[0, -1].softmax(-1), k=3)) # {' Chicago': 0.024904731661081314, ' New': 0.020505620166659355, ' Los': 0.019833678379654884} ``` The add lands on the stream leaving a block. To steer the stream entering block `i`, steer block `i - 1`. Scale the factor against the stream's norm at that block, which grows with depth and differs between models by orders of magnitude. Here the last position's stream at block 6 has norm 92.6 and the random direction 28.5, which makes `factor=3.0` about the stream's own size. [Steering](../tutorials/steering.md) derives a direction worth adding. ## project_on_vocab `project_on_vocab` returns logits for a residual-stream tensor: the final norm, `lm_head`, then whatever the model does to the head's output before it returns `logits`. Inside a trace it takes a live value, outside a saved one, and on the last block's `layer_output` it equals `logits` exactly. ```python with model.trace(prompt): mid = model.layers[6].layer_output.save() last = model.layers[-1].layer_output.save() logits = model.logits.save() print(torch.equal(model.project_on_vocab(last), logits)) # True print(model.get_topk_closest_tokens(mid[0, -1], k=3)) # [{' East': 0.05666778236627579, ' England': 0.04518940672278404, ' southwest': 0.038670022040605545}] ``` The step after the head is where families differ, and it is why a hand-rolled `lm_head(norm(hidden))` is not always the model's prediction. The plain case is the softcap when the text config sets `final_logit_softcapping`: on `google/gemma-2-2b` (cap 30.0) `project_on_vocab(last)` equals `logits`, and `lm_head(norm(last))` differs from them by up to 50.75 on this prompt. A family where the model does something else after the head defines its own `project_on_vocab`: Cohere multiplies by `logit_scale`, Granite divides by `logits_scaling`, and DeepSeek-V4 first collapses its parallel streams the way the model does. A logit lens over every block is a loop inside one trace ([Logit lens](../tutorials/logit-lens.md) is the full recipe). A list bound in the block needs `nnsight.save`; binding it before the trace, as the Logit lens tutorial does, is the equivalent idiom: ```python import nnsight with model.trace(prompt): lens = nnsight.save([model.project_on_vocab(layer.layer_output)[0, -1] for layer in model.layers]) print([next(iter(model.probs_to_dict(row.softmax(-1), k=1))) for row in lens]) # [' the', ' the', ' the', ' the', ' the', ' the', ' East', ' Ing', ' Rome', ' London', ' Paris', ' Paris'] ``` ## get_topk_closest_tokens and probs_to_dict `get_topk_closest_tokens` is `project_on_vocab`, softmax, then the `k` most likely tokens (5 by default) for each row of `hidden.reshape(-1, hidden_size)`, one `{token: probability}` dict per position. `probs_to_dict` does the last step alone, on one `[vocab_size]` distribution, with the same `k`. ```python print(model.get_topk_closest_tokens(last[0], k=1)[:3]) # the first three positions # [{'\n': 0.013599064201116562}, {'-': 0.04670916870236397}, {'el': 0.9958784580230713}] with model.trace(prompt): probs = model.next_token_probs.save() print(model.probs_to_dict(probs[0], k=5)) # {' Paris': 0.06996913999319077, ' London': 0.0583544597029686, ' New': 0.03101949580013752, ' Amsterdam': 0.029201989993453026, ' Berlin': 0.023510510101914406} ``` Each dict holds exactly `k` entries, most likely first. Where two of them decode to the same text (partial UTF-8 byte tokens all decode to the replacement character), those entries are keyed by the tokenizer's raw vocabulary token instead. Both methods call `.item()` and `tokenizer.decode`, so their result is plain Python. Inside a trace it has to be kept with `nnsight.save`, and calling them outside, on saved tensors, is the plain form. # Generation Under `model.generate(prompt, max_new_tokens=N)` the model runs one forward call for the prompt and one per new token ([nnsight: multiple token generation](https://nnsight.net/features/4_multiple_token/)), and every standard value is served once per call: `[batch, prompt_len, ...]` on the prompt call, `[batch, 1, ...]` on a decode step. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu", attn_implementation="eager") prompt = "The Eiffel Tower is in the city of" resid, probs = [], [] # bound before the block: a list bound inside it does not survive it with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[:3]: # steps 0, 1, 2 resid.append(model.layers[6].layer_output.save()) probs.append(model.next_token_probs.save()) ids = tracer.result.save() print([tuple(r.shape) for r in resid]) # [(1, 10, 768), (1, 1, 768), (1, 1, 768)] print([p.argmax(-1).item() for p in probs], ids[0, -3:].tolist(), repr(model.tokenizer.decode(ids[0, -3:]))) # [6342, 11, 290] [6342, 11, 290] ' Paris, and' ``` `next_token_probs` is the distribution each call samples from, so under greedy decoding its argmax is the token the step appended. Decoding follows the checkpoint's `generation_config`; pass `do_sample=False` for a deterministic run. ## A decode step's values | value | prompt call (step 0) | decode step | | --- | --- | --- | | `model.input_ids` | `[batch, prompt_len]` | `[batch, 1]`, the token the previous step produced | | `layers[i].layer_output`, `attention_output`, `mlp_output`, `token_embeddings` | `[batch, prompt_len, hidden_size]` | `[batch, 1, hidden_size]` | | `self_attn.attention_queries`, `attention_head_outputs` | `prompt_len` on the sequence axis | 1 on the sequence axis | | `self_attn.attention_keys`, `attention_values` | `[batch, num_kv_heads, prompt_len, head_dim]` | `[batch, num_kv_heads, prompt_len + step, head_dim]`, the cache | | `self_attn.attention_scores`, `attention_probabilities` | `[batch, num_heads, prompt_len, prompt_len]` | `[batch, num_heads, 1, prompt_len + step]` | | `model.logits` | `[batch, prompt_len, vocab_size]` | `[batch, 1, vocab_size]` | | `model.next_token_probs` | `[batch, vocab_size]` | `[batch, vocab_size]` | ```python with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[1]: # the first decode step ids_in = model.input_ids.save() keys = model.layers[6].self_attn.attention_keys.save() pattern = model.layers[6].self_attn.attention_probabilities.save() out = model.layers[6].layer_output.save() ids = tracer.result.save() print(ids_in.tolist(), tuple(keys.shape), tuple(pattern.shape), tuple(out.shape)) # [[6342]] (1, 12, 11, 64) (1, 12, 1, 11) (1, 1, 768) ``` The keys and values grow with the cache and the queries do not, so `pattern[0, h, 0]` is what head `h` of the new token attends to across the whole context so far, and a head-wise edit of `attention_keys` on a decode step touches the whole context. A read outside any loop binds to step 0, the prompt call. A write at a step reaches that step and every later one: ```python with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[1]: model.layers[6].layer_output[:] = 0 ids = tracer.result.save() print(repr(model.tokenizer.decode(ids[0, -3:]))) # ' Paris in the' ``` Step 0's token stays `' Paris'`; the two after it change. ## A read out of order inside the loop Within one step's body the reads follow the forward, as in a trace: `input_ids` first, a block's interior values before its `attention_output`, that before the block's `layer_output`, and `logits` and `next_token_probs` last. What differs under `generate` is how a misplaced read fails. !!! warning "A shifted list, not an error" In a plain trace, reading `attention_probabilities` after the same block's `attention_output` raises `OutOfOrderError`. Inside `tracer.iter` the read binds the value's next occurrence, which is the next step's, so the list comes back shifted by one step with no error. Only the read on the last step has no next occurrence. It cuts the loop short with nnsight's `was never reached` warning, which blames the loop bound, and the statements after the loop (here `tracer.result`) never run. A saved name that is missing, or a list that looks one step off, is a read-order error. ```python outs, patterns = [], [] with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[:3]: outs.append(model.layers[6].self_attn.attention_output.save()) patterns.append(model.layers[6].self_attn.attention_probabilities.save()) # after attention_output: out of order print(len(outs), len(patterns), [tuple(p.shape) for p in patterns]) # 3 2 [(1, 12, 1, 11), (1, 12, 1, 12)] ``` The key axis says which steps these are: 11 and 12 keys are steps 1 and 2, not 0 and 1. Check a saved shape against its step before trusting a per-step list. A recurrent mixer has its own order inside a decode step: on Mamba-1 and Mamba-2, `state_output` comes before `attention_head_outputs` ([Mamba](architectures/mamba.md), [Mamba-2](architectures/mamba-2.md)), and on a DeltaNet hybrid `linear_attn.state_input` is the previous step's `state_output` ([Gated DeltaNet](architectures/delta-net.md)). # Helpers `nnterp.nnsight_utils` collects one token position's activation at chosen blocks over many prompts, and `nnterp.prompt_utils` measures the probability mass on named target tokens over many prompts. Both read standard values (`layer_output`, `next_token_probs`), so one call runs on every family. The signatures are in the Documentation tab: [nnsight_utils](../documentation/nnsight_utils.md), [prompt_utils](../documentation/prompt_utils.md). ```python from nnterp import StandardizedTransformer from nnterp.nnsight_utils import get_token_activations model = StandardizedTransformer("openai-community/gpt2", device="cpu", tokenizer_kwargs={"padding_side": "left"}) prompts = ["The Eiffel Tower is in", "Hello"] acts = get_token_activations(model, prompts, layers=[4, 8]) print(acts.shape, acts.device) # torch.Size([2, 2, 768]) cpu ``` `acts[i, j]` is `model.layers[layers[i]].layer_output[j, -1]`: the stream leaving block `layers[i]` at the last token of prompt `j`. The layer axis follows `layers`, not the block numbers, and the result is on the CPU. The trace stops after the last requested block, so later blocks do not run. ## The token index and the padding side `idx` is the token position, `-1` by default. A negative index needs left padding and a positive one right padding, so that the index names the same token in every row. `0` needs neither. The check runs against `model.tokenizer.padding_side` before anything else: ```python try: get_token_activations(model, prompts, idx=2) except ValueError as error: print(error) # a positive token index needs right padding, and the tokenizer pads 'left' ``` nnsight already sets a decoder-only model's tokenizer to pad on the left, and the load above only makes that explicit with `tokenizer_kwargs`. The same rule covers every function here that batches prompts: `next_token_probs` is each row's last position, which is each prompt's own last token only under left padding. ## Another value, many prompts `get_activations(model, layer)` picks the value read at each block; any value with the sequence axis in position 1 works, because the helper indexes `[:, idx]` on it: ```python attn = get_token_activations(model, prompts, layers=[4], get_activations=lambda m, i: m.layers[i].self_attn.attention_output) print(attn.shape) # torch.Size([1, 2, 768]) ``` `collect_token_activations_batched` runs `get_token_activations` over batches and concatenates along the prompt axis. `tqdm=` takes a progress-bar factory to wrap the batch loop. With `remote=True` the batches go through `collect_last_token_activations_session`, one `model.session`, so a remote run is one request ([Remote](remote.md)). ```python from nnterp.nnsight_utils import collect_token_activations_batched many = [f"Prompt number {i}" for i in range(10)] acts = collect_token_activations_batched(model, many, batch_size=4, layers=[0, 11]) print(acts.shape) # torch.Size([2, 10, 768]) ``` `compute_next_token_probs` is the distribution at the end, `[num_prompts, vocab_size]` on the CPU: ```python from nnterp.nnsight_utils import compute_next_token_probs probs = compute_next_token_probs(model, prompts) print(probs.shape, model.probs_to_dict(probs[0], k=3)) # torch.Size([2, 50257]) {' the': 0.2401452511548996, ' a': 0.0862361341714859, ' fact': 0.02262996695935726} ``` The `tracer=` form of `get_token_activations` reads inside a trace you have open instead of opening one; the result is then a live tensor on the model's device, and yours to `.save()`. ## Target token mass over many prompts A `Prompt` pairs a text with named sets of target token ids, and `run_prompts` runs many of them in batches and returns each target's mass per prompt: ```python from nnterp.prompt_utils import Prompt, run_prompts targets = [ Prompt.from_strings("The capital of France is", {"correct": "Paris", "wrong": ["London", "Berlin"]}, model), Prompt.from_strings("The capital of Germany is", {"correct": "Berlin", "wrong": ["London", "Paris"]}, model), ] print(targets[0].target_tokens, targets[0].has_no_collisions()) # {'correct': [40313, 6342], 'wrong': [23421, 3576, 24814, 11307]} True mass = run_prompts(model, targets, batch_size=32) print(mass["correct"].shape, mass["correct"][:, 0].tolist()) # torch.Size([2, 1]) [0.03224833309650421, 0.05272921919822693] ``` `from_strings` turns each word into token ids with `get_first_tokens`: the first token of `"Paris"` and of `" Paris"`, tokenized with `model.add_prefix_false_tokenizer` so that the two differ. A target's mass is the sum over its ids, and `has_no_collisions` says whether two targets share one. The result is `[num_prompts, layers]`, the next-token distribution counting as one layer. Every prompt must name the same targets. When you pick target ids yourself instead, take `model.tokenizer(" Paris", add_special_tokens=False).input_ids` and check that it is one token: `tokenizer.encode(" Paris")[0]` is BOS on Llama and Gemma, and some vocabularies split the word, or make `"Paris"` and `" Paris"` the same token. `get_probs_func` gives the layer axis a meaning of your own. The logit lens at chosen blocks, through `project_on_vocab`: ```python import torch def lens_probs(model, batch, remote=False, layers=None): with model.trace(batch, remote=remote): rows = [model.project_on_vocab(model.layers[i].layer_output)[:, -1].softmax(-1) for i in layers] out = torch.stack(rows, dim=1).cpu().save() # [batch, len(layers), vocab_size] return out mass = run_prompts(model, targets, batch_size=32, get_probs_func=lens_probs, func_kwargs={"layers": [6, 9, 11]}) print(mass["correct"].shape, [round(p, 4) for p in mass["correct"][0].tolist()]) # torch.Size([2, 3]) [0.0001, 0.1821, 0.0322] ``` The function receives `remote=` as a keyword and must pass it to `model.trace`. `func_kwargs` carries the rest. # Remote `remote=True` on a `StandardizedTransformer` trace is nnsight's remote trace ([nnsight: remote execution](https://nnsight.net/features/15_remote_execution/)) with nnterp's names and values in the block. The server runs a plain `TransformersModel` of the same checkpoint. Your block is re-run there against your envoy tree, which carries the family's aliases and envoy classes by reference, so nnterp must be installed on the server, at the same version as yours. nnterp is never shipped by value. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") # no dispatch: the architecture on the meta device, no weights print(model.to_model_key()) # nnsight.modeling.transformers.TransformersModel:{"repo_id": "openai-community/gpt2", "revision": null, "task": "text-generation"} ``` The key that a request carries names `TransformersModel`, because that is what a server deploys, so every checkpoint a server serves is reachable through a `StandardizedTransformer` without the server deploying anything nnterp-specific. The block itself is what you would run locally, and every `.save()` is downloaded, so slice inside it: ```python # not run for this page: needs an NDIF key and a deployed model from nnsight import CONFIG from nnterp import StandardizedTransformer CONFIG.set_default_api_key("YOUR_KEY") model = StandardizedTransformer("meta-llama/Llama-3.1-70B", attn_implementation="eager") with model.trace("The Eiffel Tower is in the city of", remote=True): resid = model.layers[40].layer_output[:, -1].save() # [batch, hidden_size] pattern = model.layers[40].self_attn.attention_probabilities[0, :, -1].save() # [num_heads, key] model.layers[41].mlp.mlp_output[:] = 0 probs = model.next_token_probs.save() # [batch, vocab_size] ``` ## The request's contents The request carries the block's source, the names it saves, and the envoy tree by reference: `model.layers[i].self_attn` resolves on the server to the same module it does on your machine, and a value such as `layer_output` is a descriptor on the family's envoy classes, which the server imports from its own nnterp. A different version there can define a value differently or spell a source-located operation another way, and the error is whatever the mismatch produces server-side. Shipping nnterp by value does not work around that: `nnsight.register(nnterp)` fails with `TypeError: cannot pickle 'EProperty' object`. Register your own helper module if the block calls one. An installed package travels by reference either way. ## A dry run without a server `remote="local"` serializes the block, deserializes it with your local modules hidden, and runs it against a dispatched model, with no server and no key. A block that survives it, aliases and values included, is one that ships: ```python model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu", attn_implementation="eager") with model.trace("The Eiffel Tower is in the city of", remote="local"): pattern = model.layers[3].self_attn.attention_probabilities[0, :, -1].save() model.layers[7].mlp.mlp_output[:] = 0 probs = model.next_token_probs.save() print(tuple(pattern.shape), model.probs_to_dict(probs[0], k=2)) # (12, 10) {' Paris': 0.09203002601861954, ' London': 0.0368831530213356} ``` ## Decided by the deployment Availability is checked against your model's config, not the deployment's: `attn_implementation="eager"` at load declares the attention interior available, but whether the deployed forward runs eager is the server's choice, and on an `sdpa` deployment the read fails server-side rather than with `Unavailable` on the client. Source-located values, such as `attention_probabilities`, `attention_queries` and a mixture's `expert_outputs`, resolve against the server's transformers release, and a release that spells the operation differently fails the same way. A recurrent mixer's per-token state needs `nnterp.route_kernels` in the process that runs the model, which a request does not reach, so treat `state` and `states` as unavailable remotely ([Gated DeltaNet](architectures/delta-net.md)). The helpers batch their remote loops into one `model.session`: `collect_token_activations_batched(..., remote=True)` and `run_prompts(..., remote=True)` are one request each ([Helpers](helpers.md)). # Tutorials - [Logit lens](logit-lens.md): where in depth a prediction emerges. - [Attention patterns](attention-patterns.md): scoring attention heads. - [Contribution decomposition](contribution-decomposition.md): direct logit attribution over the standard contributions. - [Ablation](ablation.md): the effect of removing a contribution, a head or a block. - [Activation patching](activation-patching.md): a value from a clean run patched into a corrupt one. - [Steering](steering.md): a direction added to the residual stream. - [Probing](probing.md): a linear probe per layer. - [Cross-family sweep](cross-family-sweep.md): one experiment over several checkpoints. - [Expert ablation](expert-ablation.md): every expert's effect on a prediction. - [Image pathway](image-pathway.md): the image's effect in a vision-language model. - [DeltaNet state](delta-net-state.md): the recurrent state of a hybrid, read and patched. # Logit lens Where in depth a prediction emerges, on any family. `model.project_on_vocab(layer.layer_output)` reads the residual stream leaving a block through the model's own final norm, its unembedding and whatever the model does after the head, so one loop over `model.layers` is the logit lens on every family. The method is the one in nnsight's [logit lens tutorial](https://nnsight.net/tutorials/tutorials/probing/logit_lens/). What nnterp changes is that the norm, the head and any softcap are no longer yours to find per architecture. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu") prompt = "The Eiffel Tower is in the city of" paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0] lens = [] # bound outside the trace: a list made inside it is gone afterwards with model.trace(prompt): for layer in model.layers: lens.append(model.project_on_vocab(layer.layer_output)[0, -1].save()) # [vocab_size] logits = model.logits[0, -1].save() for i, row in enumerate(lens): print(f"block {i:2d} top {model.tokenizer.decode(row.argmax())!r:>10} P(Paris) {row.softmax(-1)[paris]:.3f}") print(torch.equal(lens[-1], logits)) # block 0 top ' the' P(Paris) 0.000 # block 1 top ' the' P(Paris) 0.000 # block 2 top ' the' P(Paris) 0.000 # block 3 top ' the' P(Paris) 0.000 # block 4 top ' the' P(Paris) 0.000 # block 5 top ' the' P(Paris) 0.000 # block 6 top ' East' P(Paris) 0.001 # block 7 top ' Ing' P(Paris) 0.003 # block 8 top ' Rome' P(Paris) 0.025 # block 9 top ' London' P(Paris) 0.248 # block 10 top ' Paris' P(Paris) 0.183 # block 11 top ' Paris' P(Paris) 0.070 # True ``` GPT-2 decodes ` the` for the first six blocks, a city from block 8 (` Rome`, then ` London`) and ` Paris` from block 10. The probability says more than the argmax: P(Paris) peaks at block 9, where ` London` is still the top token, and falls through the last two blocks to the 0.070 the model outputs. The loop reads the blocks in forward order and `logits` after them, which is the order one trace requires. On many newer checkpoints the early blocks decode to unrelated tokens and the answer appears only in the last few blocks; `Qwen/Qwen2.5-0.5B` puts ` Paris` on top from block 21 of 24 and nothing recognizable before it. That is not a wiring fault, and the wiring check is what decides. The final `True` is the wiring check. At the last block the lens is the model's own computation, so it must equal `model.logits` exactly. Run it first on a new checkpoint. A lens that differs there is reading through the wrong norm or head, and every earlier block is then on the wrong scale. Decoding `model.lm_head(hidden)` without the norm is not a cheaper lens. A norm is not affine, so the curve comes out confident and wrong rather than as an error. ## The step after the head `model.logits` is what the model returns and `model.lm_head.output` the raw projection. On Gemma-2 the two differ, because the model applies `cap * tanh(logits / cap)` after the head (`final_logit_softcapping`, 30.0 on `google/gemma-2-2b`, where the raw projection and the logits differ by up to 7.8 at the last position of this prompt). `project_on_vocab` reads that key from the text config and applies the same cap, so the wiring check passes there too and every block is read on the model's scale rather than an uncapped one. Cohere's `logit_scale` and Granite's `logits_scaling` are the same kind of step, and those families define their own `project_on_vocab` with it. ## On another family The recipe runs unchanged. `layer_output` is a tensor on every family, including the families with a tuple-returning block (BLOOM, MPT, Falcon), and on DeepSeek-V4, where it carries an extra stream axis, the family's `project_on_vocab` collapses the streams the way the model does. The lens reads module boundaries, so it does not need `attn_implementation="eager"`. The per-family work is all in `project_on_vocab`: the norm, the head and the step after it. [Contribution decomposition](contribution-decomposition.md) reads each sublayer's addition through the same head, in a form that sums. # Attention patterns Reading and scoring attention heads. `model.layers[i].self_attn.attention_probabilities` is the attention pattern after the softmax, `[batch, num_heads, query, key]`, on every family that runs transformers' eager attention. Load with `attn_implementation="eager"` and the per-head scores below are three lines on the saved tensor. The pattern is served in the model's dtype, and small Pythia and GPT-NeoX checkpoints load in float16, where on the CPU the row sums miss the tolerance below and the entropy is NaN, so pass `dtype=torch.float32` there. nnsight's [attention tutorial](https://nnsight.net/tutorials/tutorials/attention/attention_heads/) explains the metrics and the head types. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu", attn_implementation="eager") prompt = "When Mary and John went to the store, John gave a drink to Mary" LAYER = 4 with model.trace(prompt): pattern = model.layers[LAYER].self_attn.attention_probabilities.save() # [batch, num_heads, query, key] print(pattern.shape, torch.equal(pattern.tril(), pattern), torch.allclose(pattern.sum(-1).float(), torch.ones(pattern.shape[:-1]), atol=1e-4)) entropy = -(pattern * (pattern + 1e-12).log()).sum(-1).mean(-1)[0] # [num_heads]: low means sharp previous = pattern.diagonal(offset=-1, dim1=-2, dim2=-1).mean(-1)[0] # mass on key i-1 first = pattern[0, :, 1:, 0].mean(-1) # mass on key 0 from every later query for h in range(model.num_heads): print(f"head {h:2d} entropy {entropy[h]:.2f} previous {previous[h]:.2f} first {first[h]:.2f}") # torch.Size([1, 12, 15, 15]) True True # head 0 entropy 0.79 previous 0.36 first 0.41 # head 1 entropy 1.01 previous 0.23 first 0.56 # head 2 entropy 1.01 previous 0.11 first 0.72 # head 3 entropy 1.14 previous 0.29 first 0.36 # head 4 entropy 1.11 previous 0.15 first 0.65 # head 5 entropy 0.85 previous 0.16 first 0.73 # head 6 entropy 1.26 previous 0.16 first 0.58 # head 7 entropy 1.17 previous 0.19 first 0.31 # head 8 entropy 0.86 previous 0.11 first 0.76 # head 9 entropy 1.11 previous 0.16 first 0.63 # head 10 entropy 0.60 previous 0.09 first 0.86 # head 11 entropy 0.00 previous 1.00 first 0.07 ``` The two checks say that the pattern is causal (nothing above the diagonal) and that every row sums to one. Head 11 puts all of its mass on the previous token, with entropy zero: GPT-2's block-4 previous-token head. Head 10 puts 0.86 of its mass on the first token, which is what a head does when it has nothing to attend to on this prompt. The rows of a GPT-OSS pattern sum to less than one, by the mass a learned sink column took before it was dropped. The pattern's meaning, mass on real keys, is unchanged. ## Editing the pattern The pattern is the matrix that the attention's values are mixed with, so an in-place write reaches the model. Zero one head's rows and the head reads nothing: ```python with model.trace(prompt): clean = model.next_token_probs.save() with model.trace(prompt): model.layers[LAYER].self_attn.attention_probabilities[:, 11] = 0 # head 11 reads nothing ablated = model.next_token_probs.save() print(model.probs_to_dict(clean[0], k=3)) print(model.probs_to_dict(ablated[0], k=3)) # {' and': 0.3203795254230499, ',': 0.24333427846431732, '.': 0.1615239828824997} # {' and': 0.34432530403137207, ',': 0.21988408267498016, '.': 0.14424201846122742} ``` Zeroing the pattern and zeroing the head's slice of `attention_head_outputs` give the same logits. The head-level recipes, with the clean and ablated runs in one trace, are on the [Ablation](ablation.md) page. ## Without eager attention A checkpoint loads `sdpa` unless you ask for eager, and the pattern is then unavailable. `support()` reports that before anything runs, and reading the value anyway raises `nnterp.Unavailable` with the same text at that line: ```python sdpa = StandardizedTransformer("openai-community/gpt2", device="cpu") print(sdpa.support()["self_attn.attention_probabilities"][0]) # read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager' ``` ## On another family The layout is `[batch, num_heads, query, key]` everywhere, so the scores above do not change. What changes is which blocks have a pattern. On a hybrid (Qwen3-Next, Qwen3.5) three blocks in four carry `linear_attn` instead of `self_attn`, with no pattern, so pick the blocks from the module list outside the trace (`[i for i, layer in enumerate(model.layers) if getattr(layer, "self_attn", None) is not None]`) rather than from `num_layers // 2`. The value's shape and the eager requirement are described in [Attention](../guide/values/attention.md). # Contribution decomposition Direct logit attribution over the standard contributions. `layer_input + attention_output + mlp_output == layer_output` on every family, so direct logit attribution reads `attention_output` and `mlp_output` from each block, pushes each term through the final norm frozen at the real stream and the unembedding, and the terms sum to the model's logit. `attention_head_outputs` times the output projection's weight splits a block's attention term per head. The method is from Elhage et al. (2021), [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html). ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu", attn_implementation="eager") prompt = "The Eiffel Tower is in the city of" paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0] parts = {} # bound outside the trace with model.trace(prompt): base = model.layers[0].layer_input.save() # the stream entering block 0 for i, layer in enumerate(model.layers): # block i's attention, then its MLP: forward order parts["attn", i] = layer.self_attn.attention_output.save() parts["mlp", i] = layer.mlp.mlp_output.save() final = model.layers[-1].layer_output.save() logit = model.lm_head.output[0, -1, paris].save() # the head's output, before any step after it print(torch.allclose(base + sum(parts.values()), final, atol=1e-3)) # the stream is the sum of its contributions norm, head = model.norm._module, model.lm_head._module x = final[0, -1] # the stream the final norm reads at this position with torch.no_grad(): shift = norm(torch.zeros_like(x)) # a LayerNorm's bias; zeros on an RMSNorm centers = torch.allclose(norm(torch.ones_like(x)), shift) # a LayerNorm subtracts the mean first center = (lambda t: t - t.mean(-1, keepdim=True)) if centers else (lambda t: t) eps = getattr(norm, "variance_epsilon", getattr(norm, "eps", 0.0)) probe = torch.ones_like(x); probe[1::2] = -1 # mean 0, mean square 1 gain = (norm(probe) - shift) / probe * (1 + eps) ** 0.5 # the norm's weight, or 1 + weight on Gemma scale = (center(x).pow(2).mean() + eps).rsqrt() # the one nonlinear step, frozen at the real stream frozen = lambda t: gain * center(t) * scale # the final norm as a linear map at this position row = head.weight[paris] dla = {key: float(frozen(part[0, -1]) @ row) for key, part in parts.items()} dla["base"] = float(frozen(base[0, -1]) @ row) constant = float(shift @ row) + (float(head.bias[paris]) if head.bias is not None else 0.0) print(f"{sum(dla.values()) + constant:.4f} {float(logit):.4f}") for i in range(model.num_layers): print(f"block {i:2d} attn {dla['attn', i]:+.3f} mlp {dla['mlp', i]:+.3f}") # True # -92.0873 -92.0873 # block 0 attn -0.264 mlp -0.320 # block 1 attn -0.291 mlp +0.049 # block 2 attn +0.190 mlp -0.232 # block 3 attn +0.051 mlp -0.049 # block 4 attn -0.230 mlp -0.080 # block 5 attn +0.303 mlp -0.476 # block 6 attn -0.480 mlp -1.090 # block 7 attn +0.303 mlp -0.239 # block 8 attn -0.346 mlp -1.012 # block 9 attn +1.006 mlp -1.404 # block 10 attn -1.263 mlp -17.002 # block 11 attn -36.726 mlp -24.635 ``` The attributions plus the constant sum to the model's own ` Paris` logit, -92.0873, to four decimals. The last two blocks dominate the logit's scale (block 11's attention -36.7 and MLP -24.6, block 10's MLP -17.0), and block 9's attention is the largest positive term, +1.006. An attribution to one logit includes what a term adds to every token at once. `row = head.weight[paris] - head.weight[rome]` attributes the contrast between two tokens instead, with the same code. The norm has to be frozen because it rescales each position by that position's own statistics, so it does not distribute over the sum. Holding the scale at the real stream's value makes it a linear map (a centering, a gain, a scale), and the LayerNorm's bias plus any `lm_head` bias enters once, as `constant`. The recipe reads the centering, the gain and the bias off the norm module itself (its output on zeros is its bias, on ones says whether it centers, on a ±1 vector gives its gain), so an RMSNorm (Llama), a LayerNorm (GPT-2) and Gemma's `1 + weight` gain need no separate cases. `model.project_on_vocab(contribution)` is the other reading, the full norm applied to one term alone. It answers what the term would predict by itself and does not sum. ## Per head `attention_head_outputs` is each head's output before concatenation and the output projection, `[batch, seq, num_heads, head_dim]`. The projection is linear, so a head's share of `attention_output` is its slice times the matching rows of the projection's weight. The projection is `o_proj` on every family (GPT-2's `c_proj` is aliased), but its weight layout differs: `torch.nn.Linear` stores `[out, in]`, GPT-2's `Conv1D` stores `[in, out]`. ```python LAYER = 9 attention = model.layers[LAYER].self_attn with model.trace(prompt): heads = attention.attention_head_outputs.save() # [batch, seq, num_heads, head_dim] contribution = attention.attention_output.save() # [batch, seq, hidden_size] proj = attention.o_proj._module W = proj.weight.t() if isinstance(proj, torch.nn.Linear) else proj.weight # [num_heads * head_dim, hidden_size] H, D = heads.shape[2], heads.shape[3] per_head = torch.einsum("bshd,hdo->bsho", heads, W.reshape(H, D, -1)) # [batch, seq, num_heads, hidden_size] print(torch.allclose(per_head.sum(2) + proj.bias, contribution, atol=1e-3)) # the heads sum to the contribution with torch.no_grad(): head_dla = [float(frozen(per_head[0, -1, h]) @ row) for h in range(H)] bias_dla = float(frozen(proj.bias) @ row) print([round(v, 3) for v in head_dla]) print(round(sum(head_dla) + bias_dla, 3), round(dla["attn", LAYER], 3)) # True # [-0.084, -0.132, -0.066, 0.791, 0.116, -0.153, -0.0, -0.072, 0.872, 0.048, 0.008, 0.064] # 1.006 1.006 ``` Block 9's +1.006 is two heads: head 8 (+0.872) and head 3 (+0.791), with the other heads near zero and the projection's bias taking 0.39 back off. Head 8 of block 9 is also the head that [Ablation](ablation.md) finds by removing it. The per-head sum equals `attention_output` where the projection's output is what the block adds, as on GPT-2 and Llama. On a post-norm family (Gemma-2, Gemma-3, OLMo-2) a norm sits between the projection and the residual add, so the heads sum to `o_proj.output` and not to `attention_output`, and the split is of the pre-norm tensor. On Granite the contribution is the projection's output times `residual_multiplier`. Under grouped-query attention the head axis is still the query heads, so the slicing does not change. ## On another family The identity does the per-family work. `attention_output` and `mlp_output` are what each sublayer adds whether the block is sequential, parallel or post-norm, so the recipe runs unchanged there. Not every family is a plain sum. On Gemma-4 each block multiplies its sum by `layer_scalar`, so a term enters the final stream weighted by the product of the scalars of its block and every later block. Doge and ZAYA gate the stream per channel, and DeepSeek-V4 keeps parallel streams. Weight each term accordingly before the unembedding. The frozen norm then applies to the weighted terms unchanged. # Ablation Removing a contribution, a head or a block and measuring the effect. Zero `attention_output` or `mlp_output` and the sublayer adds nothing to the residual stream, on every family, including the families that add the residual inside the module (BLOOM, MPT) or norm the output first (Gemma-2, Gemma-3, OLMo-2). One head is one slice of `attention_head_outputs`, and whole blocks are `model.skip_layers`. The clean and ablated runs are invokes of one trace, so one forward gives the comparison. nnsight's [batching](https://nnsight.net/features/8_batching/) page covers invokes. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu", attn_implementation="eager") prompt = "The Eiffel Tower is in the city of" paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0] LAYER, HEAD = 9, 8 block = model.layers[LAYER] runs = {} # bound outside the trace with torch.no_grad(), model.trace() as tracer: with tracer.invoke(prompt): runs["clean"] = model.next_token_probs[0, paris].save() with tracer.invoke(prompt): block.self_attn.attention_output[:] = 0 # this invoke's rows only runs["attention"] = model.next_token_probs[0, paris].save() with tracer.invoke(prompt): block.mlp.mlp_output[:] = 0 runs["mlp"] = model.next_token_probs[0, paris].save() with tracer.invoke(prompt): block.self_attn.attention_head_outputs[..., HEAD, :] = 0 runs[f"head {HEAD}"] = model.next_token_probs[0, paris].save() with torch.no_grad(), model.trace(prompt): # a skip covers every row, so it gets its own trace model.skip_layers(LAYER, LAYER + 1) runs["blocks 9-10"] = model.next_token_probs[0, paris].save() for name, p in runs.items(): print(f"{name:12s} P(Paris) {float(p):.3f}") # clean P(Paris) 0.070 # attention P(Paris) 0.033 # mlp P(Paris) 0.053 # head 8 P(Paris) 0.036 # blocks 9-10 P(Paris) 0.006 ``` Block 9's attention carries half of P(Paris) (0.070 to 0.033), and one head, head 8, accounts for almost all of that (0.036). The block's MLP matters less (0.053). Skipping blocks 9 and 10 leaves 0.006. The writes are in place on the served tensors, so they reach the model, and a write in one invoke touches only that invoke's rows. `skip_layers` is in a trace of its own because a skip has to cover every row of a forward. Inside one invoke of a batched trace it raises `ValueError: A batched .skip() has to cover every row`. ## Mean ablation Replacing a contribution with its mean over a reference set, instead of zero, keeps the stream nearer to its usual state and asks whether the deviation from average matters: ```python reference = ["The capital of France is", "Berlin lies in", "London is the capital of", "Tokyo is in"] with torch.no_grad(), model.trace(reference): mean = model.layers[LAYER].mlp.mlp_output[:, -1].mean(0).save() # [hidden_size] with torch.no_grad(), model.trace(prompt): model.layers[LAYER].mlp.mlp_output[:, -1] = mean p = model.next_token_probs[0, paris].save() print(f"mean-ablated P(Paris) {float(p):.3f}") # mean-ablated P(Paris) 0.073 ``` Zeroing block 9's MLP costs 0.017 of P(Paris), and replacing it with its mean costs nothing. The two answer different questions, so say which you ran. A list in one trace is left-padded, which is what makes `[:, -1]` every reference prompt's last token. ## Every head of a block One invoke per head, in one forward: ```python rows = [] with torch.no_grad(), model.trace() as tracer: for head in range(model.num_heads): with tracer.invoke(prompt): model.layers[LAYER].self_attn.attention_head_outputs[..., head, :] = 0 rows.append(model.next_token_probs[0, paris].save()) print([round(float(p), 3) for p in rows]) # [0.077, 0.07, 0.068, 0.065, 0.065, 0.066, 0.069, 0.072, 0.036, 0.073, 0.069, 0.07] ``` Head 8 is the only head of block 9 with an effect. The others move P(Paris) by under 0.01 in either direction. Zeroing a head's rows of `attention_probabilities` gives the same logits as zeroing its slice here. Both need `attn_implementation="eager"`. The boundary values `attention_output`, `mlp_output` and `layer_output` do not. Do not slice `attention_output` into `head_dim`-wide column blocks instead, because after the output projection the hidden axis no longer decomposes per head. ## On another family The contribution values do the per-family work, so the recipe runs unchanged wherever the block has the module. OPT has no `mlp` module, and a hybrid's DeltaNet blocks have `linear_attn` instead of `self_attn`, so decide which blocks to ablate from `model.support()` or the module list, outside the trace. On Granite the block adds `residual_multiplier` times the sublayer's output, and nnterp serves the scaled copy, so a write at one position moves the other positions by rounding in bf16. Load in float32 for small edits there. The meaning of the values is in [Residual stream](../guide/values/residual-stream.md). # Activation patching Patching a value from a clean run into a corrupt one. `layer_output` is the residual stream leaving a block and `attention_output` and `mlp_output` are what its sublayers add, on every family, so a patch is one assignment at one of those values at a position, and the same line runs on a sequential, parallel or post-norm block. The method is nnsight's [activation patching tutorial](https://nnsight.net/tutorials/tutorials/causal_mediation_analysis/activation_patching/). The checkpoint is `Qwen/Qwen2.5-0.5B`, where the clean prompt gives P(Paris) 0.797 and the corrupt one 0.000 (GPT-2 gives 0.070 and 0.003, too little to read). ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("Qwen/Qwen2.5-0.5B", device="cpu", dtype=torch.float32) clean = "The Eiffel Tower is in the city of" corrupt = "The Colosseum is in the city of" assert len(model.tokenizer(clean).input_ids) == len(model.tokenizer(corrupt).input_ids) # positions must correspond paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0] LAYER, POS = 21, -1 with model.trace(clean): clean_resid = model.layers[LAYER].layer_output.save() # [1, seq, hidden_size] with model.trace() as tracer: with tracer.invoke(corrupt): baseline = model.next_token_probs[0, paris].save() with tracer.invoke(corrupt): model.layers[LAYER].layer_output[:, POS] = clean_resid[:, POS] patched = model.next_token_probs[0, paris].save() print(f"P(Paris) corrupt {float(baseline):.3f} patched {float(patched):.3f}") # P(Paris) corrupt 0.000 patched 0.778 ``` Pasting the clean stream into the last position at block 21 restores 0.778 of the clean 0.797. The write is in place on the served tensor, so it reaches the model, and it touches only its invoke's rows. The unpatched invoke in the same trace is the baseline. `clean_resid` is an ordinary tensor by the time the second trace runs, so nothing has to be synchronized. ## Across invokes The clean and corrupt rows can share one forward. A write that uses another invoke's value then needs a [barrier](https://nnsight.net/documentation/intervention/barrier/), because the assignment evaluates its right-hand side before its own invoke waits: ```python with model.trace() as tracer: barrier = tracer.barrier(2) with tracer.invoke(clean): source = model.layers[LAYER].layer_output[:, POS] barrier() # source is read with tracer.invoke(corrupt): barrier() # wait for it model.layers[LAYER].layer_output[:, POS] = source patched = model.next_token_probs[0, paris].save() print(f"across invokes {float(patched):.3f}") # across invokes 0.778 ``` ## Sweeping the layers Cache the clean run once, then one invoke per block, at the position you ask about: ```python cache = [] with model.trace(clean): for layer in model.layers: cache.append(layer.layer_output.save()) for pos in (2, -1): # the subject's second token, the last token sweep = [] with model.trace() as tracer: for i in range(model.num_layers): with tracer.invoke(corrupt): model.layers[i].layer_output[:, pos] = cache[i][:, pos] sweep.append(model.next_token_probs[0, paris].save()) print(f"pos {pos:2d}:", " ".join(f"{float(p):.2f}" for p in sweep)) # pos 2: 0.38 0.61 0.64 0.55 0.50 0.52 0.53 0.52 0.25 0.23 0.30 0.32 0.30 0.30 0.27 0.33 0.17 0.06 0.05 0.05 0.01 0.00 0.00 0.00 # pos -1: 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.00 0.32 0.78 0.80 0.80 ``` Position 2 is `iff` in the clean prompt and `os` in the corrupt one, the first token where the two differ. A patch there restores most of the answer through block 7 (0.5 to 0.6) and fades to nothing by block 20. At the last position nothing happens until block 19, then 0.32 at block 20 and 0.78 from block 21 on. The fact leaves the subject's position in the early blocks and arrives at the last position in blocks 20 and 21. The last entry of a last-position sweep is fixed by construction, since patching the last block there reproduces the clean prediction. That row says how far back a patch reaches, and the position axis is what localizes. ## Patching a contribution Patching `layer_output` at block `i` replaces everything computed at that position up to block `i`. Patching a contribution replaces one sublayer's addition and nothing else: ```python with model.trace(clean): clean_attn = model.layers[LAYER].self_attn.attention_output.save() clean_mlp = model.layers[LAYER].mlp.mlp_output.save() with model.trace() as tracer: with tracer.invoke(corrupt): model.layers[LAYER].self_attn.attention_output[:, POS] = clean_attn[:, POS] attn = model.next_token_probs[0, paris].save() with tracer.invoke(corrupt): model.layers[LAYER].mlp.mlp_output[:, POS] = clean_mlp[:, POS] mlp = model.next_token_probs[0, paris].save() print(f"attention_output {float(attn):.3f} mlp_output {float(mlp):.3f}") # attention_output 0.151 mlp_output 0.018 ``` Block 21's attention brings 0.151 of the answer to the last position on its own and its MLP 0.018. The rest of the 0.778 was already in the stream from block 20, where the `layer_output` patch gives 0.32. One head is a slice of `attention_head_outputs`, `[..., HEAD, :]`, patched the same way under `attn_implementation="eager"`. ## On another family The identity `layer_input + attention_output + mlp_output == layer_output` is the per-family work. A contribution patch means "the block adds the clean sublayer's contribution instead of the corrupt one's" on every block type, and `layer_output` is a tensor even where the block returns a tuple (BLOOM, MPT, Falcon), so there is no `[0]` and no tuple to rebuild. What does change is the tokenizer. The two prompts must tokenize to the same length for positions to correspond, which the `assert` checks, and `[:, -1]` survives a length difference where an absolute index does not. On Granite a patch at one position moves the other positions by rounding in bf16, because the block adds a scaled copy. Load in float32 there. # Steering Adding a direction to the residual stream. `model.steer(layer, vector, factor, token_positions=-1)` adds `factor * vector` in place to `layer_output` of one block, and the direction is a difference of `layer_output` between two prompts, so deriving and applying a steering vector are the same two snippets on every family. The method is activation addition (Turner et al., 2023, [Steering Language Models With Activation Engineering](https://arxiv.org/abs/2308.10248)). What the recipe adds is reading the factor against the stream's norm at the block you steer, since that norm differs several-fold between checkpoints and across a model's blocks. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", device="cpu") LAYER = 6 with model.trace() as tracer: with tracer.invoke("I love this so much"): positive = model.layers[LAYER].layer_output[0, -1].save() # [hidden_size] with tracer.invoke("I hate this so much"): negative = model.layers[LAYER].layer_output[0, -1].save() direction = positive - negative direction = direction / direction.norm() # unit norm: factors compare across directions prompt = "I went to the bakery and" with torch.no_grad(), model.trace(prompt): scale = model.layers[LAYER].layer_output[0, -1].norm().save() # the stream's norm where you steer print(float(scale)) factor = 0.25 * float(scale) # a fraction of it; sweep the fraction with model.trace() as tracer: with tracer.invoke(prompt): baseline = model.next_token_probs.save() with tracer.invoke(prompt): model.steer(LAYER, direction, factor=factor, token_positions=-1) steered = model.next_token_probs.save() print(model.probs_to_dict(baseline[0], k=3)) print(model.probs_to_dict(steered[0], k=3)) # 92.03986358642578 # {' bought': 0.09450755268335342, ' I': 0.08870173245668411, ' asked': 0.07155062258243561} # {' I': 0.1267288476228714, ' they': 0.055561766028404236, ' bought': 0.05142388120293617} ``` The stream leaving GPT-2's block 6 has norm 92 at the last position, so a factor of 23 adds a quarter of that along the direction, and the top token moves from ` bought` to ` I`. The add is in place on the served tensor, so it reaches the model, and the invoke that does not call `steer` is untouched. A raw factor means nothing on its own. On `Qwen/Qwen2.5-0.5B` the same prompt's stream has norm 13 at block 6 and 79 at block 22, against 57 to 465 across GPT-2's blocks, so measure the norm where you steer and sweep the fraction, not the factor. `steer` is one line that you could write yourself, `out[:, -1] += factor * direction.to(out)` on `layer_output`. Use the hand form when the edit is not an add, such as projecting the direction out. ## At every step of a generate Steering a generation means steering every decode step, not only the prompt. `tracer.iter[:N]` runs the body on the prefill and on each of the `N - 1` decode steps, and `min_new_tokens=N` holds the run to that bound: ```python N = 15 for f in (0.0, factor, -factor): with model.generate(prompt, max_new_tokens=N, min_new_tokens=N, do_sample=False) as tracer: for step in tracer.iter[:N]: # the prefill and every decode step model.steer(LAYER, direction, factor=f, token_positions=-1) ids = tracer.result.save() print(f"{f:+.1f}: {model.tokenizer.decode(ids[0][-N:])!r}") # +0.0: ' bought a bag of cookies. I was so excited. I was so excited' # +23.0: ' I was able to get some bread and a little bit of butter. I' # -23.0: ' asked the owner if he could make a cake for me. He said he' ``` Adding and subtracting the direction give two different continuations, and both stay fluent at a quarter of the norm. `token_positions=-1` is the right index across steps. On the prefill it is the prompt's last token and on each decode step the one token that step processes, where an absolute position would be out of range from the first decode step on. nnsight's [multiple token generation](https://nnsight.net/features/4_multiple_token/) page covers `tracer.iter` and `tracer.result`. !!! warning "A bare `steer` in a `generate` body runs once" Outside `tracer.iter`, the call fires on the prefill only. The continuation still changes, because the prompt was read differently (on this prompt and factor it becomes ` I was like, 'Oh my God, I'm so sorry. I`), which makes the mistake look like a modest success. Put the call inside the loop. ## On another family `layer_output` and `steer` do the per-family work. The stream leaving a block is `layer_output` on every block of every family, a DeltaNet block of a hybrid included, and `steer` moves the vector to the stream's dtype and device, so a CPU float32 direction steers a bf16 GPU model unchanged. The one thing that changes with the checkpoint is the norm, which is why the recipe measures it rather than fixing a factor. `steer`'s arguments, including `batch_index` for one row of a batch and a list of ascending blocks, are in [Methods](../guide/methods.md). # Probing A probe asks whether a property is linearly decodable from the residual stream at each depth: collect `layer_output` at one position of every prompt in a labeled set, fit a linear classifier per layer, and read accuracy against the layer. `get_token_activations` does the collection against the standard values, so the recipe below runs unchanged on every family. The probe itself is a few lines of torch. ## Collecting the activations ```python import torch from nnterp import StandardizedTransformer from nnterp.nnsight_utils import get_token_activations model = StandardizedTransformer("openai-community/gpt2", dispatch=True, device="cpu") positive = ["wonderful", "fantastic", "delightful", "excellent", "brilliant", "joyful", "superb", "lovely", "charming", "pleasant"] negative = ["terrible", "awful", "dreadful", "disgusting", "horrible", "miserable", "dismal", "boring", "painful", "ugly"] templates = ["The movie was {}, so afterwards we", "I found the book {} and told my", "That meal was {}, which is why the", "Their performance was {}, and the crowd"] texts = [t.format(w) for w in positive + negative for t in templates] labels = torch.tensor([1.0] * 40 + [0.0] * 40) acts = get_token_activations(model, texts, idx=-1) print(acts.shape) # torch.Size([12, 80, 768]) ``` One batched forward gives `[layers, prompts, hidden_size]` on the CPU, the residual stream `model.layers[i].layer_output` at the last position of every prompt. The last position of a batch needs left padding, which nnsight sets at load for a causal model. The helper raises `ValueError` before running if the tokenizer pads right. `layers=` picks some blocks, `get_activations=lambda m, i: m.layers[i].mlp.mlp_output` reads another site, and `collect_token_activations_batched(model, texts, batch_size)` returns the same tensor for a set too large for one batch ([`nnsight_utils`](../documentation/nnsight_utils.md)). The adjective sits mid-sentence, several tokens before the position read, and the test words never appear in training. Split by word, not by prompt: With the words shared between the splits, the probe can read the adjective's identity instead of its sentiment. ## One probe per layer ```python held_out = {"superb", "lovely", "charming", "pleasant", "dismal", "boring", "painful", "ugly"} is_test = torch.tensor([any(w in t for w in held_out) for t in texts]) train, test = (~is_test).nonzero().squeeze(1), is_test.nonzero().squeeze(1) def fit_probe(x, y, ridge=1e-2): """Ridge regression of the {-1, +1} label on [n, hidden_size]; returns predict(x) -> bool tensor.""" mean, std = x.mean(0), x.std(0) + 1e-6 design = lambda z: torch.cat([(z - mean) / std, torch.ones(len(z), 1)], dim=1) z = design(x) penalty = torch.eye(z.shape[1]) penalty[-1, -1] = 0 # no penalty on the bias w = torch.linalg.solve(z.T @ z + ridge * len(z) * penalty, z.T @ (2 * y - 1)) return lambda x_new: design(x_new) @ w > 0 for layer in range(model.num_layers): x = acts[layer].float() predict = fit_probe(x[train], labels[train]) accuracy = (predict(x[test]).float() == labels[test]).float().mean() print(f"layer {layer:2d} accuracy on held-out words {float(accuracy):.2f}") # layer 0 accuracy on held-out words 0.78 # layer 1 accuracy on held-out words 0.94 # layer 2 accuracy on held-out words 0.97 # layer 3 accuracy on held-out words 1.00 # ... # layer 11 accuracy on held-out words 1.00 ``` The probe is closed-form ridge regression, standardized on the training rows only. Keep the ridge on and report it, because with `hidden_size` far above the number of examples an unregularized probe fits anything. ## Reading the curve Sentiment is readable at the sentence's last token from layer 3 on, for words the probe never saw. The position read is the final token (`we`, `my`, `the`, `crowd`), the same in both classes, so whatever the probe finds there was carried from the adjective three to five tokens back. Layer 0's attention already moves enough of it for 0.78, and layers 1 to 3 complete it. A shuffled-label control bounds the memorization. Refit with `labels[torch.randperm(80)]` and the same probe scores 0.50, 0.44 and 0.34 at layers 0, 6 and 11. A probe is correlational. To claim that the model uses the direction, add it to the stream with `model.steer` and check that the next-token distribution moves while a random direction of the same norm does not. [Steering](steering.md) has the factor band. ## On another family Nothing in the recipe names a family. `layer_output` is the residual stream leaving block `i` as a tensor on every family, whatever the native block returns, and the helper reads nothing else. Change the checkpoint and the word lists. Accuracies are not comparable across models of different `hidden_size` unless the probe's capacity is held fixed, so compare the shape of the curves and report the layer-0 value with each. # Cross-family sweep One trace body over several checkpoints: zero each block's mixer contribution and each block's MLP contribution in turn, record how far the next-token distribution moves, and tabulate the result per family. What the loop adds to a single-model script is a `support()` read, a guard on the MLP, and the mixer decided outside the trace. The trace body is the same on a GPT-2, a Pythia, a Llama and a Mamba. ## The loop ```python import torch import torch.nn.functional as F import nnterp from nnterp import StandardizedTransformer REPOS = ["openai-community/gpt2", "EleutherAI/pythia-70m-deduped", "HuggingFaceTB/SmolLM2-135M", "state-spaces/mamba-130m-hf"] prompt = "The Eiffel Tower is in the city of" def mixer(layer): """The block's sequence mixer: self_attn, or linear_attn on a recurrent block.""" attn = getattr(layer, "self_attn", None) return attn if attn is not None else layer.linear_attn table = {} for repo in REPOS: model = StandardizedTransformer(repo, dispatch=True, device="cpu", dtype=torch.float32, attn_implementation="eager") support = model.support() # before any trace has_mlp = support.get("mlp.mlp_output", "absent") is None # a module no block has (Mamba) has no key mixers = [mixer(layer) for layer in model.layers] # decided outside the trace attention = [i for i, layer in enumerate(model.layers) if getattr(layer, "self_attn", None) is not None] if len(attention) < model.num_layers: nnterp.route_kernels(model.family, "torch") # a recurrent mixer: the torch kernels, before the first trace ablated = {} # containers made outside the trace with torch.no_grad(), model.trace() as tracer: with tracer.invoke(prompt): clean = model.logits[:, -1].log_softmax(-1).save() for i in range(model.num_layers): with tracer.invoke(prompt): mixers[i].attention_output[:] = 0 ablated["mixer", i] = model.logits[:, -1].log_softmax(-1).save() if has_mlp: with tracer.invoke(prompt): model.layers[i].mlp.mlp_output[:] = 0 ablated["mlp", i] = model.logits[:, -1].log_softmax(-1).save() entropy = {} with torch.no_grad(), model.trace(prompt): for i in attention: p = model.layers[i].self_attn.attention_probabilities entropy[i] = (-p * (p + 1e-12).log()).sum(-1).mean().save() kl = lambda logprobs: F.kl_div(logprobs, clean, log_target=True, reduction="sum") # KL(clean || ablated) mixer_kl = torch.tensor([kl(ablated["mixer", i]) for i in range(model.num_layers)]) mlp_kl = torch.tensor([kl(ablated["mlp", i]) for i in range(model.num_layers)]) if has_mlp else None table[model.family.__name__.split(".")[-1]] = { "blocks": f"{len(attention)}/{model.num_layers} attention", "top": model.tokenizer.decode(clean.argmax(-1)), "mixer": f"{float(mixer_kl.max()):.2f} nats at block {int(mixer_kl.argmax())}", "mlp": f"{float(mlp_kl.max()):.2f} nats at block {int(mlp_kl.argmax())}" if has_mlp else "no MLP", "entropy": f"{sum(float(e) for e in entropy.values()) / len(entropy):.2f}" if entropy else "-", } del model print(f"{'family':9s} {'blocks':16s} {'top':8s} {'largest mixer ablation':24s} {'largest MLP ablation':24s} entropy") for family, row in table.items(): print(f"{family:9s} {row['blocks']:16s} {row['top']!r:8s} {row['mixer']:24s} {row['mlp']:24s} {row['entropy']}") # family blocks top largest mixer ablation largest MLP ablation entropy # gpt2 12/12 attention ' Paris' 3.04 nats at block 0 4.25 nats at block 0 0.63 # gpt_neox 6/6 attention ' the' 8.42 nats at block 0 4.14 nats at block 1 0.61 # llama 30/30 attention ' Paris' 5.97 nats at block 1 16.37 nats at block 0 0.60 # mamba 0/24 attention ' Paris' 3.97 nats at block 0 no MLP - ``` `model.support()` returns, for every value, `None` when every block has it or `{block: reason}` where some do not. A module no block has (Mamba's `mlp`) has no key, which is why the guard uses `.get`. Guard on it rather than on `hasattr` inside the trace, which raises `nnterp.Unavailable` for an unavailable value ([Availability](../guide/values/availability.md)). A block has either `self_attn` or `linear_attn`, and the list comprehension outside the trace is the place to decide which. The KL is taken on log-probabilities because `next_token_probs` underflows to exact zeros, and `p * (p.log() - q.log())` is then NaN. The loads are in float32 so that one dtype serves every checkpoint. Pythia's weights are stored in float16, and loaded at that dtype its KLs come back NaN. `attn_implementation="eager"` is what `attention_probabilities` needs on the attention families, and Mamba ignores it. `route_kernels` runs only where a block has no `self_attn`. A recurrent mixer's optimized kernels have no Python source and may not run on the CPU, and the torch kernels have to be bound before the family's first trace ([Mamba](../guide/architectures/mamba.md)). ## Reading the table Zero ablation at block 0 or 1 moves the distribution by 3 to 8 nats for the mixer and 4 to 16 nats for the MLP on every checkpoint that has one. The stream leaving the first blocks is then unlike anything the later blocks were trained on, and these numbers measure the ablation more than the fact. The rows are not comparable to each other either, since vocabularies, depths and residual scales differ. Compare the shape of each checkpoint's curve (`mixer_kl` and `mlp_kl` hold it, one entry per block), or normalize per model. Pythia-70m does not predict `Paris`, and its row measures a different prediction. Mamba's row has no MLP column and no entropy, every block of it being one mixer. ## On another family Add a repo id. `attention_output` is what a block's mixer adds to the stream whether the mixer is `self_attn` or `linear_attn`, `mlp_output` is the MLP's addition, and `mixer` picks the module per block, so a hybrid such as Qwen3.5 runs through the same body with its attention blocks alone in the entropy column. A `model_type` without a family loads with the best-effort `default` family and a warning, which serves `layer_output` but not `attention_output` or `mlp_output`, and its `support()` says so before the sweep reaches it. # Expert ablation Removing one routed expert asks whether a prediction depends on it: zero `expert_weights` wherever `expert_indices == e`, and the slots that chose expert `e` add nothing while the other slots keep their weights, so the token is not renormalized onto its remaining experts and the router is not re-run. One invoke per expert gives a `[block, expert]` grid of changes in log p(target) with one trace per block, the same code on every MoE family ([Mixture of experts](../guide/architectures/mixture-of-experts.md)). Reading the routing rather than editing it is nnsight.net's [MoE Routing](https://nnsight.net/tutorials/tutorials/architectures/moe_routing/) tutorial. ## The grid ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("ibm-granite/granite-3.0-1b-a400m-base", device="cuda", dtype=torch.float32) prompt = "Big Ben is in the city of" ids = model.tokenizer(" London", add_special_tokens=False).input_ids assert len(ids) == 1, model.tokenizer.convert_ids_to_tokens(ids) # one token, or the target is not the word target = ids[0] blocks = [layer.mlp for i, layer in enumerate(model.layers) if model.support(layer=i).get("mlp.expert_weights", "absent") is None] print(len(blocks), blocks[0].num_experts, blocks[0].top_k) # 24 32 8 def sweep(exclude_first_token=False): """[block, expert]: the change in log p(target) when expert e's slots are zeroed.""" effects = torch.zeros(len(blocks), blocks[0].num_experts) for b, moe in enumerate(blocks): rows = [] # bound outside the trace with torch.no_grad(), model.trace() as tracer: with tracer.invoke(prompt): # unedited: the baseline, from the same batch clean = model.logits[0, -1].log_softmax(-1)[target].save() for e in range(moe.num_experts): with tracer.invoke(prompt): chosen = moe.expert_indices == e # [batch, seq, top_k] if exclude_first_token: chosen[:, 0] = False moe.expert_weights[:] = moe.expert_weights.masked_fill(chosen, 0) rows.append(model.logits[0, -1].log_softmax(-1)[target].save()) effects[b] = torch.stack(rows) - clean return effects, float(clean) def top(effects, k=4): for flat in effects.flatten().topk(k, largest=False).indices: b, e = divmod(int(flat), effects.shape[1]) print(f"block {b:2d} expert {e:2d} {float(effects[b, e]):+.3f}") effects, clean = sweep() print(f"clean log p(London) {clean:.3f}; experts no token chose: {int((effects == 0).sum())} of {effects.numel()}") top(effects) # clean log p(London) -0.359; experts no token chose: 205 of 768 # block 5 expert 4 -9.711 # block 0 expert 4 -2.247 # block 5 expert 17 -0.820 # block 4 expert 2 -0.555 ``` `support(layer=i)` picks the blocks with the routing pair. The dense blocks beside the mixture blocks on DeepSeek-V3, GLM-4-MoE, Llama 4 and Jamba have none. `add_special_tokens=False` keeps the BOS token out of the target, and the assertion catches a word the tokenizer splits: This tokenizer splits ` Paris` into `Par` and `is`, and the target is a city it spells in one token. Under several invokes the edit is in place (`expert_weights[:] = ...`), because each invoke is served its own rows of the flat routing tensors, and the baseline is the unedited invoke of the same batch rather than a separate single-prompt trace, since a batch rounds differently from one prompt. Load in float32 for a single-expert sweep: The median effect here is 0.009 nats, and in bfloat16 the clean run alone differs between a batch and a single prompt by 0.018 nats (by 1e-6 in float32). An expert no token of the prompt chose has no slots to zero, and its cell is exactly zero: 205 of the 768 cells here. A zero marks an unused expert, not an unimportant one. ## The first token ```python with torch.no_grad(), model.trace(prompt): chosen = (blocks[5].expert_indices[0] == 4).nonzero().save() # [slot, (position, k)] weights = blocks[5].expert_weights[0].save() print(chosen.tolist(), [round(float(weights[p, k]), 3) for p, k in chosen]) # [[0, 0], [3, 3]] [0.997, 0.073] top(sweep(exclude_first_token=True)[0]) # block 0 expert 4 -2.247 # block 5 expert 17 -0.820 # block 4 expert 2 -0.572 # block 0 expert 22 -0.516 ``` Block 5's expert 4 costs 9.7 nats, and it takes the first token with weight 0.997 and the fourth with 0.073. The effect comes through position 0 alone, and with that position excluded it drops off the list. Rank experts on positions 1 onward, or ablate per position, before reading a large effect as a mechanism. The ranking that survives puts block 0's expert 4 first at 2.2 nats, then block 5's expert 17 and block 4's expert 2 under one nat each, which is the scale of a single expert's contribution to this prediction. An edit to one token's stream also changes that token's routing at later mixture blocks, so an effect includes the rerouting it causes. As with any zero ablation, the number overstates what the expert does when present. Rerouting its slots to another expert, by assigning `expert_indices`, is the milder intervention. ## On another family `expert_weights` and `expert_indices` are `[batch, seq, top_k]` on every MoE family and `num_experts` and `top_k` are read off the mixture, so the grid's shape follows the checkpoint and the code does not change. A mixture's shared expert is not among the routed experts. Ablate it with `shared_expert_output[:] = 0` where `support()` lists it. On JetMoE the indices are computed before the weights, so bind `idx = moe.expert_indices` before reading `expert_weights`, and mask with `idx == e`. # Image pathway An image enters the text model once, as `vision.image_features` scattered into the positions that `vision.image_token_mask` marks, and from then on it is a run of positions that every text-block value covers. The recipe ablates the image at that entry and inside the tower, patches one image's features into another image's run, measures each head's attention onto the image, and edits the image positions of a text block, on a checkpoint asked for the color of a square ([Vision-language](../guide/architectures/vision-language.md)). ## The image as the text model sees it ```python import torch from PIL import Image, ImageDraw from nnterp import StandardizedTransformer model = StandardizedTransformer("llava-hf/llava-interleave-qwen-0.5b-hf", task="image-text-to-text", dispatch=True, device="cuda", attn_implementation="eager") def square(color): image = Image.new("RGB", (384, 384), "white") ImageDraw.Draw(image).rectangle([96, 96, 288, 288], fill=color) return image red, blue = square("red"), square("blue") messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "What color is the square? Answer with one word."}]}] prompt = model.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) def token(word): ids = model.tokenizer(word, add_special_tokens=False).input_ids assert len(ids) == 1, model.tokenizer.convert_ids_to_tokens(ids) # one token, or the target is not the word return ids[0] RED, BLUE = token("Red"), token("Blue") def report(name, probs): top = probs[0].argmax() print(f"{name:32s} top={model.tokenizer.decode([top])!r:8s} P(Red)={probs[0, RED]:.3f} P(Blue)={probs[0, BLUE]:.3f}") with torch.no_grad(), model.trace(prompt, images=[red]): mask = model.vision.image_token_mask.save() # first: it comes off the inputs red_features = model.vision.image_features.save() clean = model.next_token_probs.save() print(mask.shape, int(mask.sum()), red_features.shape) # torch.Size([1, 749]) 729 torch.Size([729, 1024]) report("clean (red)", clean) # clean (red) top='Red' P(Red)=0.984 P(Blue)=0.000 with torch.no_grad(), model.trace(prompt, images=[red]): model.vision.image_features[:] = 0 # the text model receives zeros at the image tokens report("image_features zeroed", model.next_token_probs.save()) # image_features zeroed top='The' P(Red)=0.001 P(Blue)=0.004 ``` The load is `task="image-text-to-text"` (the default text task has no processor, and `model.vision` then serves nothing), the processor's chat template places the image token, and the mask is read first in every trace because it comes off the inputs. `image_features` is `[image_tokens, hidden_size]`, flat over the image tokens of the batch, and `layers[0].input[mask] == image_features` exactly. One image-carrying invoke per trace, so the clean and the ablated runs are two traces, and `red_features` is an ordinary tensor by the time the second runs. With the image gone the model has nothing to answer about, and no color token takes `Red`'s place. ## Inside the tower ```python with torch.no_grad(), model.trace(prompt, images=[red]): model.vision.tower_output[:] = 0 report("tower_output zeroed", model.next_token_probs.save()) # tower_output zeroed top='Red' P(Red)=0.984 P(Blue)=0.000 with torch.no_grad(), model.trace(prompt, images=[red]): model.vision.layers[-1].layer_output[:] = 0 report("tower block -1 zeroed", model.next_token_probs.save()) # tower block -1 zeroed top='White' P(Red)=0.088 P(Blue)=0.077 ``` !!! warning "A tower edit the wrapper never reads changes nothing" Which tower value reaches the projector is the wrapper's choice. llava-interleave feeds the projector the last block's `layer_output` before the tower's final norm, so `tower_output` is computed and discarded and zeroing it leaves `Red` at 0.984. Llava 1.5 reads block -2, and Gemma 3 pools `tower_output`. `model.projector.input` is what the projector receives on every wrapper. Compare it with the value you edited when an edit has no effect. ## One image's features in another's run ```python with torch.no_grad(), model.trace(prompt, images=[blue]): report("clean (blue)", model.next_token_probs.save()) # clean (blue) top='Blue' P(Red)=0.000 P(Blue)=0.980 with torch.no_grad(), model.trace(prompt, images=[blue]): model.vision.image_features[:] = red_features # both images give 729 tokens on this tower report("blue run, red features", model.next_token_probs.save()) # blue run, red features top='Red' P(Red)=0.984 P(Blue)=0.000 ``` The assignment is activation patching across images, the method of nnsight.net's [Activation Patching](https://nnsight.net/tutorials/tutorials/causal_mediation_analysis/activation_patching/) tutorial: The text model sees the red image under the blue image's run and answers for the red one, at the clean run's probability. The assignment needs the same number of image tokens in both runs, which every image gives on a fixed-resolution tower such as this SigLIP. On a variable-resolution tower, such as the Qwen ViT or Pixtral, use images of the same size. To patch part of the image, index the rows, which are the patches in raster order here. ## Attention onto the image ```python masses = [] # made outside the trace with torch.no_grad(), model.trace(prompt, images=[red]): mask = model.vision.image_token_mask for i in range(model.num_layers): p = model.layers[i].self_attn.attention_probabilities masses.append(p[0, :, -1, mask[0]].sum(-1).save()) # [num_heads]: the last token's mass on the image masses = torch.stack(masses) # [layers, num_heads] print([round(float(m), 2) for m in masses.mean(-1)]) # [0.36, 0.12, 0.14, 0.12, 0.11, 0.09, 0.22, 0.2, 0.15, 0.12, 0.19, 0.11, 0.21, 0.21, 0.22, 0.24, 0.24, 0.26, 0.29, 0.29, 0.26, 0.25, 0.2, 0.28] for flat in masses.flatten().topk(3).indices: layer, head = divmod(int(flat), masses.shape[1]) print(f"block {layer} head {head}: {float(masses[layer, head]):.2f}") # block 20 head 9: 0.98 # block 17 head 14: 0.96 # block 19 head 13: 0.91 ``` The pattern is `[batch, num_heads, query, key]` over the whole sequence, so a head's mass on the image is its row summed over the image key columns. The last token puts 0.36 of its attention on the image at block 0 and between 0.1 and 0.3 afterwards, while single heads in blocks 17 to 20 put over 0.9 of theirs there. Those are the heads to score in [Attention patterns](attention-patterns.md). ## The image positions of a text block ```python for k in (4, 8, 12, 16, 20): with torch.no_grad(), model.trace(prompt, images=[red]): mask = model.vision.image_token_mask model.layers[k].layer_output[mask] = 0 # the image rows leaving block k; the text rows untouched report(f"image rows zeroed after block {k}", model.next_token_probs.save()) # image rows zeroed after block 4 top='The' P(Red)=0.006 P(Blue)=0.005 # image rows zeroed after block 8 top='The' P(Red)=0.000 P(Blue)=0.000 # image rows zeroed after block 12 top='White' P(Red)=0.004 P(Blue)=0.011 # image rows zeroed after block 16 top='No' P(Red)=0.062 P(Blue)=0.005 # image rows zeroed after block 20 top='Red' P(Red)=0.895 P(Blue)=0.005 ``` Boolean indexing with the mask writes the image rows and leaves the text rows alone, so the edit asks when the text positions have finished reading the image. The answer survives at 0.895 once the image rows are zeroed after block 20, and not before. The same indexing reads the two sides apart, `layer_output[mask]` as `[729, hidden_size]` and `layer_output[~mask]` as the text rows, the batch flattened in both. ## On another family `vision.image_features`, `vision.image_token_mask` and `vision.layers[i]` mean the same on every wrapper where nnterp names the tower, and `attention_probabilities` is the same pattern on every text family, so the code changes only in the checkpoint, the prompt template and the target words. Zeroing `image_features` is not the whole image on Qwen3-VL, where the first three text blocks re-add it through `layers[k].deepstack_output`. Eager attention over a large tower holds every block's pattern for autograd, so keep the traces under `torch.no_grad()`. # DeltaNet state A gated DeltaNet block keeps one `[key_dim, value_dim]` matrix per head in place of an attention pattern, and that state, with a short convolution window, is all the block carries forward about the prompt. The recipe reads the state entering and leaving a block, patches it between two prompts on a decode step and inside a prompt, and tracks its norm token by token, on Qwen3.5-0.8B, a hybrid with three DeltaNet blocks in every four ([Gated DeltaNet](../guide/architectures/delta-net.md)). nnsight.net's [State-Space Models](https://nnsight.net/tutorials/tutorials/architectures/state_space_models/) tutorial covers the recurrent view of models that carry such a state. ## The state leaving a prompt ```python import torch import nnterp from nnterp import StandardizedTransformer model = StandardizedTransformer("Qwen/Qwen3.5-0.8B", dispatch=True, device="cpu") linear = [i for i, layer in enumerate(model.layers) if getattr(layer, "linear_attn", None) is not None] print(linear) # [0, 1, 2, 4, 5, 6, 8, 9, 10, 12, 13, 14, 16, 17, 18, 20, 21, 22] mixes = [model.layers[i].linear_attn for i in linear] # decided outside the trace prompt = "The Eiffel Tower is in the city of" other = "My favourite food is pizza with" entering = [] # a None bound inside the trace would not survive it with torch.no_grad(), model.trace(prompt): state_in = mixes[0].state_input entering.append(state_in.save() if state_in is not None else None) state_out = mixes[0].state_output.save() clean = model.next_token_probs.save() print(entering[0], state_out.shape, f"{float(state_out.norm()):.3f}") # None torch.Size([1, 16, 128, 128]) 8.596 print(repr(model.tokenizer.decode(clean[0].argmax())), f"{float(clean[0].max()):.3f}") # ' Paris' 0.590 ``` A block has either `self_attn` or `linear_attn`, so the list of recurrent blocks is decided outside the trace. `state_input` is `None` on a fresh prompt and the cached state on a decode step, served as a copy because the cache overwrites its buffer in place with the step's new state. `state_output` is the state after the call's last token, `[batch, heads, key_dim, value_dim]` with `heads` the mixer's value heads (16 here, against the model's 8 attention heads), and it is what the next decode step starts from. ## Patching the state on a decode step ```python donors = [] with torch.no_grad(), model.trace(other): for mix in mixes: donors.append(mix.state_output.save()) n = len(model.tokenizer(prompt).input_ids) with model.generate(prompt, max_new_tokens=4, do_sample=False) as tracer: ids = tracer.result.save() with model.generate(prompt, max_new_tokens=4, do_sample=False) as tracer: for step in tracer.iter[1]: # decode step 1 starts from the other prompt's memory mixes[0].state_input = donors[0] ids_one = tracer.result.save() with model.generate(prompt, max_new_tokens=4, do_sample=False) as tracer: for step in tracer.iter[1]: for mix, donor in zip(mixes, donors): mix.state_input = donor ids_all = tracer.result.save() print(repr(model.tokenizer.decode(ids[0, n:])), repr(model.tokenizer.decode(ids_one[0, n:])), repr(model.tokenizer.decode(ids_all[0, n:]))) # ' Paris, France.' ' Paris, France.' ' Paris.\nThe' ``` Under `generate`, step 0 is the prompt and each later step one token with the cached state entering. `state_input` on step 1 is `state_output` of step 0, and assigning it makes the step continue from another prompt's memory. One block's memory changes nothing visible, since the other 17 recurrent blocks and the six attention blocks still hold the prompt. Swapping every recurrent block's state does change the continuation, from the second new token on. ## The state after every token ```python print(model.support()["linear_attn.states"][0]) # the state after each token is materialized only by the token-by-token kernel; the chunked kernel a prompt runs through carries it between chunks. Call nnterp.route_kernels(model.family, 'torch') before tracing this layer (slower, like attn_implementation='eager') nnterp.route_kernels(model.family, "torch") # process-wide; before the layer's first trace model = StandardizedTransformer("Qwen/Qwen3.5-0.8B", dispatch=True, device="cpu") mixes = [model.layers[i].linear_attn for i in linear] with torch.no_grad(), model.trace(prompt): states = mixes[0].states.save() # [batch, seq, heads, key_dim, value_dim] final = mixes[0].state_output.save() probs = model.next_token_probs.save() assert torch.equal(states[:, -1], final) norms = states.flatten(2).norm(dim=-1)[0] for token, norm in zip(model.tokenizer.convert_ids_to_tokens(model.tokenizer(prompt).input_ids), norms): print(f"{token!r:10} {float(norm):.3f}") # 'The' 2.802 # 'ĠEiff' 3.674 # 'el' 4.180 # 'ĠTower' 4.912 # 'Ġis' 5.792 # 'Ġin' 6.647 # 'Ġthe' 7.837 # 'Ġcity' 8.107 # 'Ġof' 8.596 print(f"{float((probs - clean).abs().max()):.4f}") # 0.0117 ``` A prompt normally runs the chunked kernel, which carries the state between chunks and never materializes it per token, and `states` says so until the family is routed through transformers' token-by-token kernel. `route_kernels` is process-wide and must run before the layer's forward is instrumented, which is why the model is loaded again after it. `route_kernels(model.family, "default")` restores the chunked kernel. The two kernels compute the same rule, and the next-token probabilities agree to bfloat16 rounding. The memory grows monotonically along this prompt, and the last token's state is the one the prompt leaves. ## Patching the state inside a prompt ```python T = 3 patched = [] with torch.no_grad(), model.trace(other): for mix in mixes: patched.append(mix.state_after(T).save()) with torch.no_grad(), model.trace(prompt): before = mixes[0].state_after(T - 1).save() # positions before the write: read first mixes[0].set_state_after(T, patched[0]) after = mixes[0].state_after(T + 1).save() # positions after it: read after for mix, donor in zip(mixes[1:], patched[1:]): # then the later blocks, in forward order mix.set_state_after(T, donor) probs_patched = model.next_token_probs.save() print(torch.equal(before, states[:, T - 1]), torch.equal(after, states[:, T + 1])) # True False print(repr(model.tokenizer.decode(probs_patched[0].argmax())), f"{float(probs_patched[0].max()):.3f}", f"{float(probs_patched[0, clean[0].argmax()]):.3f}") # ' the' 0.055 0.043 ``` `set_state_after(t, value)` writes the state after token `t`, and tokens `t + 1` onward continue from it. Reads follow the forward: in one trace, a position before the write is read before it, a position after it afterwards, and a block's reads come before any write to a later block, which is why block 0's `before` and `after` bracket its own write and the later blocks are written afterwards. With every recurrent block's memory replaced after `Tower`, `Paris` falls from 0.590 to 0.043. The state is not all a block carries forward. A width-4 convolution in front of the kernel still mixes tokens `T - 2` to `T` into the next three, and the patch replaces the recurrent memory, not everything the prompt left. ## On another family Qwen3-Next, Qwen3.5-MoE and OLMo-Hybrid serve the same values on their `linear_attn` blocks, and Kimi-Linear's mixer is the same component with a decay per key channel, so the recipe changes only in the checkpoint. The kernels have to be transformers' torch ones: With `flash-linear-attention` or `causal-conv1d` installed, the kernel has no Python source, every interior value reports unavailable in `support()`, and `attention_output` stays. Under `generate`, `states`, `state_after` and `set_state_after` count from the current call's own first token, and on a decode step they address that step's single token. # Loading and what differs `StandardizedVLLM(repo_id, dispatch=True, ...)` is nnsight's [`VLLM`](https://nnsight.net/vllm/) with the step that `StandardizedTransformer` adds before the load: it reads the checkpoint's config, takes the family from `model_type` under `nnterp.families.vllm`, and hands that family's aliases and envoy classes to nnsight. The [names](../guide/names.md), the values and their [layouts](../guide/values/layouts.md) are those of `StandardizedTransformer` with a batch axis of 1. A block written against `model.layers[i].layer_output[:, -1]` runs on either engine. ```python import torch from nnterp import StandardizedVLLM model = StandardizedVLLM("Qwen/Qwen2.5-0.5B", dispatch=True, dtype="float32", gpu_memory_utilization=0.3) print(model.family.__name__, model.num_layers, model.hidden_size, model.num_heads, model.num_kv_heads) block = model.layers[12] with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): x = block.layer_input.save() attn = block.self_attn.attention_output.save() mlp = block.mlp.mlp_output.save() out = block.layer_output.save() logits = model.logits.save() print(x.shape, attn.shape, mlp.shape, out.shape) print(logits.shape, repr(model.tokenizer.decode(logits[0, -1].argmax()))) print((x + attn + mlp - out).abs().max().item()) # nnterp.families.vllm.qwen2 24 896 14 2 # torch.Size([1, 7, 896]) torch.Size([1, 7, 896]) torch.Size([1, 7, 896]) torch.Size([1, 7, 896]) # torch.Size([1, 1, 151936]) ' Paris' # 0.0 ``` The load is in float32 so that the comparison at the end of this page holds at float32's resolution. Without `dtype` the weights load in the config's dtype, bfloat16 here. `gpu_memory_utilization` is vLLM's share of the card, which the reader sets for their own. ## Arguments `family=`, `rename=` and `envoys=` are nnterp's own and merge as on [Loading](../guide/loading.md#arguments), `envoys=` over the family's `ENVOYS` and nnsight's envoys for vLLM's parallel layers. Everything else, such as `dispatch`, `dtype`, `gpu_memory_utilization`, `max_model_len`, `tensor_parallel_size`, `mode` and `taps`, goes to `VLLM` as documented at [Loading models](https://nnsight.net/vllm/loading/), and the full signature is under the Documentation tab ([standardized_vllm](../documentation/standardized_vllm.md)). There is no `attn_implementation`: vLLM's attention kernels are its own, and the values inside attention are served another way ([below](#the-attention-interior)). In a `.py` file the engine must be built under a `__main__` guard, because vLLM spawns its worker and the spawn re-imports the main module ([nnsight](https://nnsight.net/vllm/loading/#main-guard)). The block runs in that worker, where nnterp has to be importable. ## `layer_input` and `layer_output` On most families a block is called `forward(positions, hidden_states, residual)` and returns `(hidden_states, residual)`, because vLLM fuses each residual add into the norm that follows, and the residual stream is the pair's sum ([Locations](https://nnsight.net/vllm/locations/#tuples-and-the-residual-stream)). Natively, `layers[i].input` is the positions and `layers[i].output` the pair. `layer_input` and `layer_output` are the stream on every family and on both engines. ```python with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): positions = block.input.clone().save() hidden = block.output[0].clone().save() residual = block.output[1].clone().save() stream = block.layer_output.save() print(positions) print(hidden.shape, residual.shape, torch.equal(hidden + residual, stream[0])) # tensor([0, 1, 2, 3, 4, 5, 6], device='cuda:0') # torch.Size([7, 896]) torch.Size([7, 896]) True ``` ## One sequence, batch axis 1 A request is one sequence. vLLM serves the request's own rows, `[tokens, ...]`, and nnterp puts the batch axis back, always 1. Under `tracer.iter` step 0 is the prefill, every prompt token, and each later step one token. `logits` is `[1, 1, vocab_size]`, because the engine computes the last position only. `logits[:, -1]` and `next_token_probs` read the same on both engines, and nnsight's `model.samples` is the token drawn from them ([Generation](https://nnsight.net/vllm/generation/)). Several prompts are several invokes, one prompt each, and a name saved in each comes back as a list. ```python import nnsight with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=3, ignore_eos=True) as tracer: shapes = nnsight.save([]) for step in tracer.iter[:3]: shapes.append((tuple(block.layer_output.shape), tuple(model.logits.shape))) result = tracer.result.save() print(shapes) print(repr(result.outputs[0].text)) # [((1, 7, 896), (1, 1, 151936)), ((1, 1, 896), (1, 1, 151936)), ((1, 1, 896), (1, 1, 151936))] # ' Paris, France' ``` A trace is a generation request, and the sampling settings go on `trace` or `invoke`. With none, vLLM's defaults apply, 16 tokens at temperature 1.0. The blocks on this page pass `temperature=0.0, max_tokens=1` for one deterministic step. A block without `tracer.iter` runs on the prefill whatever `max_tokens` is. ```python with model.trace("The Eiffel Tower is in") as tracer: result = tracer.result.save() print(len(result.outputs[0].token_ids)) # 16 ``` ## Private copies A tensor vLLM serves is the model's live buffer, which a later kernel may rewrite ([Clone what you keep](https://nnsight.net/vllm/locations/#clone-what-you-keep)). Every nnterp value on this engine is a private copy, handed back to the model when the block moves on: a saved value stays what it was read as without a `.clone()`, and an in-place edit and an assignment both reach the model. An assigned value must keep the shape it was served with, because the rows are spliced into a step the engine runs for every request in flight. nnterp refuses another height in the block that wrote it, the request fails, and the engine carries on. An error raised in the worker arrives as a `RuntimeError` naming the original type. ```python with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): clean = model.logits.save() with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): block.layer_output[:, -1] *= 0.5 model.layers[13].mlp.mlp_output = model.layers[13].mlp.mlp_output * 0 edited = model.logits.save() print((edited - clean).abs().max().item()) # 10.073400497436523 try: with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): block.layer_output = block.layer_output[:, :-1] except RuntimeError as error: print(str(error).splitlines()[0]) # ValueError: model.model.layers.12.layer_output is (1, 7, 896) on this request and cannot be replaced by a value of shape (1, 6, 896); write into the rows you mean instead (``value[:, positions] = ...``) ``` ## The attention interior `attention_queries`, `attention_keys`, `attention_values` and `attention_head_outputs` are what enters and leaves vLLM's attention layer, in the layouts of [Attention](../guide/values/attention.md), and writable like any value. The scores and the pattern are computed inside that layer's kernel, and nnterp recomputes them from the queries and keys with the layer's own scale, softcap and sliding window. They are read-only, and there on the prefill only, since a decode step holds one token's keys ([What is unavailable](unavailable.md#the-pattern-on-a-decode-step)). ```python attn = block.self_attn with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): q = attn.attention_queries.save() k = attn.attention_keys.save() v = attn.attention_values.save() scores = attn.attention_scores.save() pattern = attn.attention_probabilities.save() heads = attn.attention_head_outputs.save() print(q.shape, k.shape, v.shape) print(scores.shape, pattern.shape, heads.shape) print(torch.allclose(scores.float().softmax(-1), pattern), pattern[0, 0].sum(-1)) # torch.Size([1, 14, 7, 64]) torch.Size([1, 2, 7, 64]) torch.Size([1, 2, 7, 64]) # torch.Size([1, 14, 7, 7]) torch.Size([1, 14, 7, 7]) torch.Size([1, 7, 14, 64]) # True tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000], # device='cuda:0') ``` ## The same values on both engines The suite under `tests/vllm_families/` holds every value the engine serves against `StandardizedTransformer`'s on the same checkpoint and prompt, both in float32. On Qwen2.5-0.5B, block 12's stream agrees to 1e-5 of its largest magnitude, the pattern to 1e-3, the head outputs to 4e-3 and the logits to 2e-3. The two engines' attention kernels sum a sharp softmax differently, and that difference is what the later values carry. Compare at 5e-3 of a value's scale, and expect more from a bfloat16 load. ```python from nnterp import StandardizedTransformer hf = StandardizedTransformer("Qwen/Qwen2.5-0.5B", device="cuda", dtype=torch.float32, attn_implementation="eager") with torch.no_grad(), hf.trace("The Eiffel Tower is in"): hf_pattern = hf.layers[12].self_attn.attention_probabilities.save() hf_heads = hf.layers[12].self_attn.attention_head_outputs.save() hf_out = hf.layers[12].layer_output.save() hf_logits = hf.logits[:, -1:].save() def gap(vllm_value, hf_value): """The largest difference, relative to the transformers value's largest magnitude.""" return ((vllm_value.cpu() - hf_value.cpu()).abs().max() / hf_value.cpu().abs().max()).item() print(out.shape == hf_out.shape, logits.shape == hf_logits.shape) print(f"{gap(out, hf_out):.1e} {gap(pattern, hf_pattern):.1e} {gap(heads, hf_heads):.1e} {gap(logits, hf_logits):.1e}") print(logits.argmax().item() == hf_logits.argmax().item()) # True True # 1.2e-05 9.5e-04 4.1e-03 1.9e-03 # True ``` The values the engine cannot serve are on [What is unavailable](unavailable.md). # What is unavailable `model.support()` on a vLLM engine reports the values the engine cannot serve, each with its reason as a string, and `None` for the rest, as on transformers ([Availability](../guide/values/availability.md)). On a checkpoint with a gated MLP, such as Qwen2.5-0.5B, there are three. ```python from nnterp import StandardizedVLLM model = StandardizedVLLM("Qwen/Qwen2.5-0.5B", dispatch=True, dtype="float32", gpu_memory_utilization=0.3) for name, reason in model.support().items(): if reason: print(name) print(" ", reason if isinstance(reason, str) else reason[0]) # lm_head_output # vLLM never calls lm_head: its logits processor projects with lm_head's weight and applies any softcap and scale in the same call, on the last position only; read logits # attention_mask # a vLLM request is one unpadded sequence; there is no mask # mlp.mlp_activation # a gated MLP on vLLM: act(gate) is never a module's output (gate_up_proj feeds act_fn, an act-and-multiply kernel, whose output is already act(gate) * up), so it is not served; mlp_neurons is the product ``` Reading or writing one of them raises `nnterp.Unavailable` with the same reason at that line, as on transformers. Because the block runs in the engine's worker, the error arrives as a `RuntimeError` naming `Unavailable`. ## A gated MLP's `mlp_activation` vLLM fuses the gate and up projections into one `gate_up_proj`, and its `act_fn` (`SiluAndMul`, `GeluAndMul`) computes `act(gate) * up` in one kernel, and `act(gate)` alone is never a tensor. `mlp_neurons`, the down projection's input, is served and is that product. ```python mlp = model.layers[12].mlp with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): neurons = mlp.mlp_neurons.save() print(neurons.shape, type(mlp.act_fn._module).__name__) # torch.Size([1, 7, 4864]) SiluAndMul try: with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): activation = mlp.mlp_activation.save() except RuntimeError as error: print(str(error).splitlines()[0][:80]) # Unavailable: model.model.layers.12.mlp.mlp_activation is not available: a gated ``` A plain MLP, such as GPT-2's, serves both as the same tensor. Both are read at the down projection: vLLM builds one activation module by name and shares it between every block, and that module's output cannot tell the blocks apart. ```python gpt2 = StandardizedVLLM("openai-community/gpt2", dispatch=True, dtype="float32", gpu_memory_utilization=0.15) print(sorted(name for name, reason in gpt2.support().items() if reason)) # ['attention_mask', 'lm_head_output'] ``` ## `lm_head_output` vLLM never calls `lm_head`. Its logits processor projects with the head's weight and applies any softcap and scale in the same call, for the last position only. Read `logits`, or run `project_on_vocab` on a stream: on the last block's stream it is `logits`, as on transformers. The module is still on the tree. `project_on_vocab` and `get_topk_closest_tokens` call the engine's modules and work inside a trace only. ```python with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): lens = model.project_on_vocab(model.layers[-1].layer_output[:, -1:]).save() logits = model.logits.save() print(lens.shape, (lens - logits).abs().max().item()) # torch.Size([1, 1, 151936]) 4.76837158203125e-06 ``` ## The pattern on a decode step `attention_scores` and `attention_probabilities` are recomputed from the step's queries and keys ([The attention interior](index.md#the-attention-interior)). On the prefill the step holds every key. On a decode step the new token's key is the only one served, and the earlier ones are in vLLM's cache, which only its kernel reads, and the read raises before anything runs. A trace without `tracer.iter` runs its block on the prefill whatever `max_tokens` is, and `support()`, which runs outside any step, lists both as available. Both are read-only, since the kernel never takes a pattern: to change what a head attends to, edit its queries or keys, and to change what it wrote, edit `attention_head_outputs`. ```python attn = model.layers[12].self_attn with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=3, ignore_eos=True): pattern = attn.attention_probabilities.save() print(pattern.shape) # torch.Size([1, 14, 7, 7]) try: with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=2, ignore_eos=True) as tracer: for step in tracer.iter[1]: pattern = attn.attention_probabilities.save() except RuntimeError as error: print(str(error).splitlines()[0]) # Unavailable: model.model.layers.12.self_attn.attention_probabilities is not available: recomputed from this step's queries and keys, and on a decode step the earlier tokens' keys are in vLLM's cache, which only its kernel reads; read it on the prefill try: with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): attn.attention_probabilities = attn.attention_probabilities * 0 except RuntimeError as error: print(str(error).splitlines()[0]) # AttributeError: attention_probabilities is derived and read-only ``` ## `attention_mask` and the read-only inputs A request is one unpadded sequence with no mask, and `input_size` is `[1, tokens]`. `input_ids` and `input_size` are read-only on this engine, and `next_token_probs` is derived from `logits` on both engines. Assign `logits` to change what the sampler sees. ```python with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): ids = model.input_ids.save() size = model.input_size.save() print(ids, size) # tensor([[ 785, 468, 3092, 301, 21938, 374, 304]], device='cuda:0', # dtype=torch.int32) torch.Size([1, 7]) try: with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): model.next_token_probs = model.next_token_probs * 0 except RuntimeError as error: print(str(error).splitlines()[0]) # AttributeError: next_token_probs is derived from the logits and cannot be assigned; assign model.logits instead ``` ## On some families A mixture of experts serves `mlp_output` and nothing inside it: the experts run in vLLM's fused MoE kernel, so `mlp_activation` and `mlp_neurons` are unavailable, and the router has no values on this engine. DeepSeek-V2, DeepSeek-V3, Kimi-K2 and GLM-4.7-Flash run latent attention in a kernel of their own, so the six values of the attention interior are unavailable there while `attention_output` is served. The two reasons are constants, which `support()` prints for the affected values. ```python from nnterp.components.vllm.mlp import ROUTED from nnterp.families.vllm.deepseek_v2 import LATENT print(ROUTED) print(LATENT) # routed experts, inside vLLM's fused MoE kernel: there is no single activation or down projection # vLLM runs DeepSeek's latent attention in its own kernel, on compressed keys and values; not mapped on vLLM ``` ## Methods and nnsight's limits `model.order` and `model.rank` are measured on `StandardizedTransformer` only. Gradients, `.source` inside a kernel and `scan` are nnsight's limits on this engine ([Capabilities and limits](https://nnsight.net/vllm/capabilities/#not-on-vllm)). `skip_layers` has to answer for every row of a step, and a step holds whatever requests the engine batched together: it is safe on a request that runs alone. To ablate a block on a shared engine, write its contributions to zero instead. ```python try: model.order(12) except NotImplementedError as error: print(error) # model.order is not supported on the vllm engine; it is measured on a StandardizedTransformer ``` Which families run on this engine, and what each changes, is on [Families on vLLM](families.md). # Families on vLLM A vLLM family is one module under `nnterp/families/vllm/`, named after the `model_type` and written against vLLM's own implementation of that architecture. `nnterp.families.known("vllm")` lists the 60 shipped ones, and `lookup(model_type, engine="vllm")` resolves a type as the transformers lookup does ([Registering](../developing/registering.md#the-lookup-order)) with one difference. The best-effort `default` family is written against transformers' modules, and a `model_type` with no vLLM family raises `UnsupportedFamily` instead of falling back to it. ```python from nnterp import UnsupportedFamily, families on_vllm = families.known("vllm") print(len(on_vllm), len(families.known())) print(families.lookup("qwen2", engine="vllm").__name__, families.lookup("qwen2").__name__) try: families.lookup("diffllama", engine="vllm") except UnsupportedFamily as error: print(str(error)[:90]) # 60 98 # nnterp.families.vllm.qwen2 nnterp.families.qwen2 # no standardization for model_type 'diffllama' on vllm; known: ['afmoe', 'apertus', 'arcee' ``` ## The module A vLLM family declares what a transformers family declares, `RENAME`, `Layer`, `Attention`, `Mlp` and `ENVOYS` ([Adding a family](../developing/adding-a-family.md)), with `ENVOYS` keyed on vLLM's module classes and the three envoys subclassing the bases in `nnterp.components.vllm`, which read each value where vLLM keeps it and serve it as a private `[1, tokens, ...]` copy. Qwen2's module is the whole of `nnterp/families/vllm/qwen2.py`: ```python """Qwen2 and Qwen2.5 on vLLM (``vllm.model_executor.models.qwen2``). Llama's names and Llama's fused block, in classes of its own. """ from vllm.model_executor.models.qwen2 import Qwen2Attention, Qwen2DecoderLayer, Qwen2MLP from nnterp.components.vllm import Attention, FusedLayer, Mlp RENAME = { "model.embed_tokens": "embed_tokens", "model.layers": "layers", "model.norm": "norm", } class Layer(FusedLayer): """Qwen2's decoder block: ``forward(positions, hidden_states, residual) -> (hidden_states, residual)``.""" class Attention(Attention): """Qwen2's attention; its output is what the block adds, so the base holds.""" class Mlp(Mlp): """Qwen2's MLP; its output is what the block adds, so the base holds.""" #: Module type -> Envoy subclass, for nnsight's ``envoys=``. ENVOYS = {Qwen2DecoderLayer: Layer, Qwen2Attention: Attention, Qwen2MLP: Mlp} ``` What the family has to say is which convention its block follows, because the return type does not tell. `FusedLayer` is a block that takes and returns `(hidden_states, residual)` whose sum is the stream, with `HIDDEN` and `RESIDUAL` saying where the halves sit in the call. `Layer` is a block called with the stream, alone or after the positions (`STREAM`), that returns it alone or with something beside it (`returns_tuple`). EXAONE 4's block is called `forward(positions, hidden_states, residual)` and returns a pair, and is a `Layer` all the same, because the pair's first element is the whole stream. Where what the block adds is not a module's output, the family redefines the value with `Flat`, the descriptor for a `[tokens, ...]` tensor: the post-norm's output on Gemma-2 and OLMo-2, the module's output times `residual_multiplier` on Granite. The signatures are under the Documentation tab ([vllm components](../documentation/components/vllm/index.md)). ```python from nnterp.components.vllm import FusedLayer, Layer from nnterp.components.vllm.transformers_backend import Layer as BackendLayer from nnterp.families.vllm import exaone4, gpt2, llama, olmo2, phi print(issubclass(llama.Layer, FusedLayer), llama.Layer.HIDDEN, llama.Layer.RESIDUAL) print(issubclass(gpt2.Layer, Layer), gpt2.Layer.STREAM, gpt2.Layer.returns_tuple) print(issubclass(phi.Layer, Layer), phi.Layer.STREAM, phi.Layer.returns_tuple) print(issubclass(exaone4.Layer, Layer), exaone4.Layer.STREAM, exaone4.Layer.returns_tuple) print(issubclass(olmo2.Layer, BackendLayer), list(olmo2.ENVOYS)[0].__module__) # True (1, 'hidden_states') (2, 'residual') # True 0 False # True 1 False # True 1 True # True transformers.models.olmo2.modeling_olmo2 ``` ## The transformers backend vLLM has no implementation of its own for OLMo, OLMo 2, SmolLM3, StarCoder2 and GPT-BigCode and runs transformers' modules inside the engine. Those five families key `ENVOYS` on transformers' classes and build on `nnterp.components.vllm.transformers_backend`, where the activations already carry the batch axis and the engine's attention layer, which the backend keeps outside the module tree, is mounted as `self_attn.attn` inside a trace, and the queries, keys, values and head outputs are then served as on every other vLLM family. `.source` does not work on these modules, because the backend rewrites their forwards. Passing the module as `family=` is how a family is tried before it ships, as on transformers ([Checking the module](../developing/adding-a-family.md#checking-the-module-before-it-ships)): ```python import torch from nnterp import StandardizedVLLM from nnterp.families.vllm import olmo2 model = StandardizedVLLM("allenai/OLMo-2-0425-1B", family=olmo2, dispatch=True, dtype="float32", gpu_memory_utilization=0.3) print(model.family is olmo2, type(model._module).__name__, type(model.layers[0]._module).__name__) print(sorted(name for name, reason in model.support().items() if reason)) block = model.layers[8] with model.trace("The Eiffel Tower is in", temperature=0.0, max_tokens=1): raw = block.input.clone().save() x = block.layer_input.save() q = block.self_attn.attention_queries.save() pattern = block.self_attn.attention_probabilities.save() attn = block.self_attn.attention_output.save() mlp = block.mlp.mlp_output.save() out = block.layer_output.save() logits = model.logits.save() print(raw.shape, torch.equal(raw, x), q.shape, pattern.shape) print((x + attn + mlp - out).abs().max().item(), repr(model.tokenizer.decode(logits[0, -1].argmax()))) # True TransformersForCausalLM Olmo2DecoderLayer # ['attention_mask', 'attention_norm_output', 'lm_head_output', 'mlp.mlp_activation', 'mlp_norm_output'] # torch.Size([1, 7, 2048]) True torch.Size([1, 16, 7, 128]) torch.Size([1, 16, 7, 7]) # 0.0 ' Paris' ``` On OLMo 2's post-norm block, `attention_norm_output` and `mlp_norm_output` are unavailable with the post-norm reasons, as on transformers. ## Adding one Write `nnterp/families/vllm/.py` as above, reading vLLM's forward of the block before choosing the base. Then add `tests/vllm_families/test_vllm_.py`, a `VLLMFamilySuite` subclass naming a checkpoint that both engines load. The suite runs the checkpoint through `StandardizedTransformer` in float32, builds the engine in float32, and holds every value it serves against the transformers value, the contribution identity and the edits included. Each class builds one engine, and each file runs in its own process. OLMo 2's test file is the whole of it: ```python """OLMo 2 on vLLM's transformers backend, against the transformers engine: transformers' modules, post-norm.""" from vllm_suite import LLAMA_ROWS, VLLMFamilySuite from nnterp.families.vllm import olmo2 class TestVLLMOlmo2(VLLMFamilySuite): REPO = "allenai/OLMo-2-0425-1B" FAMILY = olmo2 UNAVAILABLE = VLLMFamilySuite.UNAVAILABLE | {"attention_norm_output", "mlp_norm_output"} NATIVE = LLAMA_ROWS MEMORY = 0.3 ``` ``` tests/vllm_families/run.sh olmo2 ``` The checkpoint has to be one vLLM runs: its attention kernels need a head width of at least 16, and 32 for the float32 one, and most of its MLPs take SiLU only, which rules out most of the tiny random checkpoints the transformers suite pins. Name a small real checkpoint where one is published, and otherwise a tiny one with heads at least 32 wide, or a copy of one resized for vLLM that the test writes once (`test_vllm_hyperclovax.py`). A family kept outside the package registers with `nnterp.families.register(module, engine="vllm")`, and a family registered without `engine` is a transformers family only. ## Where the engines differ vLLM's implementation of a family is not always transformers' model. Where the two disagree, nnterp serves what vLLM computes, and the family's test file aligns the two engines for the comparison and says how. On vLLM 0.27.1 and transformers 5.15: | family | difference | | --- | --- | | GLM (`glm`) | vLLM rotates half of each head whatever `partial_rotary_factor` says; GLM-Edge's config says all of it | | Ministral 3 | vLLM's YaRN ignores `mscale` and `mscale_all_dim` and scales every query and key by 1.277 at factor 16; an HF-format checkpoint needs `hf_overrides={"rope_parameters": {..., "apply_yarn_scaling": False}}` | | GLM-4.7-Flash (`glm4_moe_lite`) | the latent norms use the config's `rms_norm_eps` where transformers fixes 1e-6 | | PhiMoE | the logits leave out `lm_head.bias`, which vLLM loads and never adds | | StableLM | vLLM builds the sequential block only, ignoring `use_parallel_residual` and `qk_layernorm` | | FlexOlmo | the experts run SiLU whatever `hidden_act` says | | Llama 4 | vLLM permutes each head's query and key channels at load, which the values undo; the attention temperature tuning follows the generation config, not the model config | | DBRX | transformers reads the rotary base from `rope_parameters` only, where vLLM reads the checkpoint's `attn_config.rope_theta`; vLLM 0.27.1 loads no DBRX checkpoint | | the kernels | given the same queries, keys and values, the two engines' float32 attention kernels differ by about 1e-3 of the head outputs' scale, and by 2.4e-2 on a softmax as sharp as Qwen2.5-0.5B's first block | Each family's page in the [Documentation](../documentation/families/vllm/index.md) tab carries its own notes. # Developing - [Adding a family](adding-a-family.md): one module and one test file for a new `model_type`. - [Overriding a value](overriding-values.md): a standard value redefined where the base does not hold. - [Custom values](custom-values.md): a value of your own, through `envoys=`. - [Finding source ops](finding-source-ops.md): the operation names a `source` path needs. - [Registering](registering.md): a family from outside the package. - [Architecture](internals/architecture.md): how a `model_type` becomes a renamed tree with values. - [Descriptors](internals/descriptors.md): how `EProperty` resolves a path at run time. - [Recurrent mixers](internals/recurrent-mixers.md): how `RecurrentMixer` reaches a kernel's values and state. - [Vision](internals/vision.md): how the vision side of a wrapper is standardized. - [Testing](testing.md): the suite and what it asserts of a family. - [transformers compatibility](transformers-compatibility.md): what a transformers release can move. - [Contributing](contributing.md): the workflow for a change. # Adding a family A family is one module under `nnterp/families/`, named after the checkpoint's `config.model_type`, with a `RENAME` dict from native names to standard names, three envoy subclasses, and an `ENVOYS` dict keying those on the family's module classes. Its test file is a `FamilySuite` subclass with a pinned tiny checkpoint. The file name is the whole registration: `nnterp.families.lookup` imports `nnterp.families.` the first time a checkpoint of that type is loaded, and nothing else in nnterp changes. A shipped family also has a page in the [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/), an entry under `encyclopedia/entries/` written as `encyclopedia/README.md` describes. ## The module Llama is the family the standard vocabulary is taken from, and its module changes nothing but the container names. This is `nnterp/families/llama.py` without the names of the multimodal wrappers around a Llama text model, which the shipped file also carries: ```python """Llama (``LlamaForCausalLM``), the family the standard vocabulary is taken from.""" from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaMLP from nnterp.components import Attention, Layer, Mlp RENAME = { "model.embed_tokens": "embed_tokens", "model.layers": "layers", "model.norm": "norm", } class Layer(Layer): """Llama's decoder block; returns a bare tensor, so the base holds.""" class Attention(Attention): """Llama's attention; the shared eager forward and the residual added in the block, so the base holds.""" class Mlp(Mlp): """Llama's MLP; the residual is added in the block, so the base holds.""" #: Module type -> Envoy subclass, for nnsight's ``envoys=``. ENVOYS = {LlamaDecoderLayer: Layer, LlamaAttention: Attention, LlamaMLP: Mlp} ``` nnsight resolves each `RENAME` key relative to every envoy in the tree ([module access](https://nnsight.net/features/6_modules/) on nnsight.net). A key with several components, such as `"model.layers"`, resolves from the root, which is what lifts the containers out of the inner model: The alias is `model.layers`, not `model.model.layers`. A single-component key binds on every envoy with a child of that name, so GPT-2's `"ln_1": "input_layernorm"`, `"attn": "self_attn"` and `"ln_2": "post_attention_layernorm"` rename every block at once. A key that resolves nowhere is skipped. The three classes are defined even where nothing changes, because the suite asserts that the block envoys are exactly the family's `Layer`, `Attention` and `Mlp`, and each docstring records why the base holds. A block that returns `(hidden_states, ...)` rather than the tensor, as GPT-J's, BLOOM's, MPT's and Falcon's do, sets `returns_tuple = True` on its `Layer`. `layer_output` unwraps the tuple either way, and the flag is what `skip_layers` needs to hand back the right shape. ## Checking the module before it ships Until the file is in the package, pass the module to one load as `family=` ([Registering](registering.md)). SmolLM2 is a Llama checkpoint (`model_type` `llama`). Loaded eager, nothing on it is unavailable, and the contributions add up to the block's output: ```python import torch from nnterp import StandardizedTransformer import llama # the module above, saved as llama.py beside this script model = StandardizedTransformer("HuggingFaceTB/SmolLM2-135M", family=llama, attn_implementation="eager") print(model.family is llama) # True print({name: reason for name, reason in model.support().items() if reason}) # {} with model.trace("The Eiffel Tower is in"): block = model.layers[5] x = block.layer_input.save() attn = block.self_attn.attention_output.save() mlp = block.mlp.mlp_output.save() out = block.layer_output.save() print(torch.equal(x + attn + mlp, out)) # True ``` ## Redefining a value Because Gemma-2's block norms each sublayer's output before adding it, what the attention adds to the stream is the post-norm's output, not the module's. Its `Attention` redefines `attention_output` on that sibling norm, keeping the name, the `-> Residual` annotation and the description, and changing only where the value is read: ```python from nnterp.components import Attention, EProperty, Residual class Attention(Attention): @EProperty("../post_attention_layernorm.output", description="What the attention adds to the residual stream: the post-attention norm's output") def attention_output(self, value) -> Residual: return value ``` [Overriding a value](overriding-values.md) has every form the path takes, and the markers for a value a family does not have. ## A size the config spells its own way The root's sizes (`num_layers`, `hidden_size`, `vocab_size`, `num_heads`, `num_kv_heads`, `head_dim`, `qk_head_dim` and `intermediate_size`) are read off the config by Llama's rule, which for `intermediate_size` is `config.intermediate_size`. GPT-2's config has no such key. Its module defines a function of the same name taking the model, and the root calls it in place of the rule: ```python # nnterp/families/gpt2.py def intermediate_size(model: "StandardizedTransformer") -> int: """The MLP width is ``n_inner``, ``None`` meaning four times the hidden size; the config's ``intermediate_size`` is never read by the model.""" return model.config.n_inner or 4 * model.hidden_size ``` ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2") print(model.config.n_inner, model.intermediate_size, model.layers[0].mlp.intermediate_size) # None 3072 3072 ``` A block's own `mlp.intermediate_size` is read off the module rather than the config. Falcon's module spells `num_kv_heads` the same way. A model that scales or caps its logits after `lm_head` defines `def project_on_vocab(model, hidden)` in the module (Cohere multiplies by `logit_scale`), which is bound in place of the root's, so the logit lens on the last block still equals `logits`. ## The test file One file per family under `tests/families/`, subclassing `FamilySuite` from `tests/families/suite.py`. Llama's is the minimum: ```python """Llama, end to end: the family the vocabulary is taken from.""" from suite import FamilySuite, LLAMA_ROWS from nnterp.families import llama class TestLlama(FamilySuite): REPO = "hf-internal-testing/tiny-random-LlamaForCausalLM" FAMILY = llama NATIVE = LLAMA_ROWS ``` `NATIVE` maps standard paths to native paths, and `rows()` builds it from the container and the native names: `LLAMA_ROWS` is `rows("model", "layers", "embed_tokens", "norm")`, and GPT-2's is `rows("transformer", "h", "wte", "ln_f", attn="attn", ln1="ln_1", ln2="ln_2")`. A quirk of the checkpoint is a class attribute: | attribute | meaning | | --- | --- | | `EXPECTED_UNAVAILABLE` | `support()` key to a substring of the reason, for the values this checkpoint lacks; every other value must report `None` | | `REFUSES_IN_PLACE_QKV` | torch refuses in-place edits on queries that are views out of one `split` (GPT-2, GPT-BigCode, MPT) | | `ATTENTION_NORM`, `MLP_NORM` | the block's own module whose output enters each sublayer, where it is not `input_layernorm` and `post_attention_layernorm` (Gemma-2's `MLP_NORM` is `pre_feedforward_layernorm`) | | `LOAD_KWARGS` | extra load arguments the checkpoint needs | Run it: ``` pytest tests/families/test_llama.py -q # 41 passed, 12 skipped, 2 warnings in 12.62s ``` The suite loads the checkpoint eager and asserts, among other things, that: - every alias reaches its native module, and the block envoys are the family's classes - `support()` matches what reads and what raises - `layer_input + attention_output + mlp_output == layer_output` on every block - the standardized model equals a raw `TransformersModel` - every source-located value resolves on every attention block - each value's tensor matches its layout, and the sizes match the weights [Testing](testing.md) lists all of it. The skipped tests are the mixture-of-experts ones, which a dense model has none of. A failure in `test_every_source_value_resolves_on_every_layer` names the operation the forward does not have, and [Finding source ops](finding-source-ops.md) is how to find the one it does. # Overriding a value When a family's forward puts a standard value somewhere other than where the base `Layer`, `Attention` or `Mlp` reads it, the family's subclass redefines the value under the same name with an `EProperty` keyed on a path from the host envoy, and the name keeps its meaning on that family too. The name, the layout annotation and the description stay; only the path changes. [EProperty's signature](../documentation/components/eproperty.md) is in the Documentation tab. ## A sibling module's output Gemma-2's block is a sandwich, `x + post_attention_layernorm(attn(input_layernorm(x)))` and then the same around the MLP. Each sublayer adds its post-norm's output, and the family points `attention_output` and `mlp_output` at the sibling norms, and `layer_mid` and `mlp_norm_output` at the MLP's pre-norm, which is not `post_attention_layernorm` here: ```python """Gemma 2 (``Gemma2ForCausalLM``): Llama's names, a sandwich block.""" from transformers.models.gemma2.modeling_gemma2 import Gemma2Attention, Gemma2DecoderLayer, Gemma2MLP from nnterp.components import Attention, EProperty, Layer, Mlp, Residual RENAME = { "model.embed_tokens": "embed_tokens", "model.layers": "layers", "model.norm": "norm", } class Layer(Layer): @EProperty("pre_feedforward_layernorm.input", description="The residual stream after the attention add") def layer_mid(self, value) -> Residual: return value @EProperty("pre_feedforward_layernorm.output", description="The normed stream entering the MLP") def mlp_norm_output(self, value) -> Residual: return value class Attention(Attention): @EProperty("../post_attention_layernorm.output", description="What the attention adds to the residual stream: the post-attention norm's output") def attention_output(self, value) -> Residual: return value class Mlp(Mlp): @EProperty("../post_feedforward_layernorm.output", description="What the MLP adds to the residual stream: the post-feedforward norm's output") def mlp_output(self, value) -> Residual: return value ENVOYS = {Gemma2DecoderLayer: Layer, Gemma2Attention: Attention, Gemma2MLP: Mlp} ``` With the module passed as `family=` for one load ([Registering](registering.md)), the redefined value is the norm's output and the contribution identity holds, where the module's own return would not. The smallest public Gemma-2 checkpoint is a tiny random one: ```python import torch from nnterp import StandardizedTransformer import gemma2 # the module above, saved as gemma2.py beside this script model = StandardizedTransformer("trl-internal-testing/tiny-Gemma2ForCausalLM", family=gemma2, attn_implementation="eager") block = model.layers[0] with model.trace("The Eiffel Tower is in"): x = block.layer_input.save() returned = block.self_attn.output.save() attn = block.self_attn.attention_output.save() normed = block.post_attention_layernorm.output.save() mlp = block.mlp.mlp_output.save() out = block.layer_output.save() print(torch.equal(attn, normed), torch.equal(attn, returned[0])) # True False print(torch.equal(x + attn + mlp, out), torch.allclose(x + returned[0] + mlp, out, atol=1e-5)) # True False ``` A `../` segment steps to the parent module, and what follows it is a native name, because above the host the path is joined onto the host's own path as a string: On GPT-2's attention, `"../post_attention_layernorm.output"` names a module that does not exist and the read fails with `OutOfOrderError`, while `"../ln_2.output"` works. Below the host, and from the root after a leading `/`, aliases resolve. | key | what it serves | a shipped use | | --- | --- | --- | | `"output"` | the host's own output, the same location as `.output` | the base `layer_output`, `attention_output`, `mlp_output` | | `"input"`, `"inputs"` | the host's first argument; the whole `(args, kwargs)` pair | the base `layer_input`; the root's `input_ids` | | `".output"`, `".input"` | a child module, aliases resolved | the base `layer_mid`, `"post_attention_layernorm.input"` | | `"../.output"` | a module named from the parent, native names only | Gemma-2's `attention_output` | | `"/.output"`, `"/inputs"` | a path from the model's root, aliases resolved | the vision tower's `image_token_mask`, `"/inputs"` | | `"source..input"`, `"source..output"` | an operation of the host's forward | BLOOM's `attention_output`, `"source.dropout_add_0.input"` | | `"source..inputs"` with `select=` | one argument of a call the forward makes | the base `attention_queries`, `select=1` | | `"source..source..output"` | an operation inside a function the forward calls | the base `attention_probabilities` | | a function of the host | a path decided at read time, inside the trace | Falcon's `by_alibi(...)`, one operation per `alibi` setting | ## An operation inside the forward BLOOM's sublayers take the residual as an argument and add it inside the module (`dropout_add(x, residual, ...)`). The module's output is then a state of the residual stream, and the contribution is the first argument of that call: ```python # nnterp/families/bloom.py class Attention(Attention): @EProperty( "source.dropout_add_0.input", description="What the attention adds to the residual stream: the tensor entering dropout_add", ) def attention_output(self, value) -> Residual: return value ``` ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("hf-internal-testing/tiny-random-BloomForCausalLM") block = model.layers[0] with model.trace("The Eiffel Tower is in"): x = block.layer_input.save() attn = block.self_attn.attention_output.save() # inside the attention's forward: before the module's own output returned = block.self_attn.output.save() mlp = block.mlp.mlp_output.save() out = block.layer_output.save() print(torch.equal(returned[0], x + attn), torch.equal(x + attn + mlp, out)) # True True ``` The segment after `source` is the operation's label in `print(envoy.source)` ([Finding source ops](finding-source-ops.md)). The descriptor walks the path before every read or write, which instruments the operation for the current run. A value inside a forward is read before the module's own `.output`, as the block above does. In the other order the read fails with `OutOfOrderError`. With `inputs`, `select=n` picks a positional argument and `select="name"` a keyword one, and with `output`, `select=i` indexes a returned tuple. A write repacks the element, and assigning one argument replaces that argument alone. The base's `attention_queries` is `EProperty(f"source.{INTERFACE}.inputs", select=1, ...)`, the second argument of the shared attention call, and BLOOM's is `EProperty("source.self__reshape_0.output", select=0, description=Attention.attention_queries.description)`, the first element of the tuple its `_reshape` returns. `Attention.attention_queries` on the class is the descriptor itself, which lets the redefinition reuse the base's description, and `-> Queries` keeps the base's layout. ## A value the family does not have A marker in the class body takes the place of the inherited descriptor. It keeps the name in the tree and in the repr, as `(layer_mid): Unavailable: `, makes `support()` report the reason, and makes any read or write raise `nnterp.Unavailable` with it before the model runs. Cohere's block is parallel and has no stream between its sublayers: ```python # nnterp/families/cohere.py class Layer(Layer): layer_mid = unavailable(PARALLEL) ``` When the checkpoint decides, `unavailable=` takes a function of the envoy returning a reason or `None`. Every value the base `Attention` reads inside the shared attention forward takes `unavailable=interface_reason`, which calls the envoy's `off_interface()`. The default answer is `needs_eager`, the reason given when the model runs `sdpa`. A family with another reason overrides the method once rather than each value. GPT-2's `reorder_and_upcast_attn` routes around the interface: ```python # nnterp/families/gpt2.py class Attention(Attention): def off_interface(self): if self._module.config.reorder_and_upcast_attn: return "this checkpoint sets reorder_and_upcast_attn, which takes GPT-2's own upcast attention path" return super().off_interface() ``` ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") print(model.support()["self_attn.attention_scores"]) # None model.config.reorder_and_upcast_attn = True print(model.support()["self_attn.attention_scores"][0]) # this checkpoint sets reorder_and_upcast_attn, which takes GPT-2's own upcast attention path ``` A predicate must not raise `AttributeError`. The descriptor re-raises one as a `RuntimeError` naming the check, because a descriptor's `AttributeError` would otherwise come back as "no attribute" and hide the predicate's own bug. ## A value that has to be a copy Falcon's parallel block adds the attention output into the MLP's output tensor in place, and a plain `mlp_output` would change under a saved read. The value reads a clone, and because a clone is invisible to the model, a `transform` hands the edited copy back to be swapped in once the block is done with the read: ```python # nnterp/families/falcon.py from nnterp.components import EProperty, Mlp, Residual, first_tensor, rewrap class Mlp(Mlp): @EProperty(key="output", description="What the MLP adds to the residual stream (a copy, since the block adds the attention into the live tensor in place)") def mlp_output(self, value) -> Residual: return first_tensor(value).clone() @mlp_output.postprocess def mlp_output(self, value): return rewrap(self, value) @mlp_output.transform def mlp_output(self, value, raw): return value.clone() ``` The three hooks answer three different events. The decorated function is the getter, and what it returns is what a read serves. `postprocess` runs on an assigned value, and what it returns is what is written to the location. `transform` runs on the served view after an in-place edit, and what it returns is what the model receives when the block moves on. `postprocess` is the write side of a value keyed on `"output"`: The base's `layer_output` and `mlp_output` read `first_tensor(value)` and write `rewrap(self, value)`, both from `nnterp.components`, where `first_tensor` is the tensor of a module output, its first element when the output is a tuple, and `rewrap` puts a tensor back into that tuple, so a module that returns a tuple reads as a tensor and an assignment puts the tensor back into its tuple. A value redefined on a path that serves a bare tensor, such as a sibling norm's output or an operation's, needs neither. nnsight's [eproperty](https://nnsight.net/documentation/intervention/eproperty/) reference has `postprocess` and `transform` in full. A root size the config spells its own way is not a descriptor override but a module-level function named after the size ([Adding a family](adding-a-family.md)), and a value that is new rather than relocated is [Custom values](custom-values.md). # Custom values A value nnterp does not ship is one descriptor on a subclass of the family's `Attention`, `Layer` or `Mlp`, installed for one load with `envoys=`: an `EProperty` keyed on a path from the host, or a `DerivedEProperty` computed from other values. It then reads, writes and saves like a standard value, is listed in the repr with its description, and answers `support()`. ```python import torch from jaxtyping import Float from torch import Tensor from nnterp import StandardizedTransformer, DerivedEProperty, EProperty from nnterp.components import BATCH, HEADS, QUERY, Pattern, interface_reason from nnterp.families import gpt2 from transformers.models.gpt2.modeling_gpt2 import GPT2Attention # after nnterp def attention_entropy(self) -> Float[Tensor, f"{BATCH} {HEADS} {QUERY}"]: probs = self.attention_probabilities return -(probs * torch.log(probs.clamp_min(1e-12))).sum(-1) class Attention(gpt2.Attention): """GPT-2's attention plus two values of my own.""" @EProperty( "source.attention_interface_1.source.nn_functional_softmax_0.output", description="The softmax output before the dropout", unavailable=interface_reason, ) def attention_softmax(self, value) -> Pattern: return value attention_entropy = DerivedEProperty(attention_entropy, description="Per-row entropy of the pattern", unavailable=interface_reason) model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager", envoys={GPT2Attention: Attention}) attn = model.layers[0].self_attn print([line.strip() for line in repr(attn).splitlines() if "attention_softmax" in line or "attention_entropy" in line]) # ['(attention_softmax) -> Pattern [batch num_heads query key]: The softmax output before the dropout', # "(attention_entropy) -> Float[Tensor, 'batch num_heads query'] [batch num_heads query]: Per-row entropy of the pattern"] print(attn.support()["attention_softmax"], model.support()["self_attn.attention_entropy"]) # None None with model.trace("The Eiffel Tower is in"): soft = attn.attention_softmax.save() ent = attn.attention_entropy.save() print(soft.shape, ent.shape) # torch.Size([1, 12, 7, 7]) torch.Size([1, 12, 7]) ``` The softmax is an operation inside transformers' shared eager attention, and the path has a `source` segment for the call and another for the operation inside it ([Finding source ops](finding-source-ops.md)). `unavailable=interface_reason` is the predicate the base's interior values carry: Under the checkpoint's default `sdpa`, both values then report "read inside the eager attention forward, but this model runs 'sdpa'; load with attn_implementation='eager'" in `support()` instead of failing at read time with `SourceNotAvailable`. A `DerivedEProperty` runs its function on the envoy at read time, inside the trace, and forward order applies to what it reads: `attention_entropy` reads the pattern, and a trace that reads it and then `attention_queries` is out of order. A derived value is read-only, and assigning raises `AttributeError: attention_entropy is derived and read-only`. ## On the block or the MLP The same pattern on the family's other classes, keyed on the module's own output and on a child module: ```python from jaxtyping import Float from torch import Tensor from nnterp import StandardizedTransformer, EProperty from nnterp.components import BATCH, HIDDEN, Residual from nnterp.families import gpt2 from transformers.models.gpt2.modeling_gpt2 import GPT2Block, GPT2MLP # after nnterp class Mlp(gpt2.Mlp): @EProperty(key="output", description="The MLP output at the last position") def last_position(self, value) -> Float[Tensor, f"{BATCH} {HIDDEN}"]: return value[:, -1] class Layer(gpt2.Layer): @EProperty("ln_2.input", description="The residual stream after the attention sublayer") def mid_stream(self, value) -> Residual: return value model = StandardizedTransformer("openai-community/gpt2", envoys={GPT2MLP: Mlp, GPT2Block: Layer}) with model.trace("The Eiffel Tower is in"): mid = model.layers[0].mid_stream.save() last = model.layers[0].mlp.last_position.save() print(mid.shape, last.shape) # torch.Size([1, 7, 768]) torch.Size([1, 768]) print([name for name in model.support() if name in ("mid_stream", "mlp.last_position")]) # ['mid_stream', 'mlp.last_position'] ``` `key="output"` shares the module's own output location, and a stub that returns a slice hands the reader a view the model does not see: An in-place edit to `last_position` needs a `transform` to land, and an assignment a `postprocess` that rebuilds the served shape ([Overriding a value](overriding-values.md)). A path below the host resolves aliases, and `"ln_2.input"` and `"post_attention_layernorm.input"` are the same location on GPT-2's block, where the base reads `layer_mid`. Above the host, after `../`, only native names resolve: from GPT-2's attention `"../ln_1.input"` works, `"../input_layernorm.input"` does not. `model.support()` walks each block's children that carry values under their standard names, where a value on an `Attention` or `Mlp` appears as `self_attn.` or `mlp.`, and a load without your `envoys=` lists none of them. ## The layout annotation The return annotation is the value's declared layout, read by `.layout` and `.dims` and shown in the repr. The standard layouts, such as `Residual`, `Pattern` and `Keys`, are exported by `nnterp.components`, and a shape none of them has takes an inline `Float[Tensor, ...]` composed from the same axis constants, such as `BATCH`, `HEADS`, `QUERY` and `HIDDEN`. `attention_softmax` above is annotated `-> Pattern`, and `Attention.attention_softmax.layout is Pattern`, the object the base's `attention_probabilities` carries. `Attention.attention_entropy.dims` is `('batch', 'num_heads', 'query')`. The family suite checks every value's tensor against its annotation, so a value that will ship carries one. Without an annotation `layout` is `None`, and the repr lists `(name): description`. ## Key on the module type `envoys=` entries merge over the family's `ENVOYS`, and nnsight tries type keys before path keys: `envoys={"attn": Attention}` on GPT-2 leaves the family's `gpt2.Attention`, keyed on `GPT2Attention`, in place. Key yours on the type, and subclass the family's class rather than `nnterp.Attention` so the family's own overrides stay. To make the class the family's for every load of the model type in the process, [Registering](registering.md) is the next page. # Finding source ops A `source.` path names an operation inside a module's forward by the label nnsight gives it, and the only correct list of labels is the one printed from the transformers version you run: `print(envoy.source)` outside a trace for the module's own forward, `.source` inside a trace for a function the forward calls. nnsight's [intermediate operations](https://nnsight.net/features/11_source/) page is the full account of `.source`. The listings below are transformers 5.17's. ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") print(model.layers[0].self_attn.source) ``` ``` * def forward( ... self_c_attn_1 -> 40 query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2) split_1 -> + ... ... query_states_view_0 -> 46 query_states = query_states.view(shape_q).transpose(1, 2) transpose_4 -> + ... query_states_1 -> + ... ... 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 -> + ... 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( 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, ... 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 -> + ... 81 return attn_output, attn_weights ``` The labels on the left are the operation names, trimmed here to the lines a path uses. A call is `_`, the attribute chain joined with `_`: `self.c_proj(...)` is `self_c_proj_0`, and `self._upcast_and_reordered_attn(...)` is `self__upcast_and_reordered_attn_0`, the underscore of the method name kept. A binding is an operation too, `_` for the n-th assignment to that name in the forward, and the `+ ...` lines under a call are the bindings of its result. Calls and bindings share one counter per name, so `attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(...)` is `attention_interface_0` and the call `attention_interface(...)` is `attention_interface_1`, which is the string `nnterp.components.INTERFACE` holds. The counter follows the order in the source, not the order the forward runs, and the listing includes branches this forward never takes: GPT-2's cross-attention path owns `self_c_attn_0`, and the live call is `self_c_attn_1`. `.output` is what a call returned or a binding bound, `.input` a call's first argument, `.inputs` its `(args, kwargs)` pair. ## Inside a call `attention_interface_1` calls a function, transformers' `eager_attention_forward`. Its operations exist only under `attention_interface_1.source`, and that drill works inside a trace only, because the callee is a local bound at run time: nnsight resolves it from the live call the first time a trace drills into it, and clears what it built when the next run starts. ```python import torch from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") attn = model.layers[0].self_attn with model.trace("The Eiffel Tower is in"): args = attn.source.attention_interface_1.inputs.save() # the call's arguments: served first inner = attn.source.attention_interface_1.source # the drill into the callee, inside the trace only print(inner) probs = inner.nn_functional_dropout_0.output.save() # then an operation inside it pattern = attn.attention_probabilities.save() print(args[0][1].shape, probs.shape, torch.equal(probs, pattern)) # torch.Size([1, 12, 7, 64]) torch.Size([1, 12, 7, 7]) True ``` ``` * def eager_attention_forward(module, query, key, value, attention_mask, scaling=None, dropout=0.0, **kwargs): ... key_transpose_0 -> 3 attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling torch_matmul_0 -> + ... attn_weights_0 -> + ... 5 if attention_mask is not None: attn_weights_1 -> 6 attn_weights = attn_weights + attention_mask nn_functional_softmax_0 -> 8 attn_weights = nn.functional.softmax(attn_weights, dim=-1) attn_weights_2 -> + ... 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 -> + ... torch_matmul_1 -> 14 attn_output = torch.matmul(attn_weights, value) attn_output_0 -> + ... attn_output_transpose_0 -> 15 attn_output = attn_output.transpose(1, 2) attn_output_1 -> + ... 17 return attn_output, attn_weights ``` This listing is where the base `Attention`'s values come from: The queries, keys and values are the call's arguments 1 to 3, `attention_scores` is `nn_functional_softmax_0`'s input, `attention_probabilities` is `nn_functional_dropout_0`'s output, and `attention_head_outputs` is element 0 of the call's own output. An `EProperty` does the drill at every read, with a `source` segment between the call and the operation inside it: `"source.attention_interface_1.source.nn_functional_dropout_0.output"`. Outside a trace, `attn.source.attention_interface_1.source` raises `SourceNotAvailable: recursive .source is only available inside a trace`. ## The operation to pick The pattern is the dropout after the softmax wherever one exists, because that is the tensor the values are mixed with, cast back to the model's dtype and, on an attention-sink model, with the sink column dropped. In eval the two tensors are equal, and the rule is about what the location means. A contribution added inside the module is the tensor entering the add, BLOOM's `dropout_add_0` input. Queries, keys and values are what the attention arithmetic receives, after the rotary embedding where there is one and before `repeat_kv`, so the head axis is `num_kv_heads` wide. ## Errors from a wrong name ```python from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", attn_implementation="eager") model.layers[0].self_attn.source.nope_0 # AttributeError: 'model.transformer.h.0.attn.source' has no operation 'nope_0'; available: is_cross_attention_0, isinstance_0, ... ``` An `EProperty` on a missing operation raises `SourceNotAvailable` with the value's path, its key and that list: "model.transformer.h.0.attn.broken reads 'source.no_such_op_0.output', which this run does not have: ...". A label that is in the listing but under a branch this forward does not take fails differently: `attn.source.self__upcast_and_reordered_attn_0.output` raises `OutOfOrderError: '...self__upcast_and_reordered_attn_0.output.i0' was requested but the model already ran past it`, because the model ran to the end without reaching it. A call to a submodule is not drilled: `attn.source.self_c_proj_0.source` raises `SourceNotAvailable: 'self_c_proj_0' calls a submodule; call .source on that submodule directly`, which is `attn.c_proj.source`. Requests are served in the forward's order, and for a call and its inside that order is the call's `.inputs`, then the drill, then the operations inside, then the call's `.output`. The block above reads them that way. Reading `attention_interface_1.inputs` after an operation inside the call raises `OutOfOrderError` on `...attention_interface_1.input.i0`. Two values read in one trace follow the same rule, `attention_queries` (the call's inputs) before `attention_probabilities` (inside the call), and the family suite reads interior values one per trace for this reason. The labels are what transformers releases change: A renamed local renames a binding, and a refactor that routes a family onto the shared interface replaces its own softmax with `attention_interface_1`. Print the source for the version you run rather than copying a label from another family or from this page, and let `test_every_source_value_resolves_on_every_layer` and `test_written_pattern_moves_the_logits` in the family suite be the guard: every source-located value must resolve on every attention block, and a written pattern must move the logits, since a label can read a perfectly good tensor that nothing downstream uses. The descriptors that take these labels are on [Overriding a value](overriding-values.md). # Registering `nnterp.families.register(family, *model_types)` makes a family the one every later load of those model types resolves to, for the rest of the process, whether the family is kept outside the package or replaces a shipped one. `StandardizedTransformer(repo_id, family=module)` uses a family for one load and leaves the registry alone. Either takes any module or object with `RENAME` and `ENVOYS`. The signatures of `register` and `lookup` are under [families](../documentation/families/index.md) in the Documentation tab. ```python import types from nnterp import StandardizedTransformer, families from nnterp.families import gpt2 variant = types.SimpleNamespace(**vars(gpt2)) # everything the shipped module defines variant.RENAME = {**gpt2.RENAME, "mlp": ["mlp", "ffn"]} # one more alias for the feed-forward families.register(variant, "gpt2") model = StandardizedTransformer("openai-community/gpt2") print(model.family is variant, model.layers[0].ffn is model.layers[0].mlp, model.intermediate_size) # True True 3072 del families.REGISTRY["gpt2"] # back to the shipped module print(families.lookup("gpt2") is gpt2) # True ``` `vars(gpt2)` copies everything the shipped module defines, and the variant carries GPT-2's `intermediate_size` function along with its classes. A list as a `RENAME` value gives one native module several aliases, here `mlp` and `ffn` for GPT-2's `mlp`. A namespace built from `RENAME` and `ENVOYS` alone also loads and answers `support()` with every block value, because `support()` walks the envoy tree the `ENVOYS` build and reads no class off the family. The root's sizes and `project_on_vocab`, though, look for a function of their name on `model.family` before falling back to their own rule, and with that namespace `model.intermediate_size` falls back to `config.intermediate_size`, which GPT-2's config does not have, and raises `AttributeError`. A variant of a family that defines a size function, such as GPT-2, Falcon, BLOOM or DeepSeek-V2, passes it on, which the spread above does. ## One load `family=` skips the lookup, leaving the config unread for it and `REGISTRY` untouched: ```python import types from nnterp import StandardizedTransformer, families from nnterp.families import gpt2 variant = types.SimpleNamespace(**vars(gpt2)) variant.RENAME = {**gpt2.RENAME, "mlp": ["mlp", "ffn"]} model = StandardizedTransformer("openai-community/gpt2", family=variant) print(model.family is variant, model.layers[0].ffn is model.layers[0].mlp) # True True print(families.lookup("gpt2") is gpt2, families.REGISTRY) # True {} ``` No model type is named or checked against the checkpoint, and a family can be applied to a checkpoint with a `model_type` it was not written for. `rename=` and `envoys=` still apply on top of either route: The aliases are the family's `RENAME` and then `rename=`, the envoys nnsight's tensor-parallel defaults, then the family's `ENVOYS`, then `envoys=`, and a key given later wins. ## A module of your own A family written as a file ([Adding a family](adding-a-family.md)) and named after its model type registers with no types given: `register` takes the type from the last component of the module's `__name__`, as a shipped module's file name gives it. ```python from nnterp import StandardizedTransformer, families import llama # a family module of my own, named after the model_type it covers families.register(llama) # no types given: covers "llama", from the module's name print(families.lookup("llama") is llama) # True model = StandardizedTransformer("HuggingFaceTB/SmolLM2-135M") print(model.family is llama) # True del families.REGISTRY["llama"] ``` A module named otherwise, or one covering several types, names them, `families.register(my_family, "zamba", "zamba2")`. A `SimpleNamespace` built from scratch has no `__name__`, and `register(ns)` raises `TypeError: register() needs the model types the family covers: a SimpleNamespace has no __name__ to take one from`. One spread from `vars(gpt2)` inherits `nnterp.families.gpt2` as its `__name__` and registers under `gpt2`. ## The lookup order `lookup(model_type)` returns, in order, `REGISTRY[model_type]` if `register` put one there, else `nnterp.families.` imported on first use, else `nnterp.families.default`, the best-effort family, with this warning: ``` UserWarning: nnterp has no family for model_type 'nanochat'; the default family standardizes it as a best-effort guess. Check model.support() for what it found, and add nnterp/families/nanochat.py (or nnterp.families.register(family, 'nanochat')) for a standardization you can rely on. ``` Registering a family for the type, or shipping its module, takes it off the default. Register before loading: The lookup runs in `StandardizedTransformer.__init__`, and a model built earlier keeps the family it resolved to. `register` overrides silently, and a test that registers cleans up with `del families.REGISTRY[...]` in a `finally`. `families.known()` lists the shipped modules' names only, not registered families and not `default`. `model.family` is whatever the load resolved to. Another engine has its own registry: `register(family, *model_types, engine="vllm")` registers a family for `StandardizedVLLM`, which `lookup(model_type, engine="vllm")` finds, and a family registered without `engine` is a transformers family only. # Architecture A checkpoint's `model_type` selects one family module, and `StandardizedTransformer.__init__` hands that module's `RENAME` to nnsight as `rename=` and its `ENVOYS` as `envoys=`. nnsight then builds a tree in which every native envoy also answers to its standard alias and every block, attention and MLP is the family's `Layer`, `Attention` or `Mlp` subclass, and that class's `EProperty` descriptors are the values. The root class adds the whole-model values and the sizes. ```python import nnterp from nnterp import StandardizedTransformer from nnterp.families import gpt2 model = StandardizedTransformer("openai-community/gpt2", dispatch=True, attn_implementation="eager", device="cpu") print(model.family is gpt2) # True print(model._aliases) # {'embed_tokens': 'transformer.wte', 'layers': 'transformer.h', 'norm': 'transformer.ln_f'} print(model.layers[0]._aliases) # {'input_layernorm': 'ln_1', 'self_attn': 'attn', 'post_attention_layernorm': 'ln_2'} print(model.layers[0].self_attn is model.transformer.h[0].attn) # True print(type(model.layers[0]) is gpt2.Layer, type(model.layers[0].self_attn) is gpt2.Attention) # True True print(sorted(gpt2.Attention.values())) # ['attention_head_outputs', 'attention_keys', 'attention_output', 'attention_premix', 'attention_probabilities', 'attention_queries', 'attention_scores', 'attention_values'] with model.trace("The Eiffel Tower is in"): pattern = model.layers[2].self_attn.attention_probabilities.save() resid = model.layers[2].layer_output.save() logits = model.logits.save() print(tuple(pattern.shape), tuple(resid.shape), tuple(logits.shape)) # (1, 12, 7, 7) (1, 7, 768) (1, 7, 50257) ``` An alias is the same envoy object, so `model.layers[0].self_attn` is the `gpt2.Attention` that nnsight built for `transformer.h.0.attn`, and the descriptors on that class serve values at that module's native location. The pattern is read inside block 2's attention and the stream as block 2 returns, which is the forward order. ## From `model_type` to a family `StandardizedTransformer.__init__` in `nnterp/standardized.py` reads the config before anything is built. `_read_config` returns a ready `nn.Module`'s own `config`, else `AutoConfig.from_pretrained` with the call's `revision` and `trust_remote_code`, and a config transformers cannot parse fails there with transformers' own error. On a multimodal config the lookup takes `text_config.model_type`, because the text model is what the family standardizes. `families.lookup(model_type)` in `nnterp/families/__init__.py` consults `REGISTRY` first and then imports `nnterp.families.`. A `ModuleNotFoundError` naming exactly that module means there is no such family, and `lookup` warns and returns `nnterp.families.default`. Any other import error comes from inside a family module and is re-raised as it is. A family module imports its transformers modeling module at its top, which keeps `import nnterp` free of modeling code and makes the first lookup of a type load only that family's module. `tests/test_registry.py` checks this in a subprocess. ```python import sys from nnterp import families print(len(families.known()), families.lookup("gpt2") is gpt2) # 98 True print(sorted(m for m in sys.modules if m.startswith("nnterp.families."))) # ['nnterp.families.default', 'nnterp.families.gpt2'] ``` `known()` lists the package's modules by file name with `default` left out, so a module's name is its `model_type`. `default` is imported with nnterp itself, since the root compares `family is default` after the build. `register(family, *model_types)` writes into `REGISTRY`, and `family=` on the constructor skips the lookup for one load ([Registering](../registering.md)). Another engine's families form a package under this one: `lookup("llama", engine="vllm")` imports `nnterp.families.vllm.llama`, and since there is no default there a type without a module raises `UnsupportedFamily`. ## The two maps nnsight receives The constructor merges `rename={**family.RENAME, **rename}` and `envoys={**tp_envoys, **family.ENVOYS, **envoys}`, later keys winning, so a caller's entry replaces the family's on the same key. `_base_envoys` starts from nnsight's tensor-parallel envoys on a sharded load because `TransformersModel` only defaults to them: A map passed by nnterp would otherwise drop them. Both maps are nnsight features ([Module Access](https://nnsight.net/features/6_modules/)). Every `rename` key is resolved from every envoy after that envoy's children exist, and an alias that resolves is bound as another attribute of the same envoy object and recorded in `_aliases`. A key that resolves nowhere is skipped. A multi-component key such as `"transformer.h"` resolves only from the root, which is what puts `layers` on `model` rather than on `model.transformer`, while a single name such as `"attn"` resolves in every block that has one. A dotted alias such as `"attn.c_proj": "attn.o_proj"` binds on the attention. The envoy class is chosen per module at construction: nnsight tries the module type's MRO against the type keys of `envoys`, then the string keys as dotted path suffixes, in the native spelling or any alias spelling of the path. A type key always beats a path key, which is why `envoys={"self_attn": X}` does nothing on GPT-2 while the family's `GPT2Attention` key stands, and why a caller displaces a family envoy by keying on the same type. The default family is the one that keys on names (`"layers.*"`, `"self_attn"`, `"mlp"`), which works because its own `RENAME` has aliased them first. ## Which layer owns what | layer | holds | never knows | |---|---|---| | `nnterp/families/.py` | `RENAME`; the `Layer`, `Attention`, `Mlp` subclasses (and `Moe`, a `RecurrentMixer`, the vision classes where the family has them) that point a value at this family's operation or sibling module; `ENVOYS` keyed on transformers types; `def (model)` for a root size this config spells its own way; `def project_on_vocab(model, hidden)` where the model does something of its own after the head | what a value means or how nnsight serves it | | `nnterp/components/` | what each value means and how it is read and written (`EProperty` paths, `select`, transforms); the availability predicates; the shared operation names such as `INTERFACE` | any one family's module names or classes | | `nnterp/standardized.py` | the root values and methods, the sizes, `support()`, the load path, the remote key | an operation name inside a forward; a family's config keys | Gemma-2's `attention_output` is the post-attention norm's output: The family keys it `"../post_attention_layernorm.output"` on its `Attention`, and the component only knows how to walk such a path. BLOOM's contribution is the first argument of its `dropout_add` call, so the family names `"source.dropout_add_0.input"`. A family on transformers' shared path, such as `nnterp/families/llama.py`, overrides nothing and is three container keys, three empty subclasses and the type keys. ## The root `Standardized`, in `nnterp/standardized.py`, is the engine-independent part: the sizes, `support()`, `skip_layers`, `steer`, `get_topk_closest_tokens` and `_read_config`. `StandardizedTransformer` mixes it over `TransformersModel`, and `StandardizedVLLM` over nnsight's `VLLM` ([vLLM](../../vllm/index.md)). The root values are `EProperty` descriptors like a block's. `logits` reads the model's output and picks `.logits`, `token_embeddings` is keyed `"embed_tokens.output"`, `final_norm_output` and `lm_head_output` read `norm` and `lm_head`, `next_token_probs` is derived and refuses assignment, and `input_ids`, `attention_mask` and `input_size` are keyed `"inputs"` and read the call's keyword arguments. Each size is a `StandardizedProperty`. On read it looks for a function of its own name on `model.family` and calls it with the model, else runs the plain rule over `config.get_text_config()`. Assignment raises `AttributeError`, which leaves the family module as the only place a size is spelled otherwise. `project_on_vocab` is a `StandardizedCapability`, the same lookup for a method. ```python print(model.intermediate_size, model.config.n_inner) # 3072 None ``` GPT-2's config leaves `n_inner` unset, and `gpt2.py` defines `intermediate_size` as four times the hidden size. `support()` collects the root values' reasons, then `blocks_support(model.layers)`: every `Standard` child of a block is listed under its standard name, the set of names is the union over the blocks, and a value missing on some blocks reads as `{layer: reason}`. A `Standard` child of the root, such as a vision tower, adds its own rows under its alias ([Vision](vision.md)). ## The default family `nnterp/families/default.py` serves a `model_type` with no module. Its `RENAME` is the union of the spellings the shipped families use for the containers, the embedding, the blocks, the final norm, the head and a block's modules, and its `ENVOYS` key on the standard names. After the build the root calls `default.check(model)`, which raises `UnsupportedFamily` when a required root module is missing, when no block is a `Layer`, or when a shape-only `scan` shows that the blocks do not pass a `[batch, seq, hidden_size]` stream. It serves `layer_output`, the root values, the sizes and the attention interior wherever the module calls the shared interface. `attention_output`, `mlp_output` and `project_on_vocab` are unavailable, since what a sublayer adds and what follows the head depend on code a name does not show. How a value's path becomes a served location is on the next page, [Descriptors](descriptors.md). # Descriptors Every value is an `EProperty`, nnsight's [`eproperty`](https://nnsight.net/documentation/intervention/eproperty/) with two additions: its key is a path from the host envoy, resolved into a served location on every access, and an `unavailable=` reason is checked before anything runs. `select` picks one element of what is served, a transform carries edits back from a reshaped view, and `DerivedEProperty` computes a read-only value from other served ones. The class is in `nnterp/components/eproperty.py` ([eproperty](../../documentation/components/eproperty.md)). ```python import torch import nnterp from nnterp import StandardizedTransformer, Unavailable from nnterp.components import DerivedEProperty, EProperty, Pattern, Residual, interface_reason, unavailable from nnterp.families import gpt2 from transformers.models.gpt2.modeling_gpt2 import GPT2Attention def entropy(attn): p = attn.attention_probabilities return -(p * (p + 1e-12).log()).sum(-1) class Attention(gpt2.Attention): @EProperty("source.attention_interface_1.inputs", select=4, description="The mask the interface is called with", unavailable=interface_reason) def mask(self, value: torch.Tensor) -> Pattern: return value @EProperty("source.attention_interface_1.source.torch_matmul_0.output", description="q @ k^T before the scaling", unavailable=interface_reason) def raw_scores(self, value: torch.Tensor) -> Pattern: return value @EProperty("../ln_2.output", description="The MLP's pre-norm output, read from the attention") def sibling_norm(self, value: torch.Tensor) -> Residual: return value pattern_entropy = DerivedEProperty(entropy, description="Entropy of each row of the pattern") attention_head_outputs = unavailable("dropped on purpose for this example") model = StandardizedTransformer("openai-community/gpt2", dispatch=True, attn_implementation="eager", device="cpu", envoys={GPT2Attention: Attention}) attn = model.layers[0].self_attn print(Attention.sibling_norm.key, Attention.sibling_norm.inside_forward(), Attention.raw_scores.inside_forward()) # ../ln_2.output False True print(Attention.raw_scores.layout is Pattern, Attention.raw_scores.dims) # True ('batch', 'num_heads', 'query', 'key') print(str(Attention.raw_scores)) # (raw_scores) -> Pattern [batch num_heads query key]: q @ k^T before the scaling print(attn.support()["attention_head_outputs"]) # dropped on purpose for this example with model.trace("The Eiffel Tower is in"): mask = attn.mask.save() raw = attn.raw_scores.save() scores = attn.attention_scores.save() ent = attn.pattern_entropy.save() normed = attn.sibling_norm.save() ln2 = model.layers[0].ln_2.output.save() print(torch.equal(normed, ln2), torch.allclose(raw * attn._module.scaling + mask, scores), tuple(ent.shape)) # True True (1, 12, 7) ``` The reads are in forward order: the interface's arguments, then an operation inside it, then its softmax, then the pattern, then the block's second norm. Passing the subclass through `envoys=` keyed on `GPT2Attention` replaces the family's entry, and every block's attention carries the five new names. [Custom values](../custom-values.md) is the user-facing version of this page. ## The key is a path A key is dotted segments ending in `output`, `input` or `inputs`, walked from the host, or from the model's root after a leading `/`. A function of the host is also a key: It runs at read time, inside the trace, and returns such a path, which is how a forward that branches names the call that fires (`RecurrentMixer.KERNEL`) and how Falcon picks an operation by `config.alibi`. `None` means the attribute's own name, which is what a bare `unavailable("...")` marker gets. | segment | meaning | example | |---|---|---| | `output`, last | the current node's output | `"output"`: `Layer.layer_output`, a view over the same location as `.output` | | `input`, last | the first argument of the current node's call | `"source.dropout_add_0.input"`: BLOOM's contribution | | `inputs`, last | the `(args, kwargs)` pair, or one element of it with `select` | `"source.attention_interface_1.inputs"` with `select=1`: `attention_queries` | | `/`, leading | the model's root, then a relative path from there, aliases included | `"/inputs"`: `Vision.image_token_mask` | | `../`, leading, repeatable | the parent module, by native name | `"../post_attention_layernorm.output"`: Gemma-2's `attention_output` | | a name | a child module, aliases included, or under `source` an operation | `"embed_tokens.output"`: `token_embeddings` | | `source` | the current module's or operation's forward, instrumented for this run | `"source.attention_interface_1.source.nn_functional_dropout_0.output"`: `attention_probabilities` | Operation names are nnsight's labels, `{callable}_{occurrence}`, and `print(envoy.source)` lists them ([Finding source ops](../finding-source-ops.md)). `inside_forward()` answers whether a key has a `source` segment or is a function, which the suite uses to find the values a transformers release can move. ## Availability `unavailable=` is a reason string or a predicate `f(envoy) -> str | None`. `reason(obj)` evaluates it on the instance, which lets a checkpoint's config decide (`needs_eager` reads `_attn_implementation`) and a hybrid answer per block. `_check` runs before every read and write and raises `Unavailable`, a `RuntimeError`, naming the host path, the value and the reason. A predicate that itself raises `AttributeError` is re-raised as a `RuntimeError` saying the availability check failed, because an `AttributeError` escaping a descriptor falls through to nnsight's `Envoy.__getattr__` and comes back as "no attribute", hiding the predicate's own bug. For the same reason `Unavailable` is not an `AttributeError`, which is why `hasattr(envoy, name)` and `getattr(envoy, name, None)` raise it instead of answering `False` or `None`. Decide what a checkpoint has with `support()`, outside the trace ([Availability](../../guide/values/availability.md)). A marker such as `attention_head_outputs = unavailable("...")` is never called on a stub, so nnsight's `eproperty.__call__` never fills its `name` and `key`. `EProperty.__set_name__` fills them from the class body, which is what makes the marker show in `support()` and the repr with its reason. ## Resolving a path at run time `_resolve(obj, key)` turns the path into the location string that nnsight's `Mediator` serves, and it runs before every read and write rather than once, because an operation inside a called function exists only after a drill into that call in the current run. nnsight resolves the callee from the live value and clears what it built at the start of every run ([source](https://nnsight.net/documentation/intervention/source/)). A leading `/` swaps the host for `obj.root` and drops the slash, and a `../` right after it is refused, since nothing is above the root. Leading `../` segments are counted and stripped. Below the host every remaining segment but the last is handed to `obj.get(...)`, and each is an attribute: a child module (an alias works), `source`, or an operation of the `Source` a drill returned. The result is `f"{node.path}.{attribute}"`, with `inputs` folded into `input`, since the pair and its first argument are one served location and picking the argument is `_pick`'s job. An `AttributeError` anywhere on the walk becomes `SourceNotAvailable` naming the value, the path and what nnsight found there. ```python class Probe(Attention): @EProperty("source.attention_interface_1.source.no_such_op_0.output", unavailable=interface_reason) def broken(self, value): return value probe = StandardizedTransformer("openai-community/gpt2", dispatch=True, attn_implementation="eager", device="cpu", envoys={GPT2Attention: Probe}) try: with probe.trace("The Eiffel Tower is in"): probe.layers[0].self_attn.broken.save() except Exception as e: print(type(e).__name__, str(e)[:261]) # SourceNotAvailable model.transformer.h.0.attn.broken reads 'source.attention_interface_1.source.no_such_op_0.output', which this run does not have: 'model.transformer.h.0.attn.source.attention_interface_1.source' has no operation 'no_such_op_0'; available: query_size_0, scaling_0 try: with model.trace("The Eiffel Tower is in"): normed = attn.sibling_norm.save() raw = attn.raw_scores.save() except Exception as e: print(type(e).__name__, e) # OutOfOrderError 'model.transformer.h.0.attn.source.attention_interface_1.fn.i0' was requested but the model already ran past it ``` The second failure is the drill's own ordering. Drilling into a call parks the worker on that call's `.fn`, a location the forward passes once, just before the call runs. A value inside the call read after the module has returned asks for a location the model passed already, and the error names the call's `.fn`, not the value. On the module's own `source` nothing is asked of the model, and an operation already drilled this run is reused, so one trace that reads several interior values in forward order drills once. A function key is evaluated once per access and its last segment passed down, never re-derived: `RecurrentMixer.KERNEL` reads the forward's branch variable as a step's first, pinned read, and a second evaluation would run with the pin relaxed ([Recurrent mixers](recurrent-mixers.md)). ## Above the host A path that starts with `../` is not walked through envoys. `_resolve` drops as many trailing segments from the host's own path as there are `../`, appends the rest of the key, and returns the string, so the parent and the sibling are named natively: `"../ln_2.output"` on `model.transformer.h.0.attn` is `model.transformer.h.0.ln_2.output`, which is why the example above spells `ln_2` and not `post_attention_layernorm`. Nothing is drilled. A `source` segment above the host names an operation of the parent's own forward, served by its string because the family instrumented that forward at build. A second `source` on such a path would need the parent drilled at read time, and `_resolve` refuses it with a `ValueError`. A value that reads something far from its host by standard name anchors its key at `/` instead. Instrumenting at build is `Standard.sourced`, in `nnterp/components/standard.py`. A value inside a forward is served only when that forward was instrumented before the call began. Drilling at read time is enough when the read comes before the module runs, and too late when the worker is already parked inside it, as it is when a block's `mlp_output` is an operation of the block's forward read after the block's `attention_output` (Llama 4). The family sets `sourced = True` on the envoy that owns the forward, and `Standard.__init__` and `_update` then touch `self.source` at build and again when real weights replace meta ones, which reinstalls the plain forward. Nothing infers `sourced` from a path. `inside_forward()` says where a value is, not when its forward has to be instrumented. ## `select` and writes `_pick(attribute, value, select)` takes one element of the served value and `_put` puts one back. `input` is `first_input(args, kwargs)` on read and `replace_first_input` on write, nnsight's own rule for `.input`, so a stub on an `input` path receives a tensor and the call's other arguments stay intact. `inputs` with an int `select` is `args[n]`, with a str `kwargs[name]`, and with none the pair itself (`input_ids` takes `kwargs["input_ids"]` in its stub). `output` with an int indexes the returned tuple (`attention_head_outputs` is `select=0`), and a module that returns a bare tensor takes no `select`. `select` may also be a function of the host, evaluated before the read, for a value that sits at different positions in the calls a forward branches to (`StateSpace`'s two kernels). `__get__` is `_check`, one evaluation of the key, `_resolve`, `Mediator.value`, `_pick`, then the preprocess. `__set__` is `_check`, the postprocess, `_resolve`, and when the key selects or ends in `input` a re-read of the current value with the element put back, then `Mediator.swap` of the whole. Assigning `attn.mask` above swaps the interface's fifth argument and keeps the other four. ## Layouts and the repr `layout` is the return annotation of the stub, resolved with `typing.get_type_hints`, and `-> Pattern` written in a family module gives the identical alias object the base uses, and `Attention.attention_keys.layout is Keys` holds across families. A `State | None` annotation is a union, and the non-`None` member is taken. `dims` splits the alias's `dim_str`, and `__str__` renders the repr line as `(name) -> Layout [axes]: description`, which is what `print(model)` shows ([Layouts](../../guide/values/layouts.md)). ## Transforms A value that serves a copy or a reshaped view carries edits back with a transform, nnsight's write-back of `(view, raw)` to the served value. `__get__` binds a `WriteBack` with the raw served value riding along, nnsight flushes it when the worker next moves on, and a repeated read in between is answered with the view it holds. It fires whether or not the view was edited, so a transform must be an identity on an unedited view. Falcon's `mlp_output` uses one to carry in-place edits on a clone back into the model. ## `DerivedEProperty` `DerivedEProperty(compute)` has no location. `__get__` runs `_check` and then `compute(obj)`, which may read any number of served values in forward order. `__set__` raises `AttributeError` saying the value is derived and read-only, which is the one place an `AttributeError` is right, since it is the assignment that fails. `compute` is stored as `_preprocess` only so that `layout` can read its annotation. `pattern_entropy` above and `RecurrentMixer.states` are instances. ```python try: with model.trace("The Eiffel Tower is in"): attn.pattern_entropy = torch.zeros(1) except AttributeError as e: print(e) # pattern_entropy is derived and read-only ``` The function keys, `pinned` and `per_call` that a forward with two kernels needs are the next page, [Recurrent mixers](recurrent-mixers.md). # Recurrent mixers A recurrent mixer's values are the arguments and results of a kernel call its forward makes, and `RecurrentMixer` in `nnterp/components/recurrent.py` holds everything about reaching that call that does not depend on which mixer it is: which of the two kernels fires, decided once per call by the forward's own test; the requirement that the kernels be transformers' pure-torch ones, so `.source` has a body to read inside; `route_kernels`, which rebinds a family's kernels process-wide; and the state after every token, read by occurrence arithmetic. `LinearAttention` (gated DeltaNet), `SelectiveScan` (Mamba-1) and `StateSpace` (Mamba-2) name their kernels in class constants and declare their values. ```python import torch import nnterp from nnterp import StandardizedTransformer, route_kernels model = StandardizedTransformer("Qwen/Qwen3.5-0.8B", dispatch=True, device="cpu") mix = model.layers[0].linear_attn cls = type(mix) print(cls.__name__, cls.CHUNK_KERNEL, cls.RECURRENT_KERNEL, cls.STATE_OP) # LinearAttention torch_chunk_gated_delta_rule_0 torch_recurrent_gated_delta_rule_0 last_recurrent_state_3 print(cls.attention_queries.key, cls.state.key, cls.states.key) # <_token_state_op> <_states> print(model.support()["linear_attn.states"][0]) # the state after each token is materialized only by the token-by-token kernel; the chunked kernel a prompt runs through carries it between chunks. Call nnterp.route_kernels(model.family, 'torch') before tracing this layer (slower, like attn_implementation='eager') route_kernels(model.family, "torch") # before the first trace that touches a mixer of this family prompt = "Hello world there" n = len(model.tokenizer(prompt).input_ids) with model.trace(prompt): q = mix.attention_queries.save() prompt_states = mix.states.save() out = mix.state_output.save() print(n, tuple(q.shape), tuple(prompt_states.shape), torch.equal(prompt_states[:, -1], out)) # 3 (1, 3, 16, 128) (1, 3, 16, 128, 128) True firsts, seqs, stacks, outs = [], [], [], [] # bound outside: a list first bound in the block does not survive it with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer: for step in tracer.iter[:]: op, first, seq = mix._token_op() # (the state op, the occurrence its first token is, this call's length) seqs.append(seq) firsts.append(first) stacks.append(mix.states.save()) outs.append(mix.state_output.save()) print(seqs, firsts) # [3, 1, 1] [0, 0, 1] print(torch.equal(stacks[0], prompt_states), torch.equal(stacks[2][:, 0], outs[2]), torch.equal(stacks[2][:, 0], outs[1])) # True True False route_kernels(model.family, "default") ``` Qwen3.5 is a hybrid, and block 0 holds a `linear_attn`. `support()` lists the value as a dict because every fourth block has a `self_attn` instead. The open `tracer.iter[:]` ends with nnsight's "never reached" warning, which is how an open slice ends. Without the `first` offset, step 2 would ask for occurrence 0 of the recurrent kernel's state op, which step 1 consumed. ## Which kernel fires The mixer's forward binds a flag before it branches and then calls the token-by-token kernel when it has a cached state and one token, the chunked kernel otherwise. A subclass names the pieces in class constants. | constant | meaning | `LinearAttention` | |---|---|---| | `BRANCH` | the binding made before the branch, `True` when the call continues a cached state | `"use_precomputed_states_0"`, the base's default | | `SEQ_OP` | the forward's masking of its input, once per call, with the call's length on axis 1 | `"apply_mask_to_padding_states_0"`, the base's default | | `CHUNK_KERNEL` | the call a prompt runs through | `"torch_chunk_gated_delta_rule_0"` | | `RECURRENT_KERNEL` | the call a decode step runs through | `"torch_recurrent_gated_delta_rule_0"` | | `STATE_OP` | inside the token-by-token kernel, the binding of the state after each token's update, or `None` | `"last_recurrent_state_3"` | | `STEP_STATE_OP` | set when the decode kernel updates the state once instead of looping: the binding of the new state inside it | `None` | `KERNEL` is one `staticmethod` on the base. It reads the `BRANCH` and `SEQ_OP` outputs in the order the forward binds them, sorted by each operation's source line, and returns `RECURRENT_KERNEL` when the call continues a cached state with one token, `CHUNK_KERNEL` otherwise, which means that several tokens over a cached state read the prompt's kernel. `kernel(attribute)` is the key function the values are declared on and returns `source..`: `attention_queries` is `kernel("inputs")` with `select=0`, `state_output` is `kernel("output")` with `select=1`. `attention_output` is the module's own output, unwrapped from a tuple and rewrapped on write, as on `Attention`. ## Once per call The two bindings are served once per module call, and every value of the call depends on the kernel they pick, so `KERNEL` makes its choice once, inside `per_call(envoy, "kernel", choose)`. `per_call` keeps one record per `(envoy.path, key)` on the worker greenlet running the intervention code and recomputes when the record's call is not the current one. The current call is `mediator.iteration or mediator.occurrence(f"{envoy.path}.output")`. A read pinned by `tracer.iter` to step k is served at the k-th occurrence of its location, and a mixer's kernel fires once per call, so the read is in call k. Relaxed, on step 0, or outside `tracer.iter`, the call is how many times the module's output has been passed, which nnsight counts whether or not anything read it ([interleaver](https://nnsight.net/documentation/intervention/interleaver/)). Every run of every invoke has its own worker, so two invokes reading one mixer never share a record and nothing stays on the envoy. The records are the kernel choice, `_token_op`'s `(op, first, seq)` under `"call"`, and `StateSpace`'s kernel arguments under `"arguments"`. ## Pure-torch kernels and routing transformers decorates each kernel with its hub dispatcher, which binds the module-level name to a closure holding `torch_function`, the pure-torch body, and `implementation`, the optimized kernel when its package is installed. `_dispatch` reads that closure with `inspect.getclosurevars`, and `needs_torch_kernels`, the `unavailable` predicate of every kernel value, refuses when `implementation is not torch_function`, naming the package: A compiled kernel has no Python source to drill into. `route_kernels(family, "torch")` rebinds the family's two kernel names in its modeling module to their pure-torch functions, stashing the originals in the module's `_nnterp_kernels`, from which `"default"` restores them. On a mixer with a `STATE_OP` the prompt's name is bound to the loop kernel too, which on a DeltaNet is the decode kernel: after routing both names run the token-by-token rule, the operation keeps its label `torch_chunk_gated_delta_rule_0`, and `STATE_OP` under it fires once per token. The routing is process-wide and must come before the first trace that touches a mixer of that family, because nnsight builds an instrumented forward over a copy of the module's globals on the first drill. A mixer traced before routing keeps calling the kernel it was compiled with, and a read of `states` on it fails as nnsight's "no attribute `states`" while an untouched mixer of the same model serves it. `route_delta_rule(family, "recurrent" | "chunked")` is the DeltaNet spelling. ```python mamba = StandardizedTransformer("state-spaces/mamba-130m-hf", dispatch=True, device="cpu") scan = mamba.layers[0].linear_attn print(type(scan).__name__, type(scan).STATE_OP, type(scan).STEP_STATE_OP) # SelectiveScan ssm_state_3 ssm_state_0 print(mamba.support()["linear_attn.attention_queries"][0]) # read inside transformers' pure-torch mamba_selective_scan, but this process dispatches it to an optimized kernel (mamba_ssm) with no Python source; uninstall it, or call nnterp.route_kernels(model.family, 'torch'), to read these route_kernels(mamba.family, "torch") with mamba.trace(prompt): states = scan.states.save() out = mamba.layers[0].layer_output.save() print(mamba.support()["linear_attn.states"], tuple(states.shape), tuple(out.shape)) # None (1, 3, 1536, 16) (1, 3, 768) route_kernels(mamba.family, "default") ``` With `mamba_ssm` installed, a Mamba-1 or Mamba-2 model needs the routing before any trace on a CPU, including one that reads only `layer_output`: The compiled kernel refuses a CPU tensor (`Expected u.is_cuda() to be true`). ## The state after every token `state` is an `EProperty` with the key function `_token_state_op`, which returns `source..source..output`, a location with one occurrence per token, so nnsight's own `tracer.iter` walks it: `for t in tracer.iter[:n]: mix.state` is the state after every prompt token, and an assignment there is a write the following tokens continue from ([iterator](https://nnsight.net/documentation/intervention/iterator/)). Under `tracer.iter` the pin at the read is a token index while the bindings `KERNEL` reads and the kernel call's `.fn` fire once per forward. `_token_state_op` decides the kernel and drills into its call inside `with pinned(None)`, which relaxes the pin and restores it on exit, and the pinned read that follows asks for that token's occurrence. `pinned(n)` sets the worker mediator's `iteration` the way `tracer.iter[n]` does, and one read relaxes it: A `with` holds one read. `states`, `state_after(t)` and `set_state_after(t, value)` address the same location by index and need `_token_op`'s three numbers, computed once per call: the state op under the drilled kernel, its sequence length from `_seq()` (the queries' axis 1 on a DeltaNet), and `first`, the op's occurrence count at the moment the kernel's arguments are read. Reading an argument parks the worker at the call's start, the one moment that count is the number of tokens earlier calls put through the op. Occurrences are counted per location since the op was drilled this run, so on step k of a `generate` the recurrent op's count is k minus 1, and the prompt's tokens, under the chunk kernel's label, do not offset it. `_states` then loops `with pinned(first + t): op.output` and stacks on axis 1. `states` is a `DerivedEProperty`, read-only. `needs_recurrent_routing` guards `state` and `states`, and answers in this order: a mixer with no `STATE_OP` says its kernels do not materialize the state per token, an optimized kernel gives `needs_torch_kernels`' reason, and a `CHUNK_KERNEL` name still bound to the chunked function says to route. The predicate reads the live binding, so `support()` follows the routing, as the first block shows. ## Mamba-1 and Mamba-2 `SelectiveScan`, in `nnterp/components/selective_scan.py`, differs from the DeltaNet at the points the base leaves open to a subclass. Its forward binds the branch flag before the mask, which `KERNEL`'s line ordering handles. Its pure-torch prompt kernel is itself the token loop (`STATE_OP = "ssm_state_3"`) while the decode kernel updates once (`STEP_STATE_OP = "ssm_state_0"`): `route_kernels` binds each name to its own function, and `_state_op(kernel)` picks the binding for the kernel that fires. Its arguments sit at different positions in the two kernels, and `select` is a function looking the name up in `ARGUMENTS` by whether the call decodes. Its tensors are channel-first, and each preprocess returns a tokens-first view that the postprocess lays back. `betas` and `decays` are `DerivedEProperty` values computed from `dt`, `dt_bias` and `A` the way the kernel computes them. `StateSpace`, in `nnterp/components/state_space.py`, has `STATE_OP = None`, since neither Mamba-2 kernel binds the state once per token. `state` and `set_state_after` are `unavailable(...)` with reasons, and `states` is a `DerivedEProperty` over the chunk scan's `new_states_0` binding, the state at every chunk boundary. `chunk_per_token(model)` sets `chunk_size` to 1 on each `StateSpace` module of one model, which makes every token a boundary. It is per model, not process-wide, and `needs_per_token_chunks` reports the chunk size until then. The two kernels take their arguments in different places (`CHUNK_ARGUMENTS`, `RECURRENT_ARGUMENTS`, through `argument(name)`). The update has no sequence axis, and a decode step's values gain one on read and lose it on write. `betas` and `decays` are the kernel's `dt` seen through `dt_bias`, the softplus and `A`, with an inverse postprocess, and an assignment reaches the kernel. Values that need another argument read the call's `(args, kwargs)` once per call through `_arguments`. ```python mamba2 = StandardizedTransformer("AntonV/mamba2-130m-hf", dispatch=True, device="cpu") ssd = mamba2.layers[0].linear_attn route_kernels(mamba2.family, "torch") print(mamba2.support()["linear_attn.states"][0]) # the chunk scan materializes the state only at chunk boundaries, every 256 tokens (chunk_size=256); call nnterp.chunk_per_token(model) to set every mixer's chunk_size to 1, so every token is a boundary (slower on long prompts) nnterp.chunk_per_token(mamba2) with mamba2.trace(prompt): states = ssd.states.save() out = ssd.state_output.save() print(mamba2.support()["linear_attn.states"], tuple(states.shape), torch.equal(states[:, -1], out)) # None (1, 3, 24, 128, 64) True nnterp.chunk_per_token(mamba2, enabled=False) route_kernels(mamba2.family, "default") ``` ## Adding a mixer A new mixer subclasses `RecurrentMixer`, sets `CHUNK_KERNEL` and `RECURRENT_KERNEL` (and `BRANCH`, `SEQ_OP` or `KERNEL` itself, as a `staticmethod`, when its forward spells the test otherwise), sets `STATE_OP` only when a kernel binds the state once per token, and declares its values at `kernel("inputs")` and `kernel("output")` with `unavailable=needs_torch_kernels`. A kernel with tensors not laid out `[batch, seq, ...]` overrides `_seq`. `tests/test_base.py` holds the base's own tests, and a family's file mixes in `SelectiveScanSuite` or the tests in `tests/families/ssd.py` ([Testing](../testing.md)). # Vision On a checkpoint loaded with `task="image-text-to-text"`, the family's `RENAME` names the vision tower `vision` and the last module before the scatter `projector`, and its `ENVOYS` wrap the tower in `Vision` and the tower's blocks in `VisionLayer`, `VisionAttention` and `VisionMlp`, the text classes with their stream values re-annotated as `Patches`. The two values where the image meets the text model, `image_token_mask` and `image_features`, are on `Vision` and read outside the tower: the mask off the model's inputs, the features at the scatter in the forward of the wrapper model the family keys `ImageScatter` on. The classes are in `nnterp/components/vision.py` ([vision](../../documentation/components/vision.md)), and the user-facing page is [Vision-language](../../guide/architectures/vision-language.md). ```python import numpy as np import torch from PIL import Image import nnterp from nnterp import StandardizedTransformer from nnterp.components.vision import scatter_call, scatter_host model = StandardizedTransformer("llava-hf/llava-interleave-qwen-0.5b-hf", task="image-text-to-text", dispatch=True, attn_implementation="eager", device="cpu") print(model.family.__name__, type(model._module).__name__) # nnterp.families.qwen2 LlavaForConditionalGeneration print(model._aliases) # {'embed_tokens': 'model.language_model.embed_tokens', 'layers': 'model.language_model.layers', 'norm': 'model.language_model.norm', 'vision': 'model.vision_tower', 'projector': 'model.multi_modal_projector'} print(model.vision._aliases, model.vision.layers[0]._aliases) # {'patch_embed': 'embeddings.patch_embedding', 'layers': 'encoder.layers', 'norm': 'post_layernorm'} {'input_layernorm': 'layer_norm1', 'post_attention_layernorm': 'layer_norm2'} print(model.vision.num_layers, model.vision.vision_hidden_size, model.vision.patch_size, model.num_layers, model.hidden_size) # 26 1152 14 24 1024 print(scatter_host(model)[0], type(scatter_host(model)[1]).__name__, scatter_call(model)) # model ImageScatter ('model.source.inputs_embeds_masked_scatter_0', 1) print(type(model.vision).image_features.path(model.vision), type(model.vision).image_token_mask.key) # /model.source.inputs_embeds_masked_scatter_0.inputs /inputs image = Image.fromarray((np.random.RandomState(0).rand(64, 64, 3) * 255).astype("uint8")) messages = [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "What is in this image?"}]}] prompt = model.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) with model.trace(prompt, images=[image]): mask = model.vision.image_token_mask.save() tower = model.vision.tower_output.save() proj = model.projector.output.save() features = model.vision.image_features.save() first = model.layers[0].input.save() print(tuple(mask.shape), int(mask.sum()), tuple(tower.shape), tuple(proj.shape), tuple(features.shape)) # (1, 744) 729 (1, 729, 1152) (1, 729, 1024) (729, 1024) print(torch.equal(first[mask], features)) # True ``` The family is `qwen2`, chosen from `text_config.model_type`, and the wrapper's spellings of the text stack and the SigLIP tower's paths are in that module's `RENAME` beside the plain Qwen2 ones. The reads follow the forward: the mask comes off the inputs, the tower and projector run, the scatter writes the features into the embeddings, and block 0 receives them. ## How the tower's keys bind nnsight resolves every rename key from every envoy and binds the alias on the envoy it resolves from ([Architecture](architecture.md)). The tower root and the projector are keyed from the model root (`"model.vision_tower": "vision"`), and the tower's inner names are keyed relative to the tower: multi-component keys no text model has (`"encoder.layers": "layers"`, `"embeddings.patch_embedding": "patch_embed"`) and single names no text block has (`"post_layernorm": "norm"`, `"layer_norm1": "input_layernorm"`). A bare name a text block also has is never a tower key. Where a family hosts the same tower under two wrappers at two paths (`llama`: `model.vision_tower` and `model.vision_model`) both are keyed, and where two towers spell their insides differently (`mistral`: CLIP and Pixtral) both sets are keyed and each binds only on its own tower. A rename key cannot tell CLIP's `post_layernorm`, which norms the pooled CLS token, from SigLIP's, which norms the patches, so `llama` keys neither and its `SiglipVision` serves `norm` as a property. Because `projector` is the last module before the scatter, `projector.input` is what the host feeds its projector, not always `tower_output` (Llava 1.5 feeds block -2's stream, llava-interleave the last block's stream before `post_layernorm`), and `projector.output` is not always what the text model receives (LLaVA-NeXT unpads it and appends newline rows). That is why `image_features` is read at the scatter. ## The blocks and the tower's values `Patches` is `[images, patches, vision_hidden_size]`, the tower's own batch by its tokens. `VisionLayer`, `VisionAttention` and `VisionMlp` redefine the stream values of `Layer`, `Attention` and `Mlp` with that annotation and inherit the attention interior unchanged, so `attention_probabilities` is a `Pattern` over the patches with no causal mask. A tower that runs on `[patches, vision_hidden_size]` (the Qwen ViT) is served through the same classes: `as_patches` gives the 2-D native tensor a leading images axis of 1 as a view, and `as_native` drops it on a write, which leaves no need for a packed variant of the classes. Every block value carries `no_tower_run` as its `unavailable=` predicate, which walks up the parent links to the `Vision` and asks `no_images()`. `Vision` serves `patch_embeddings` (the patch embedding's output, a convolution's grid flattened to rows as a view) and `tower_output`, the last block's stream after the tower's final norm where it has one, before any pooling, CLS dropping or adapter. By default `tower_output` is the output's `last_hidden_state`, and a tower that returns a tensor from after one of those steps reads it elsewhere, as `llama4_text`'s subclass does. Its sizes are read off the tower's own config and shadow the root's on purpose, since a tower value's axes are the tower's sizes. `vision_hidden_size` is named apart from `hidden_size` because `image_features` is the text model's width served on the tower, and `axis_size` walks past the tower to the root for it. A tower with no fixed resolution sets `image_size = property(variable_resolution)`, which raises `Unavailable` naming where each image's grid is. ## Where `image_features` is read The key is a function of the host: `image_scatter` returns `f"/{path}.inputs"` and `scattered_argument` the `select`, both from `scatter_call(vision.root)`. `scatter_host` finds the root's child the family keyed `ImageScatter` on (`LlavaModel`, `Qwen2VLModel`, `Gemma3Model`), or the root itself where the family sets `ROOT_SCATTER` because the wrapper scatters in its own top-level forward (Llama 4, NemotronH Omni). `ImageScatter.scatter` names the operation, `inputs_embeds_masked_scatter_0` on every wrapper with a `masked_scatter`, and `scatter_argument` the argument, 1. Idefics 3 and SmolVLM write the features in through a helper call, and `llama`'s `InputsMerger` subclass names `self_inputs_merger_0` and its `image_hidden_states` keyword. The preprocess flattens the argument to `[image_tokens, hidden_size]` as a view, so in-place edits land, and the postprocess reshapes an assignment back to the argument's shape. `ImageScatter` sets `sourced = True`: Its forward is instrumented at build, because the tower runs inside that forward and a drill on first read, after a tower value, would ask for a location the model passed already ([Descriptors](descriptors.md)). For the root case `StandardizedTransformer._source_root_scatter` touches the root's own `source` at build and again after dispatch, when the family names `ROOT_SCATTER` and the load has a `projector`. ## Availability `Vision.no_images()` is the one place that decides whether an image reaches the model: the family names no `projector` on this wrapper, or the load has no processor. Then every tower value is `Unavailable` with that reason, `model.vision.support()` is empty, and `model.support()` has no `vision` rows. A text-only checkpoint has no `model.vision` at all. Otherwise the root lists the tower's rows the way it lists a block's, under the alias (`"vision.image_features"`, `"vision.self_attn.attention_probabilities"`), through the same `standard_children` walk. ```python text_only = StandardizedTransformer("llava-hf/llava-interleave-qwen-0.5b-hf", dispatch=True, device="cpu") print(type(text_only._module).__name__, text_only.vision.no_images()) # LlavaForConditionalGeneration a text-only load: no processor, so no image reaches the model; load with task='image-text-to-text' print([name for name in text_only.support() if name.startswith("vision")]) # [] ``` ## Loading `task` defaults to `"text-generation"`, and `task="image-text-to-text"` is the only switch: It loads the wrapper with its processor, and `model.tokenizer` is then the processor's tokenizer, which keeps the text helpers working. Under the default task transformers builds the causal-LM class the config maps to, the text-only class where there is one (Llama 4's `Llama4ForCausalLM`, with the weights read out of the wrapper checkpoint). On a config with no causal-LM class, as above, the lazy build is refused and `dispatch=True` falls back through the pipeline to the wrapper with a warning and no processor. A trace on a wrapper takes `model.trace(prompt, images=[image])` with the placeholder the processor's chat template puts in the prompt, or an encoding built with `model.processor`. One image-carrying invoke goes in a trace, since nnsight does not batch `pixel_values` across invokes. ## Which file holds what `nnterp/components/vision.py` holds everything shared: the three layouts, `Vision`, the three block classes, `ImageScatter`, `scatter_host` and `scatter_call`, `no_tower_run`, `variable_resolution`, and the subclasses several families need (`QwenVision`, `QwenVisionAttention`, `PixtralVision`). A family file holds the paths. A tower on the shared attention interface (SigLIP, CLIP, Llama 4's ViT) is about five `RENAME` lines and four `ENVOYS` entries keying the base classes on the tower's module types, plus one `ENVOYS` entry per wrapper keying `ImageScatter` on the wrapper's model, repeated in each family that hosts the tower. A tower that differs gets a small subclass named as the base in its family file, such as `gemma4_text`'s `VisionAttention`, which points the contributions at a sandwich block's post-norms. ## Towers the shared classes do not fit The Qwen ViT splits the queries, keys and values per image and calls the interface once per image, under eager and sdpa alike, so no one tensor is the block's pattern. `QwenVisionAttention` serves the queries, keys and values whole at the `unsqueeze` operations that end their preparation and the head outputs at the `torch.cat` that joins the calls, and marks the scores and pattern `Unavailable` with `PER_IMAGE`, which says to split at `cu_seqlens` and never to load eager. Pixtral's convolution runs on the batch padded to its largest image and the grids are cropped and concatenated afterwards, and `PixtralVision.patch_embeddings` is the packed row entering `ln_pre`. Qwen2.5-VL serves its block values in the tower's window order and does not reorder them. Qwen3-VL's three tower taps re-enter the text stream after text blocks 0 to 2, and `qwen3_vl_text` adds `layers[k].deepstack_output`, a `DeepstackEProperty` pinned to block k's occurrence of the one `_deepstack_process` call, with the text model's envoy `sourced`. Gemma 4 unified has no tower: `vision` is its embedder, a `Vision` with no `layers`. Mllama's cross-attention blocks and Gemma 4's audio tower have no family and no values yet ([Contributing](../contributing.md)). # Testing `HF_HUB_OFFLINE=1 python -m pytest` from the repository root runs every family's file under `tests/families/` against its pinned tiny checkpoint, and the root tests under `tests/`. `FamilySuite` in `tests/families/suite.py` is the contract: 53 methods every family inherits, each an end-to-end statement about a standardized model. A family's file names its checkpoint, its native paths and its quirks as class attributes and adds only what is specific to it. ```bash CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 python -m pytest tests/families/test_gpt2.py -q -p no:cacheprovider # 42 passed, 12 skipped, 2 warnings in 5.72s CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 python -m pytest tests/test_registry.py -q -p no:cacheprovider # 23 passed, 7 warnings in 14.24s ``` The checkpoints are cached, and `HF_HUB_OFFLINE=1` keeps a run from touching the Hub. The fixtures load with `dispatch=True` and no `device=`, which puts every model on a visible GPU, and `CUDA_VISIBLE_DEVICES=""` keeps a run on the CPU. The twelve skips on GPT-2 are the mixture-of-experts methods, which skip on a dense model. `pyproject.toml` sets `testpaths = ["tests"]` and `addopts = "--dist loadgroup"`, so a `pytest -n` run with xdist keeps a class marked `xdist_group` on one worker. `tests/conftest.py` imports nnsight before anything else, since a transformers modeling module imported first can crash the interpreter on import, and pins one torch thread per xdist worker. pytest imports `tests/families/suite.py` as the bare module `suite`, and no other test helper may take that name. ## What `FamilySuite` asserts Two class-scoped fixtures load the checkpoint once per class: `model`, a `StandardizedTransformer` with `dispatch=True, attn_implementation="eager"` plus `LOAD_KWARGS`, and `raw_model`, the same checkpoint as a plain `TransformersModel`. The methods, by what they pin: | group | methods | assert | |---|---|---| | names | `test_family_resolved`, `test_standard_names_alias_native_envoys`, `test_no_inner_model_alias`, `test_every_layer_is_renamed`, `test_envoy_classes`, `test_trace_through_standard_names` | `model.family is FAMILY`; every `NATIVE` row resolves to the native envoy itself; nothing is bound at `model.model`; every block carries the block rows; every block, attention, MLP and mixer is exactly the family's class and subclasses the component base; a trace through the standard names gives `hidden_size` and `vocab_size` widths | | availability | `test_support_lists_every_standard_value`, `test_support_is_what_this_family_expects`, `test_support_matches_what_reads` | `support()` has exactly the expected keys; each `EXPECTED_UNAVAILABLE` key carries its reason substring on every block and every other key is `None`; on block 0 an available name is an `EProperty` and an unavailable one raises `Unavailable` with that reason | | boundaries | `test_layer_output_is_the_tensor`, `test_layer_input_is_the_block_input`, `test_every_value_reads_a_tensor_on_every_layer`, `test_contribution_identity`, `test_sublayer_inputs_are_the_normed_stream`, `test_module_boundaries`, `test_sublayer_interiors`, `test_renamed_model_equals_raw_model`, `test_boundary_writes_land` | `layer_output` equals `.output` or its first element; `layer_input + attention_output + mlp_output == layer_output` on every block at 8 ulp of the block's dtype; what enters each sublayer is `ATTENTION_NORM`'s and `MLP_NORM`'s output; `layer_mid`, the norm outputs, `attention_premix`, `mlp_activation` and `mlp_neurons` are the module boundaries they name and a write to each moves the logits; the renamed model equals the raw one; assigning and zeroing the three boundary values moves the logits | | the pattern | `test_probabilities_are_a_pattern`, `test_attention_interior_under_generate`, `test_pattern_across_layers_and_traces`, `test_written_pattern_moves_the_logits`, `test_every_source_value_resolves_on_every_layer` | `[batch, num_heads, seq, seq]`, the model dtype, rows summing to one (or below one with `ATTENTION_SINK`), lower-triangular; a decode step's pattern has one query row; two traces agree; a written pattern and a zeroed head move the logits; every value with `inside_forward()` reads on every attention block | | the interior | `test_interior_shapes`, `test_interior_writes_are_causal`, `test_interior_in_place_edits` | the queries, keys, values, scores and head outputs have their layouts and `softmax(scores)` is the pattern (`pattern_from_scores` is overridable); zeroing each moves the logits and zeroed head outputs leave a position-independent contribution; in-place edits land, or queries raise `RuntimeError` mentioning `view` with `REFUSES_IN_PLACE_QKV` | | methods | `test_skip_layers_hands_the_stream_straight_through`, `test_skip_layers_with_a_given_stream`, `test_steer`, `test_project_on_vocab_is_the_logit_lens` | skipping blocks 1 to the last leaves block 0's stream as the final one; `skip_with` zeros reach the next block; `steer` moves the last position by `2 * vector` and no other; `project_on_vocab` on the last block's output is `logits` | | layouts | `test_values_match_their_annotations` | every value with a `layout` on the root, a block, its attention, its MLP, its mixture and its mixer is an instance of the alias and each named axis matches `axis_sizes` | | the input | `test_input_accessors`, `test_attention_mask_under_generate`, `test_assigning_input_ids_runs_other_ids` | `input_ids`, `input_size` and `attention_mask` are `[1, n]`; a decode step's mask covers the cache; assigning another prompt's ids reproduces its logits and assigning `input_size` raises | | the root | `test_logits_are_the_models_output`, `test_assigning_logits_replaces_the_result`, `test_token_embeddings_are_the_embedding_output`, `test_next_token_probs`, `test_sizes_match_the_model`, `test_per_module_sizes_match_each_block`, `test_repr_lists_the_values` | `logits` is `output.logits`; assigning it zeros `tracer.result.logits`; `next_token_probs` is the last position's softmax and refuses assignment; the root sizes match the pattern, the stream and the projections, and each block's own sizes match its tensors; the block repr lists the values | | a mixture | the twelve `test_moe_*`, `test_expert_*`, `test_routed_*`, `test_writes_to_the_shared_*`, `test_zeroing_*`, `test_ablating_*`, `test_rerouting_*`, `test_writing_router_logits_*` methods | on the first block with an `mlp` that is a `Moe` running a mixture: the routing pair's shapes, `expert_outputs` summing to `routed_output`, `routed_output + shared_expert_output` being the module's output, a zeroed slot and an ablated expert changing exactly the tokens they touch, a rerouted slot matching a hand computation, written router logits deciding the routing, and per-invoke rows; skipped on a dense model | Reads in one trace follow the forward, and families bind the interior values at different points, so the interior and layout tests read one value per trace. Containers filled inside a trace block are bound outside it, since a name first bound in the block does not survive it. ## How a family file adds specifics Most files are the class attributes alone. `tests/families/test_gpt2.py` sets `REPO`, `FAMILY`, `NATIVE = rows("transformer", "h", "wte", "ln_f", attn="attn", ln1="ln_1", ln2="ln_2")` and `REFUSES_IN_PLACE_QKV = True`, and adds one test that flipping `reorder_and_upcast_attn` on the config makes the interface unavailable. `rows(container, layers, embed, norm, ...)` builds the standard-to-native map, and `LLAMA_ROWS` is the Llama layout. | attribute | states | |---|---| | `EXPECTED_UNAVAILABLE` | `support()` key to a substring of the reason, for the values this checkpoint lacks; a module no block has is absent from `support()` and is not listed here | | `ATTENTION_NORM`, `MLP_NORM`, `MLP_NORM_BEFORE_ATTENTION` | which block module's output enters each sublayer (`None` when the block input does), and whether the MLP's norm fires first | | `REFUSES_IN_PLACE_QKV`, `ATTENTION_SINK`, `KV_HEADS_EXPANDED`, `QUERY_GATED` | queries that are split views, a sink column, keys and values already expanded to `num_heads`, a `q_proj` twice as wide | | `MLP_WIDTH_KEY`, `LOAD_KWARGS` | a config key naming the MLP width when it is not `intermediate_size`; extra load arguments (DBRX needs `dtype=torch.float32`) | | `MOE_UNAVAILABLE`, `ROUTER_EXTRA_CLASSES` | mixture values the mixture lacks; router classes beyond the experts | | `DECODE_MASK_COVERS_CACHE`, `MID_IDENTITY`, `PREMIX_IS_HEAD_OUTPUTS`, `NEURONS_FROM_ACTIVATION` | a pure state-space model's decode mask, a gated stream, a norm between the heads and the output projection, a fused `gate_up_proj` | A family with a checkpoint that needs another config or another kernel path makes a second class rather than mutating the shared fixture. `test_falcon.py` has three classes for the 7B layout, the 7B layout with alibi written into a copy of `config.json`, and the 40B layout, and `test_olmo3.py` symlinks the snapshot into a temporary directory and rewrites only `config.json`. The hybrid files override `test_every_layer_is_renamed` for the per-block `self_attn` and `linear_attn` split. The Mamba-1 files mix `SelectiveScanSuite` (`tests/families/scan_suite.py`) in before `FamilySuite`: A class-scoped autouse fixture routes the family to the pure-torch kernels for the class and restores the default afterwards, which matters because the routing is process-wide. The Mamba-2 files take their mixer tests from `tests/families/ssd.py`. A multimodal wrapper runs the whole `FamilySuite` loaded with `LOAD_KWARGS = {"task": "image-text-to-text"}`, and its vision side runs `VisionSuite` from `tests/families/vision_suite.py`: the tower names, the tower's identity, `layers[0].input[vision.image_token_mask] == vision.image_features`, causal edits, and no `vision` host on a text-only checkpoint or load. `tests/families/test_default.py` runs the default family on several shipped families' checkpoints with the contributions listed unavailable, and checks that every value it serves is bit-identical to the dedicated family's. ## The root tests ```bash CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 python -m pytest tests/test_base.py tests/test_prompt_utils.py tests/test_nnsight_utils.py tests/test_order.py -q -p no:cacheprovider # 51 passed in 26.67s ``` | file | covers | |---|---| | `tests/test_registry.py` | every module under `nnterp/families/` is named after its type; `import nnterp` pulls in no `transformers.models.*.modeling_*` module and `lookup("gpt2")` imports only that family, in a subprocess; an unknown `model_type` falls back to the default family with a warning; `register` adds and overrides; a preloaded module uses its own config; a caller's `rename` and `envoys` merge over the family's; the remote key is `TransformersModel`; a size or `project_on_vocab` defined on a family wins over the root's; `family=` takes a module or a `SimpleNamespace` for one load; a size refuses assignment | | `tests/test_base.py` | a tuple-returning block is unwrapped and rewrapped; an `unavailable(...)` marker is listed, raises on read and makes `hasattr` raise; paths with `../`, `/` and `select` functions resolve and write; `route_kernels` round-trips a module's bindings; `per_call` separates invokes and replays an edit; `axis_size` resolves by name up the tree; `probs_to_dict` keeps tokens that decode alike | | `tests/test_prompt_utils.py`, `tests/test_nnsight_utils.py` | the helpers, on GPT-2 with left padding | | `tests/test_order.py` | the measured forward order and its ranks | | `tests/test_encyclopedia.py` | every encyclopedia entry builds from its pinned checkpoint | | `tests/vllm_families/` | the vLLM engine against the transformers one; collected only where vLLM imports and a GPU is visible | The procedure for a transformers release that moves an operation name is on [transformers compatibility](transformers-compatibility.md). # transformers compatibility `pyproject.toml` requires `transformers>=5.17`, `nnsight>=0.8` and `jaxtyping`, and the pages on this site ran on transformers 5.17.0 and nnsight 0.8.0. The module names and classes a family binds to (`RENAME`, `ENVOYS`) rarely move between releases. The operation names inside a forward that every `source.` path pins are what a release renames. ```python import nnterp from nnterp import StandardizedTransformer model = StandardizedTransformer("openai-community/gpt2", dispatch=True, attn_implementation="eager", device="cpu") attn = model.layers[0].self_attn print([op.name for op in attn.source][-12:]) # outside a trace: the module's own forward # ['using_eager_0', 'ALL_ATTENTION_FUNCTIONS_get_interface_0', 'attention_interface_0', 'self__upcast_and_reordered_attn_0', 'attention_interface_1', 'attn_output_reshape_0', 'contiguous_0', 'attn_output_0', 'self_c_proj_0', 'attn_output_1', 'self_resid_dropout_0', 'attn_output_2'] with model.trace("Hello world there"): inner = [op.name for op in attn.source.attention_interface_1.source].save() # inside: drilled from the live callee print(inner) # ['query_size_0', 'scaling_0', 'key_transpose_0', 'torch_matmul_0', 'attn_weights_0', 'attn_weights_1', 'nn_functional_softmax_0', 'attn_weights_2', 'attn_weights_type_0', 'attn_weights_3', 'nn_functional_dropout_0', 'attn_weights_4', 'torch_matmul_1', 'attn_output_0', 'attn_output_transpose_0', 'attn_output_1'] print({name: attr.key for name, attr in type(attn).values().items() if attr.inside_forward()}) # {'attention_queries': 'source.attention_interface_1.inputs', 'attention_keys': 'source.attention_interface_1.inputs', 'attention_values': 'source.attention_interface_1.inputs', 'attention_scores': 'source.attention_interface_1.source.nn_functional_softmax_0.input', 'attention_probabilities': 'source.attention_interface_1.source.nn_functional_dropout_0.output', 'attention_head_outputs': 'source.attention_interface_1.output'} ``` `print(attn.source)` renders the forward with every operation labeled at its line, and that listing is what to diff across releases ([Finding source ops](finding-source-ops.md)). ## What a release can move An operation name is `{callable}_{occurrence}`: the dotted name in the source joined with `_`, and a count over the calls and bindings of that name in the order they appear. Three things move a name: the call is renamed, a binding of the same name is added or removed before it, or the arithmetic moves onto or off transformers' shared attention interface. The occurrence counts bindings as well as calls, which is why `attention_interface_1` is the call: `attention_interface_0` is the binding `attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(...)` two lines above it. nnsight numbers every call in the source, executed or not, so GPT-2's `self__upcast_and_reordered_attn_0`, a branch that does not run, still takes a label. ### The shared interface `INTERFACE = "attention_interface_1"` in `nnterp/components/attention.py` is the call every family on transformers' shared attention path makes, and inside `eager_attention_forward` the pattern is `nn.functional.softmax(...)` then `nn.functional.dropout(...)`. The base `Attention` reads the interface's arguments 1, 2 and 3 as the queries, keys and values, `nn_functional_softmax_0.input` as the scores, `nn_functional_dropout_0.output` as the pattern, and `output[0]` as the head outputs. GPT-OSS keeps the interface but concatenates a sink column before its softmax, so `nnterp/families/gpt_oss.py` reads the scores at `attn_weights_1`, the masked scores bound just before the column joins. That binding exists only when a mask is passed, as one is on every prompt. ### Families with their own attention arithmetic | family | value to operation, in `nnterp/families/.py` | |---|---| | BLOOM | queries, keys, values `self__reshape_0.output[0..2]`; scores `F_softmax_0.input`; head outputs `torch_bmm_0.output`; pattern `self_attention_dropout_0.output`; both contributions at `dropout_add_0.input` | | MPT | queries, keys, values the bindings `query_states_0`, `key_states_0`, `value_states_0`; scores `nn_functional_softmax_0.input`; head outputs `torch_matmul_1.output` | | Falcon | each key chosen by `config.alibi` through `by_alibi(without, with_alibi)`: queries and keys `apply_rotary_pos_emb_0.output` or the bindings `query_layer_0`, `key_layer_0`; values `value_layer_0`; scores `F_softmax_0.input` or `F_softmax_1.input`; head outputs `attn_output_1` or `flatten_0`; pattern `F_softmax_0` or `self_attention_dropout_0` | | GPT-J, GPT-Neo, GPT-NeoX-Japanese | around the `self._attn(...)` call: `self__attn_0.inputs[0..2]`, `self__attn_0.source.nn_functional_softmax_0.input`, `self__attn_0.output[0]`, the dropout inside `self__attn_0.source` | | CodeGen | GPT-J's operations, except the scores at `self__attn_0.source.call_0.input`, the call of the `nn.Softmax` instance built on the same line | | XGLM | queries, keys, values at `*_states_reshape_0.output`; scores at `nn_functional_softmax_0` or `nn_functional_softmax_1` by the weights' dtype, a function key; head outputs `torch_bmm_1.output` | ### Recurrent mixers `torch_chunk_gated_delta_rule_0` and `torch_recurrent_gated_delta_rule_0` are the two DeltaNet kernel calls, `use_precomputed_states_0` the branch binding that picks between them, `apply_mask_to_padding_states_0` the masking that gives a call's length, and `last_recurrent_state_3` the per-token state binding inside the recurrent loop, where the `_3` counts the `last_recurrent_state = ...` bindings before the one after a token's update. A release that adds or removes one binding in that loop moves it. The kernels' argument names `g`, `beta` and `initial_state` are the `select` keys of `decays`, `betas` and `state_input`. Mamba-1 pins `mamba_selective_scan_0`, `mamba_selective_state_update_0`, `ssm_state_3`, `ssm_state_0`, `ssm_state_to_0`, `scan_output_5` and `out_1`. Mamba-2 pins `mamba2_chunk_scan_0`, `mamba2_selective_state_update_0`, `new_states_0` and `ssm_states_0`, with each argument's position at the two call sites. The hub dispatcher's closure variables `torch_function` and `implementation` are what `route_kernels` and `needs_torch_kernels` read ([Recurrent mixers](internals/recurrent-mixers.md)). ### Names, not operations A `RENAME` key that resolves nowhere is skipped silently, so a renamed module is caught by `test_standard_names_alias_native_envoys`, not at load. Where a release has renamed a module both spellings are keys, as with GPT-NeoX's `embed_out` and `lm_head`, and nnsight binds whichever resolves. ## How the suite catches a move `test_every_source_value_resolves_on_every_layer` reads every available value on the family's `Attention` with `inside_forward()` on every attention block, one trace per value. A renamed operation fails there with `SourceNotAvailable` naming the value, the path it pinned and the operations the forward has now. `test_written_pattern_moves_the_logits` catches an operation that still resolves but is no longer what the values are mixed with, since a read can be causally inert. `test_interior_shapes` checks `softmax(scores) == pattern`, which catches a `_0` to `_1` slip onto another tensor of the same shape. `test_contribution_identity` catches a moved contribution operation on BLOOM, MPT and Falcon, and the DeltaNet and Mamba tests catch the kernel names, the branch binding, the argument names and the state bindings ([Testing](testing.md)). `support()` cannot know an operation moved, since it does not run the forward, and a `SourceNotAvailable` is raised only inside a trace. ## Upgrading transformers 1. Upgrade in a scratch environment and run `HF_HUB_OFFLINE=1 python -m pytest -q`. Each `SourceNotAvailable` names the value, the operation it pinned and the operations that exist now. 2. For each failing family, `print(model.layers[0].self_attn.source)` on the old and new versions, and inside a trace `print(attn.source.attention_interface_1.source)` for the interface, and diff the labeled listings. Decide whether the value moved (a renamed call, a binding added before it) or the family moved onto or off the shared interface. 3. Update the operation string in the family module, or delete the override when the family now uses the shared interface. Keep the standard name's meaning fixed: the pattern is the tensor the values are mixed with, the scores are the softmax's input, the contribution is what is added to the stream. 4. Re-run the family's file, then the whole suite. Families get no version conditionals: One module targets the version nnterp is developed on, and this procedure moves it. 5. A checkpoint with a config that no longer parses is a test-data problem, not a family problem: Patch the config the way `tests/families/test_olmo3.py` does, into a copy of the snapshot. Releases also move behavior no operation name pins. From 5.18 `generate` calls the model with no `attention_mask` when the batch has no padding, so `attention_mask` installs the all-ones mask it stands for into the call on read. `generate` also encodes a wrapper's images before the first forward, so under `generate` the tower's values come before `vision.image_token_mask`. ## What nnterp needs from nnsight Every envoy keeps a link to the envoy that holds it, as `parent`, and `root` follows the links to the model envoy; a value reads its block through `self.parent` and the model through `self.root`, and a key with a leading `/` is walked from `root`. `envoys=` chooses the envoy class per module at construction, by type first and path suffix second. `eproperty.transform` takes `(view, raw)`, the edited view and the value as served. Beyond the documented surface nnterp also relies on `Mediator.current`, `Mediator.iteration`, `Mediator.occurrence`, `Interleaver.sourced` and the iterator's pin, exactly as [Descriptors](internals/descriptors.md) and [Recurrent mixers](internals/recurrent-mixers.md) describe. nnsight's `rename` must take a dotted alias (`"attn.c_proj": "attn.o_proj"` binds on the attention), which GPT-2's two `c_proj` modules need. Import nnterp, or nnsight, before any `transformers.models...` module in a script and in every test process. `tests/conftest.py` is the one line that enforces it under pytest. # Contributing A change is a branch off `main`, the family's test file run before and after, the whole suite before the pull request, and the docs page the change touches in the same pull request. Code follows nnsight's `STYLE.md` with the additions below, and every value, family and helper is written so that `support()`, the repr and the suite describe it without anyone opening the source. The example is a family for one load, written the way a shipped module is. ```python from types import SimpleNamespace import torch import nnterp from nnterp import StandardizedTransformer from nnterp.components import Attention, EProperty, Layer, Mlp, Residual from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2Block, GPT2MLP RENAME = {"transformer.wte": "embed_tokens", "transformer.h": "layers", "transformer.ln_f": "norm", "ln_1": "input_layernorm", "attn": "self_attn", "ln_2": "post_attention_layernorm"} class Layer(Layer): """GPT-2's decoder block; returns a bare tensor, so the base holds.""" class Attention(Attention): """GPT-2's attention; the shared eager forward and the residual added in the block, so the base holds.""" class Mlp(Mlp): """GPT-2's MLP, with one value of its own.""" @EProperty("c_fc.output", description="The up projection's output, before the activation") def mlp_preactivation(self, value: torch.Tensor) -> Residual: return value def intermediate_size(model): """The MLP width is ``n_inner``, ``None`` meaning four times the hidden size.""" return model.config.n_inner or 4 * model.hidden_size family = SimpleNamespace(RENAME=RENAME, ENVOYS={GPT2Block: Layer, GPT2Attention: Attention, GPT2MLP: Mlp}, intermediate_size=intermediate_size) model = StandardizedTransformer("openai-community/gpt2", family=family, dispatch=True, attn_implementation="eager", device="cpu") print(model.family is family, model.intermediate_size) # True 3072 print(model.support()["mlp.mlp_preactivation"], str(type(model.layers[0].mlp).mlp_preactivation)) # None (mlp_preactivation) -> Residual [batch seq hidden_size]: The up projection's output, before the activation with model.trace("The Eiffel Tower is in"): pre = model.layers[0].mlp.mlp_preactivation.save() print(tuple(pre.shape)) # (1, 7, 3072) ``` A shipped family is the same content as a module under `nnterp/families/` named after the `model_type`, with the transformers import at its top and a module docstring that teaches the family ([Adding a family](adding-a-family.md)). The annotation `-> Residual` is what gives the value its layout and its repr line. ## House style - Present tense, current behavior. The source and the docs say what is: no "used to", no issue numbers, no "fixed", no TODO in prose. Where code exists because something once broke, write the constraint and its reason, not the incident. - A docstring states the contract; a comment says why. A module docstring teaches the family or the concept, a value's docstring says what the tensor is, its layout and how writes behave, and an inline comment names the failure averted or the alternative rejected. A comment that paraphrases the line under it is deleted. - One word, one meaning: family, name, value, contribution, available and unavailable, interface, hybrid, as the [Glossary](../reference/glossary.md) fixes them. - A description on a value is a capitalized sentence without a trailing period that names what the tensor is, which is what `repr(model)` and `support()` show. - One family per module, named after its `model_type`: The file name is the registry entry. Families import their modeling module at the top of the family module only, never from `nnterp/components/` or `nnterp/standardized.py`, so `import nnterp` loads no modeling code. - Every value is annotated with a layout name (`-> Residual`, `-> Pattern`, `-> Keys`), never an inline `Float[Tensor, "..."]`. A name is defined in the file of the envoy that serves it, with a comment above it, and exported from `nnterp.components`. A family that redefines a value writes the base's name, so `layout` and `dims` cannot drift. Axis names are the constants in `nnterp/components/standard.py`, and an axis that is a fact about the model is named as the size that holds it, which lets `axis_size` resolve it. - Every value located inside a forward carries an `unavailable=` predicate (`needs_eager`, `interface_reason`, `needs_torch_kernels`, or the family's own), so `support()` answers before a trace runs and a read raises `Unavailable` rather than `SourceNotAvailable` where the reason is known in advance. A value a family lacks is declared, `attention_probabilities = unavailable("...")`, never omitted. - A root size a config spells its own way is `def (model)` in the family module, and a model that does something of its own after the head is `def project_on_vocab(model, hidden)` there. Neither is a subclass of the root. - One method per value in `FamilySuite`, one file per family. A new value gets a method that reads it, checks its layout and writes it, and runs on every family. A family that lacks it declares `unavailable(...)` and lists the key in its `EXPECTED_UNAVAILABLE`. - `from __future__ import annotations`, `X | None`, `list[str]`. Raise by default, and warn only where the program can proceed. ## Workflow 1. Branch from `main`. 2. Make the change: a family module, a component value, or both, and the test that pins it. 3. Run the family's file, then the whole suite, offline and on the CPU ([Testing](testing.md)): ```bash CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 python -m pytest tests/families/test_gpt2.py -q -p no:cacheprovider # 42 passed, 12 skipped, 2 warnings in 5.72s ``` 4. For a new family, add `tests/families/test_.py` with a tiny checkpoint that is cached offline, preferring the `hf-internal-testing`, `trl-internal-testing`, `yujiepan` and `hf-tiny-v2` organizations. A config that does not parse is patched in a copy of the snapshot, the way `tests/families/test_olmo3.py` does. 5. For a new value, add a method to `FamilySuite` and expect it to run on every family. 6. Update the page under `site/docs/` that the change touches, and run every snippet on it before publishing. A page other than [transformers compatibility](transformers-compatibility.md) names no version. 7. A test that routes kernels restores `"default"` in a `finally`, since the routing is process-wide. A degenerate tiny checkpoint gets `LOAD_KWARGS` or a family-specific override, not a weaker assertion. 8. Commit as `area: lowercase imperative summary`, with a body that says why. The skills in [ndif-team/skills](https://github.com/ndif-team/skills) state the same facts as these pages for an agent. A change that alters what an agent should do, a renamed value, a new condition on one, a new family's quirk, is a change to the matching skill as well. ## Design decisions and open items Design decisions nnterp holds to, stated as what nnterp does not do: - Values are on the envoy that serves them, `model.layers[i].layer_output`, not on the model as an indexed row, and there is no table of them to query apart from `support()` and `order()`. - Nothing runs at construction except the default family's shape check, and the suite carries the causal checks. There is no load-time refusal of a block with a class not in `ENVOYS`: Such a block is a plain `Envoy` with no `layer_output`. - Eager attention is not forced at load, and `support()` says `eager` is needed when the checkpoint runs `sdpa`. - A remote run needs nnterp installed on the server. The family travels by reference, since an `EProperty` cannot be pickled by value. - The kernels of a recurrent mixer are not routed at load: With `flash-linear-attention`, `causal-conv1d` or `mamba_ssm` installed the kernel values are unavailable until the caller routes them. - The sequence axis is 2 on softmax `attention_queries`, `attention_keys` and `attention_values`, as transformers hands them to its interface, and 1 on every other value. Nothing normalizes it. - `hasattr(envoy, name)` and `getattr(envoy, name, None)` raise `Unavailable`, and there is no `require` or `available` helper on the model that answers for a set of values at once. - There is no wider tolerance for half precision: The suite loads each checkpoint in its own dtype and holds the contribution identity and the pattern's row sums to 8 ulp of it, and a bf16 checkpoint that needs float32 for the identity is loaded that way through `LOAD_KWARGS`. Not done yet: - Interventions and display helpers over the values: a logit lens over layers, patchscopes, prompt batches with targets, repeated-prompt builders, a top-k plot, a prompts-to-dataframe table. `nnterp/prompt_utils.py` and `nnterp/nnsight_utils.py` hold target-token probabilities and activation collection only, typed on `StandardizedTransformer`. - Index lists of a hybrid's attention and recurrent blocks, and a block-structure constant per family. The blocks are read off `model.layers` with `getattr(layer, "self_attn", None)`. - `device_map="auto"` as a default: nnsight's default stands. - Mllama's cross-attention blocks, which need their own family and a cross-attention component. Gemma 4's and NemotronH Omni's audio towers, and video values, have no names or values. MoonViT has no family until a matching tiny checkpoint exists. `scatter_host` walks nnsight's private `_named_children`. # API at a glance Every public symbol of nnterp, with the page that explains it and the generated page that holds its signature; `model` is a `StandardizedTransformer` unless a row says otherwise. ## The model classes The constructors of [`StandardizedTransformer`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer) and [`StandardizedVLLM`](../documentation/standardized_vllm.md#nnterp.standardized_vllm.StandardizedVLLM), which share the [`Standardized`](../documentation/standardized.md#nnterp.standardized.Standardized) mixin, take nnsight's keywords plus these. | keyword | `StandardizedTransformer` | `StandardizedVLLM` | page | | --- | --- | --- | --- | | `repo_id` | a Hub repo id, or an already-loaded `torch.nn.Module` | a Hub repo id | [Loading](../guide/loading.md#an-already-loaded-module) | | `family=` | the family for this load, a module or any object with `RENAME` and `ENVOYS`; the registry is not consulted | the same, written against vLLM's module classes | [Registering](../developing/registering.md#one-load) | | `rename=` | extra aliases, merged over the family's `RENAME`; a key given here wins | the same | [Loading](../guide/loading.md#rename-and-envoys) | | `envoys=` | extra entries over nnsight's tensor-parallel envoys and the family's `ENVOYS`; a key given here wins | the same, over nnsight's envoys for vLLM's parallel layers | [Custom values](../developing/custom-values.md#key-on-the-module-type) | | `tokenizer_kwargs=` | attributes set on the loaded tokenizer, such as `padding_side` and `pad_token` | not taken | [Loading](../guide/loading.md#tokenizer_kwargs) | | `**kwargs` | passed to `TransformersModel`: `dispatch`, `device`, `dtype`, `attn_implementation`, `revision`, `trust_remote_code`, `task` (default `"text-generation"`) | passed to nnsight's `VLLM`: `dispatch` and the engine's arguments | [Loading](../guide/loading.md#arguments), [vLLM](../vllm/index.md) | ## Methods and attributes | member | what | where | page | | --- | --- | --- | --- | | [`support(layer=None)`](../documentation/standardized.md#nnterp.standardized.Standardized.support) | which values this checkpoint has: `None`, or the reason, per block without `layer` | outside a trace; nothing runs | [Availability](../guide/values/availability.md) | | [`order(layer=None)`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.order), [`rank(name, layer=None)`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.rank) | the forward order of the available values, name to rank; one value's `(side, rank)` | outside; one probe run on first use; `StandardizedTransformer` only | [Attention: read order](../guide/values/attention.md#read-order) | | [`skip_layers(start, end, skip_with=None)`](../documentation/standardized.md#nnterp.standardized.Standardized.skip_layers) | blocks `start` to `end` inclusive do not run | inside a trace, before block `start` | [Methods](../guide/methods.md#skip_layers) | | [`steer(layers, vector, factor=1.0, token_positions=None, batch_index=None)`](../documentation/standardized.md#nnterp.standardized.Standardized.steer) | adds `factor * vector` to `layer_output` in place | inside; `layers` ascending | [Methods](../guide/methods.md#steer) | | [`project_on_vocab(hidden)`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.project_on_vocab) | the logit lens: the final norm, `lm_head`, and what the model does past the head | inside on a live value, outside on a saved one; inside only on vLLM | [Methods](../guide/methods.md#project_on_vocab) | | [`get_topk_closest_tokens(hidden, k=5)`](../documentation/standardized.md#nnterp.standardized.Standardized.get_topk_closest_tokens) | `project_on_vocab`, softmax, the `k` most likely tokens per position | either | [Methods](../guide/methods.md#get_topk_closest_tokens-and-probs_to_dict) | | [`probs_to_dict(probs, k=5)`](../documentation/standardized.md#nnterp.standardized.Standardized.probs_to_dict) | the `k` most likely tokens of one `[vocab_size]` distribution | either | [Methods](../guide/methods.md#get_topk_closest_tokens-and-probs_to_dict) | | `to_model_key()` | nnsight's remote key, naming `TransformersModel` (or `VLLM`) | outside | [Remote](../guide/remote.md) | | `num_layers`, `hidden_size`, `vocab_size`, `num_heads`, `num_kv_heads`, `head_dim`, `qk_head_dim`, [`intermediate_size`](../documentation/standardized.md#nnterp.standardized.Standardized.intermediate_size) | the sizes, read off the config and read-only; a family defines `def (model)` to spell one otherwise | outside; no `dispatch` needed | [The whole model](../guide/values/whole-model.md#sizes) | | `layers[i].self_attn.num_heads`, `.num_kv_heads`, `.head_dim`, `.qk_head_dim`; `layers[i].mlp.intermediate_size`; `layers[i].mlp.num_experts`, `.top_k` | a block's own sizes, read off the module | either | [The whole model](../guide/values/whole-model.md#sizes) | | [`family`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.family) | the family module the load resolved to, or the `family=` passed | | [Loading](../guide/loading.md) | | `layers`, `embed_tokens`, `norm`, `lm_head`; `layers[i].self_attn`, `.mlp`, `.linear_attn` | the standard names: the root's modules, and a block's `Attention`, `Mlp` and recurrent mixer | | [Names](../guide/names.md) | | [`vision`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.vision), [`projector`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.projector) | the vision tower and the last module before the scatter, on a wrapper whose family names them | | [Vision-language](../guide/architectures/vision-language.md) | | [`add_prefix_false_tokenizer`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.add_prefix_false_tokenizer) | the tokenizer loaded with `add_prefix_space=False`, on first use | outside | [Helpers](../guide/helpers.md#target-token-mass-over-many-prompts) | ## Root values The whole-model values, in forward order; the layouts are `Logits`, `NextTokenProbs` and `Tokens` from `nnterp.standardized` and `Residual` from `nnterp.components`. | value | layout | writes | on `StandardizedVLLM` | page | | --- | --- | --- | --- | --- | | [`input_ids`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.input_ids) | `Tokens`, `[batch, seq]` | assign | read-only, `[1, tokens]` | [The whole model](../guide/values/whole-model.md#the-input) | | [`attention_mask`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.attention_mask) | `Tokens` | assign | unavailable | [The whole model](../guide/values/whole-model.md#the-input) | | [`input_size`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.input_size) | a `torch.Size` | read-only | read-only | [The whole model](../guide/values/whole-model.md#the-input) | | [`token_embeddings`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.token_embeddings) | `Residual` | assign | `[1, tokens, hidden_size]` | [The whole model](../guide/values/whole-model.md#token_embeddings-and-block-0) | | [`final_norm_output`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.final_norm_output) | `Residual` | assign, in place | `[1, tokens, hidden_size]` | [The whole model](../guide/values/whole-model.md) | | [`lm_head_output`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.lm_head_output) | `Logits` | assign, in place | unavailable | [The whole model](../guide/values/whole-model.md#logits-and-lm_head_output) | | [`logits`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.logits) | `Logits` | assign | `[1, 1, vocab_size]`, this step's last position | [The whole model](../guide/values/whole-model.md#logits-and-lm_head_output) | | [`next_token_probs`](../documentation/standardized.md#nnterp.standardized.StandardizedTransformer.next_token_probs) | `NextTokenProbs`, `[batch, vocab_size]` | read-only | `[1, vocab_size]` | [The whole model](../guide/values/whole-model.md#next_token_probs) | ## Block values | value | host | layout | available when | page | | --- | --- | --- | --- | --- | | [`layer_input`](../documentation/components/layer.md#nnterp.components.layer.Layer.layer_input) | `layers[i]` | `Residual` | always | [Residual stream](../guide/values/residual-stream.md) | | [`attention_norm_output`](../documentation/components/layer.md#nnterp.components.layer.Layer.attention_norm_output) | `layers[i]` | `Residual` | the block has an `input_layernorm` | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`layer_mid`](../documentation/components/layer.md#nnterp.components.layer.Layer.layer_mid) | `layers[i]` | `Residual` (`FlatResidual` on OPT) | a stream between the sublayers: not on a parallel block | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`mlp_norm_output`](../documentation/components/layer.md#nnterp.components.layer.Layer.mlp_norm_output) | `layers[i]` | `Residual` (`FlatResidual` on OPT) | the block norms the MLP's input | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`mlp_input`](../documentation/components/layer.md#nnterp.components.layer.Layer.mlp_input) | `layers[i]` | `Residual` | the block has an `mlp` | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`layer_output`](../documentation/components/layer.md#nnterp.components.layer.Layer.layer_output) | `layers[i]` | `Residual` (`Streams` on DeepSeek-V4) | always | [Residual stream](../guide/values/residual-stream.md) | | [`attention_output`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_output) | `self_attn` | `Residual` | always | [Residual stream](../guide/values/residual-stream.md) | | [`attention_queries`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_queries) | `self_attn` | `Queries` | eager attention | [Attention](../guide/values/attention.md) | | [`attention_keys`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_keys) | `self_attn` | `Keys` | eager attention | [Attention](../guide/values/attention.md) | | [`attention_values`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_values) | `self_attn` | `Values` | eager attention | [Attention](../guide/values/attention.md) | | [`attention_scores`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_scores) | `self_attn` | `Pattern` | eager attention | [Attention](../guide/values/attention.md) | | [`attention_probabilities`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_probabilities) | `self_attn` | `Pattern` | eager attention | [Attention](../guide/values/attention.md) | | [`attention_head_outputs`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_head_outputs) | `self_attn` | `HeadOutputs` | eager attention | [Attention](../guide/values/attention.md) | | [`attention_premix`](../documentation/components/attention.md#nnterp.components.attention.Attention.attention_premix) | `self_attn` | `HeadOutputs` | the attention has an `o_proj`; any implementation | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`mlp_output`](../documentation/components/mlp.md#nnterp.components.mlp.Mlp.mlp_output) | `mlp` | `Residual` | the block has an `mlp` | [Residual stream](../guide/values/residual-stream.md) | | [`mlp_activation`](../documentation/components/mlp.md#nnterp.components.mlp.Mlp.mlp_activation) | `mlp` | `Neurons` | the MLP has an `act_fn`; not on a mixture | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`mlp_neurons`](../documentation/components/mlp.md#nnterp.components.mlp.Mlp.mlp_neurons) | `mlp` | `Neurons` | the MLP has a `down_proj`; not on a mixture | [Residual stream](../guide/values/residual-stream.md#between-the-sublayers) | | [`router_logits`](../documentation/components/moe.md#nnterp.components.moe.Moe.router_logits) | `mlp`, a `Moe` | `RouterLogits` | the module runs a mixture | [Mixture of experts](../guide/architectures/mixture-of-experts.md) | | [`expert_weights`](../documentation/components/moe.md#nnterp.components.moe.Moe.expert_weights), [`expert_indices`](../documentation/components/moe.md#nnterp.components.moe.Moe.expert_indices) | `mlp`, a `Moe` | `ExpertWeights`, `ExpertIndices` | the module runs a mixture | [Mixture of experts](../guide/architectures/mixture-of-experts.md) | | [`expert_outputs`](../documentation/components/moe.md#nnterp.components.moe.Moe.expert_outputs) | `mlp`, a `Moe` | `ExpertOutputs` | `experts_implementation` is `grouped_mm` or `batched_mm` | [Mixture of experts](../guide/architectures/mixture-of-experts.md#expert_outputs-needs-the-grouped-experts) | | [`routed_output`](../documentation/components/moe.md#nnterp.components.moe.Moe.routed_output) | `mlp`, a `Moe` | `Residual` | the module runs a mixture | [Mixture of experts](../guide/architectures/mixture-of-experts.md) | | [`shared_expert_output`](../documentation/components/moe.md#nnterp.components.moe.Moe.shared_expert_output) | `mlp`, a `Moe` | `Residual` | the mixture has a shared expert | [Mixture of experts](../guide/architectures/mixture-of-experts.md#read-order) | ## Recurrent mixer values The values of `layers[i].linear_attn` on a [gated DeltaNet](../guide/architectures/delta-net.md) (`LinearAttention`), a [Mamba-1](../guide/architectures/mamba.md) (`SelectiveScan`) and a [Mamba-2](../guide/architectures/mamba-2.md) (`StateSpace`) mixer, with the kernel's own name for the tensor where it has one. | value | [`LinearAttention`](../documentation/components/linear_attention.md#nnterp.components.linear_attention.LinearAttention) | [`SelectiveScan`](../documentation/components/selective_scan.md#nnterp.components.selective_scan.SelectiveScan) | [`StateSpace`](../documentation/components/state_space.md#nnterp.components.state_space.StateSpace) | available when | | --- | --- | --- | --- | --- | | `attention_output` | `Residual` | `Residual` | `Residual` | always | | `attention_queries` | `LinearQK` | `ScanQK` (`C`) | `SSDQueries` (`C`) | the kernels are transformers' pure-torch ones | | `attention_keys` | `LinearQK` | `ScanQK` (`B`) | `SSDKeys` (`B`) | pure-torch kernels | | `attention_values` | `LinearV` | `ScanValues` (`x`) | `SSDValues` (`x`) | pure-torch kernels | | `betas` | `Gates` | `ScanSteps` (`dt`), read-only | `Gates` (`dt`) | pure-torch kernels | | `decays` | `Gates` (`ChannelGates` on Kimi-Linear) | `ScanDecays` (`dt * A`), read-only | `Gates` (`A * dt`) | pure-torch kernels | | `state_input` | `State`, or `None` on a prompt | `ScanState`, or `None` on a prompt | `State`, or `None` on a prompt | pure-torch kernels | | `attention_head_outputs` | `LinearV` | `ScanValues` (`y`) | `SSDHeadOutputs` (`y`) | pure-torch kernels, read inside them on Mamba-1 | | `state_output` | `State` | `ScanState` | `State` | pure-torch kernels, read inside them on Mamba-1 | | [`state`](../documentation/components/recurrent.md#nnterp.components.recurrent.RecurrentMixer.state) | `State`, one occurrence per token | `ScanState`, one per token | unavailable | `route_kernels(family, "torch")` | | [`states`](../documentation/components/recurrent.md#nnterp.components.recurrent.RecurrentMixer.states) | `States`, read-only | `ScanStates`, read-only | `States`, read-only | `route_kernels`; on Mamba-2 `chunk_per_token(model)` | | [`state_after(t)`](../documentation/components/recurrent.md#nnterp.components.recurrent.RecurrentMixer.state_after), [`set_state_after(t, value)`](../documentation/components/recurrent.md#nnterp.components.recurrent.RecurrentMixer.set_state_after) | both | both | `state_after` only | as `states` | ## Vision values The values of `model.vision` on a wrapper loaded with `task="image-text-to-text"`, on the [Vision-language](../guide/architectures/vision-language.md) page; `vision.support()` and `vision.no_images()` answer for them. | value | layout | read at | available when | | --- | --- | --- | --- | | [`image_token_mask`](../documentation/components/vision.md#nnterp.components.vision.Vision.image_token_mask) | `ImageTokenMask`, `[batch, seq]` bool | the model's inputs, so first in a trace; read-only | an image reaches the model and the config names an `image_token_id` | | [`patch_embeddings`](../documentation/components/vision.md#nnterp.components.vision.Vision.patch_embeddings) | `Patches` | `patch_embed.output`, one row per patch | an image reaches the model | | [`tower_output`](../documentation/components/vision.md#nnterp.components.vision.Vision.tower_output) | `Patches` | the last block's stream after `vision.norm` | an image reaches the model | | [`image_features`](../documentation/components/vision.md#nnterp.components.vision.Vision.image_features) | `ImageFeatures`, `[image_tokens, hidden_size]` | the scatter into the token embeddings | an image reaches the model and the family keys an `ImageScatter` | | `vision.layers[i].layer_output` and the other block values | `Patches` for the stream values; the attention interior as on a text block | the tower's blocks ([`VisionLayer`](../documentation/components/vision.md#nnterp.components.vision.VisionLayer), [`VisionAttention`](../documentation/components/vision.md#nnterp.components.vision.VisionAttention), [`VisionMlp`](../documentation/components/vision.md#nnterp.components.vision.VisionMlp)) | an image reaches the model; eager attention for the interior | | `vision.num_layers`, `vision_hidden_size`, `num_heads`, `head_dim`, `intermediate_size`, `patch_size`, `image_size` | sizes, off the tower's config | outside a trace | `image_size` raises `Unavailable` on a tower with no fixed resolution | ## Helpers | function | what | page | | --- | --- | --- | | [`nnsight_utils.get_token_activations(model, prompts, layers=None, get_activations=None, remote=False, idx=None, tracer=None)`](../documentation/nnsight_utils.md#nnterp.nnsight_utils.get_token_activations) | `[len(layers), num_prompts, hidden_size]` at token `idx` (default `-1`) of every prompt | [Helpers](../guide/helpers.md) | | [`nnsight_utils.collect_token_activations_batched(model, prompts, batch_size, ..., tqdm=None, use_session=True)`](../documentation/nnsight_utils.md#nnterp.nnsight_utils.collect_token_activations_batched) | the same over batches; a remote run goes through one session | [Helpers](../guide/helpers.md#another-value-many-prompts) | | [`nnsight_utils.collect_last_token_activations_session(model, prompts, batch_size, ...)`](../documentation/nnsight_utils.md#nnterp.nnsight_utils.collect_last_token_activations_session) | the same inside one `model.session` | [Remote](../guide/remote.md#decided-by-the-deployment) | | [`nnsight_utils.compute_next_token_probs(model, prompt, remote=False)`](../documentation/nnsight_utils.md#nnterp.nnsight_utils.compute_next_token_probs) | `next_token_probs` per prompt, `[num_prompts, vocab_size]` on the CPU | [Helpers](../guide/helpers.md#another-value-many-prompts) | | [`nnsight_utils.layer_output(model, layer)`](../documentation/nnsight_utils.md#nnterp.nnsight_utils.layer_output) | the default `get_activations`: `model.layers[layer].layer_output` | [Helpers](../guide/helpers.md) | | [`prompt_utils.Prompt`](../documentation/prompt_utils.md#nnterp.prompt_utils.Prompt): `from_strings(prompt, target_strings, model_or_tokenizer)`, `has_no_collisions()`, `get_target_probs(probs)`, `run(model, get_probs)` | a prompt with named sets of target token ids | [Helpers](../guide/helpers.md#target-token-mass-over-many-prompts) | | [`prompt_utils.get_first_tokens(words, model_or_tokenizer)`](../documentation/prompt_utils.md#nnterp.prompt_utils.get_first_tokens) | the first token of `word` and of `" word"`, through `add_prefix_false_tokenizer` | [Helpers](../guide/helpers.md#target-token-mass-over-many-prompts) | | [`prompt_utils.run_prompts(model, prompts, batch_size=32, get_probs_func=None, func_kwargs=None, remote=False, tqdm=None)`](../documentation/prompt_utils.md#nnterp.prompt_utils.run_prompts) | each target's mass per prompt, `[num_prompts, layers]` | [Helpers](../guide/helpers.md#target-token-mass-over-many-prompts) | | [`prompt_utils.next_token_probs_unsqueeze(model, prompt, remote=False)`](../documentation/prompt_utils.md#nnterp.prompt_utils.next_token_probs_unsqueeze) | the default `get_probs_func`, with a layer axis of one | [Helpers](../guide/helpers.md#target-token-mass-over-many-prompts) | | [`prompt_utils.TokenizationError`](../documentation/prompt_utils.md#nnterp.prompt_utils.TokenizationError) | a word has no standalone first token under the tokenizer | [Helpers](../guide/helpers.md#target-token-mass-over-many-prompts) | ## Families, kernels and exceptions | name | what | page | | --- | --- | --- | | [`families.lookup(model_type, engine=None)`](../documentation/families/index.md#nnterp.families.lookup) | the family: a registered one, else `nnterp.families.` imported on first use, else `default` with a warning; on another engine `UnsupportedFamily` | [Registering](../developing/registering.md#the-lookup-order) | | [`families.register(family, *model_types, engine=None)`](../documentation/families/index.md#nnterp.families.register) | a family for every later load of those types, or of the type its `__name__` ends in; also overrides a shipped one | [Registering](../developing/registering.md) | | [`families.known(engine=None)`](../documentation/families/index.md#nnterp.families.known), [`families.all_families(engine=None)`](../documentation/families/index.md#nnterp.families.all_families) | the shipped model types (98; 60 on vLLM); every shipped family, imported | [Loading](../guide/loading.md) | | [`families.REGISTRY`](../documentation/families/index.md#nnterp.families.REGISTRY), `families.default`, `families.` | what `register` added; the best-effort family; a family module, imported on first access | [Architecture](../developing/internals/architecture.md#from-model_type-to-a-family) | | [`route_kernels(family, kernel="torch")`](../documentation/components/recurrent.md#nnterp.components.recurrent.route_kernels) | rebind a family's recurrent kernels, process-wide, to transformers' pure-torch functions, or `"default"` back | [Gated DeltaNet](../guide/architectures/delta-net.md#kernels-and-route_kernels) | | [`route_delta_rule(family, kernel="recurrent")`](../documentation/components/recurrent.md#nnterp.components.recurrent.route_delta_rule) | `route_kernels` in the delta rule's words: `"recurrent"` or `"chunked"` | [Recurrent mixers](../developing/internals/recurrent-mixers.md#pure-torch-kernels-and-routing) | | [`chunk_per_token(model, enabled=True)`](../documentation/components/state_space.md#nnterp.components.state_space.chunk_per_token) | set every Mamba-2 mixer's `chunk_size` to 1 on this model, so `states` reads every token | [Mamba-2](../guide/architectures/mamba-2.md#the-state-after-every-token) | | [`axis_size(envoy, axis)`](../documentation/components/standard.md#nnterp.components.standard.axis_size) | the size of a layout axis for a value on `envoy`; `ValueError` for an axis that is not a fact about the model | [Layouts](../guide/values/layouts.md#axes-named-as-sizes) | | [`Unavailable`](../documentation/components/eproperty.md#nnterp.components.eproperty.Unavailable) | a `RuntimeError` at the read or write of a value the checkpoint lacks, before the model runs | [Availability](../guide/values/availability.md#nnterpunavailable-at-the-read) | | [`UnsupportedFamily`](../documentation/families/index.md#nnterp.families.UnsupportedFamily) | a `ValueError` at load: no family covers the `model_type` and the default cannot standardize it | [Loading](../guide/loading.md#a-model_type-with-no-family) | | nnsight's `OutOfOrderError`, `SourceNotAvailable` | a read after the model ran past its location; a `source` path this run does not have | [Attention: read order](../guide/values/attention.md#read-order), [Finding source ops](../developing/finding-source-ops.md#errors-from-a-wrong-name) | ## For family authors What `nnterp.components` exports beyond the layouts, grouped by role. | names | what | page | | --- | --- | --- | | [`Standard`](../documentation/components/standard.md#nnterp.components.standard.Standard) (`values()`, `support()`, `sourced`), [`Layer`](../documentation/components/layer.md#nnterp.components.layer.Layer) (`returns_tuple`, `skip_with`), [`Attention`](../documentation/components/attention.md#nnterp.components.attention.Attention) (`off_interface()`, `SINK`), [`Mlp`](../documentation/components/mlp.md#nnterp.components.mlp.Mlp), [`Moe`](../documentation/components/moe.md#nnterp.components.moe.Moe) (`SCORING`, `no_mixture()`) | the envoy bases a family subclasses and keys on its module types in `ENVOYS` | [Adding a family](../developing/adding-a-family.md) | | [`RecurrentMixer`](../documentation/components/recurrent.md#nnterp.components.recurrent.RecurrentMixer) (`CHUNK_KERNEL`, `RECURRENT_KERNEL`, `STATE_OP`, `STEP_STATE_OP`, `BRANCH`, `SEQ_OP`, `KERNEL`), `LinearAttention`, `SelectiveScan`, `StateSpace` | the recurrent mixer bases | [Recurrent mixers](../developing/internals/recurrent-mixers.md) | | [`Vision`](../documentation/components/vision.md#nnterp.components.vision.Vision), `VisionLayer`, `VisionAttention`, `VisionMlp`, `QwenVision`, `QwenVisionAttention`, `PixtralVision`, [`ImageScatter`](../documentation/components/vision.md#nnterp.components.vision.ImageScatter) | the vision tower's envoys, and the marker class for the module that scatters the image features | [Vision](../developing/internals/vision.md) | | [`EProperty(key=None, description=None, unavailable=None, select=None)`](../documentation/components/eproperty.md#nnterp.components.eproperty.EProperty), [`DerivedEProperty(compute, description=None, unavailable=None)`](../documentation/components/eproperty.md#nnterp.components.eproperty.DerivedEProperty), [`TokenEProperty`](../documentation/components/tokens.md#nnterp.components.tokens.TokenEProperty) | the descriptors a value is made of; the third serves a tensor flat over tokens as this invoke's rows | [Descriptors](../developing/internals/descriptors.md) | | [`unavailable(reason)`](../documentation/components/eproperty.md#nnterp.components.eproperty.unavailable) | the marker a family assigns in place of a value it lacks | [Overriding a value](../developing/overriding-values.md#a-value-the-family-does-not-have) | | `needs_eager`, `interface_reason`, `mixture_reason`, `needs_grouped_experts`, `no_shared_expert`, `needs_torch_kernels`, `needs_recurrent_routing`, `needs_kernel_source`, `needs_token_loop`, `needs_per_token_chunks`, `image_token_id` | `unavailable=` predicates, a function of the envoy returning the reason or `None` | [Overriding a value](../developing/overriding-values.md#a-value-the-family-does-not-have) | | `PARALLEL`, `POST_NORM_ATTENTION`, `POST_NORM_MLP`, `NOT_ON_INTERFACE`, `PER_IMAGE` | reason strings for the blocks a family marks | [Overriding a value](../developing/overriding-values.md#a-value-the-family-does-not-have) | | `INTERFACE`, `LOGITS`, `DISPATCH`, `PER_SLOT` | operation names: the shared attention call, the router's projection, the experts call and the per-slot view | [transformers compatibility](../developing/transformers-compatibility.md#the-shared-interface) | | `first_tensor`, `rewrap`, `seq_first`, `per_call`, `pinned`, `blocks_support` | a tuple output's tensor and its inverse, a heads-first transpose, the once-per-call record, the occurrence pin, what `support()` builds the block rows with | [Descriptors](../developing/internals/descriptors.md#select-and-writes) | | `BATCH`, `SEQ`, `TOKENS`, `HIDDEN`, `VOCAB`, `HEADS`, `KV_HEADS`, `HEAD_DIM`, `QK_HEAD_DIM`, `INTERMEDIATE`, `QUERY`, `KEY`, `EXPERTS`, `TOP_K`, `STREAMS`, `MIXER_HEADS`, `KEY_DIM`, `VALUE_DIM`, `CHANNELS`, `STATE_DIM`, `GROUPS`, `IMAGES`, `PATCHES`, `IMAGE_TOKENS`, `VISION_HIDDEN` | the axis names every layout is composed from | [Layouts](../guide/values/layouts.md#axes-named-as-sizes) | ## Layouts The 34 layout names `nnterp.components` exports and the root's three in `nnterp.standardized`, each defined in the generated page of the envoy that serves it. | layout | axes | carried by | defined in | | --- | --- | --- | --- | | `Residual` | `batch seq hidden_size` | the stream values, the contributions, the module boundaries, `token_embeddings`, `final_norm_output`, `routed_output`, `shared_expert_output` | [layer](../documentation/components/layer.md) | | `FlatResidual` | `tokens hidden_size` | OPT's `layer_mid`, `mlp_norm_output` | [layer](../documentation/components/layer.md) | | `Streams`, `StreamWeights`, `StreamMixing` | `batch seq streams hidden_size`; `batch seq streams`; `batch seq streams streams` | DeepSeek-V4's `layer_output`, `layer_input`; its `attention_post`, `mlp_post`; its `attention_comb`, `mlp_comb` | [layer](../documentation/components/layer.md) | | `Logits`, `NextTokenProbs`, `Tokens` | `batch seq vocab_size`; `batch vocab_size`; `batch seq` (integer) | `logits`, `lm_head_output`; `next_token_probs`; `input_ids`, `attention_mask` | [standardized](../documentation/standardized.md) | | `Queries` | `batch num_heads seq qk_head_dim` | `attention_queries` | [attention](../documentation/components/attention.md) | | `Keys`, `Values` | `batch num_kv_heads seq qk_head_dim`; `batch num_kv_heads seq head_dim` | `attention_keys`; `attention_values` | [attention](../documentation/components/attention.md) | | `Pattern` | `batch num_heads query key` | `attention_scores`, `attention_probabilities` | [attention](../documentation/components/attention.md) | | `HeadOutputs` | `batch seq num_heads head_dim` | `attention_head_outputs`, `attention_premix` | [attention](../documentation/components/attention.md) | | `Neurons` | `batch seq intermediate_size` | `mlp_activation`, `mlp_neurons` | [mlp](../documentation/components/mlp.md) | | `RouterLogits` | `batch seq num_experts` | `router_logits` | [moe](../documentation/components/moe.md) | | `ExpertWeights`, `ExpertIndices` | `batch seq top_k` (`ExpertIndices` integer) | `expert_weights`, `expert_indices` | [moe](../documentation/components/moe.md) | | `ExpertOutputs` | `batch seq top_k hidden_size` | `expert_outputs` | [moe](../documentation/components/moe.md) | | `State`, `States` | `batch heads key_dim value_dim`; `batch seq heads key_dim value_dim` | a DeltaNet's or Mamba-2's `state_input`, `state_output`, `state`; `states` | [recurrent](../documentation/components/recurrent.md) | | `LinearQK`, `LinearV` | `batch seq heads key_dim`; `batch seq heads value_dim` | a DeltaNet's queries, keys; its values, head outputs | [linear_attention](../documentation/components/linear_attention.md) | | `Gates`, `ChannelGates` | `batch seq heads`; `batch seq heads key_dim` | `decays`, `betas` on a DeltaNet or Mamba-2; Kimi-Linear's `decays` | [linear_attention](../documentation/components/linear_attention.md) | | `ScanQK` | `batch seq groups state_dim` | a Mamba-1's queries (`C`), keys (`B`) | [selective_scan](../documentation/components/selective_scan.md) | | `ScanValues`, `ScanSteps` | `batch seq channels` | a Mamba-1's values (`x`), head outputs (`y`); its `betas` | [selective_scan](../documentation/components/selective_scan.md) | | `ScanDecays` | `batch seq channels state_dim` | a Mamba-1's `decays` | [selective_scan](../documentation/components/selective_scan.md) | | `ScanState`, `ScanStates` | `batch channels state_dim`; `batch seq channels state_dim` | a Mamba-1's `state_input`, `state_output`, `state`; `states` | [selective_scan](../documentation/components/selective_scan.md) | | `SSDQueries`, `SSDKeys` | `batch seq groups state_dim` | a Mamba-2's queries (`C`), keys (`B`) | [state_space](../documentation/components/state_space.md) | | `SSDValues`, `SSDHeadOutputs` | `batch seq heads head_dim` | a Mamba-2's values (`x`), head outputs (`y`) | [state_space](../documentation/components/state_space.md) | | `Patches` | `images patches vision_hidden_size` | a tower block's stream values, `patch_embeddings`, `tower_output` | [vision](../documentation/components/vision.md) | | `ImageTokenMask`, `ImageFeatures` | `batch seq` (bool); `image_tokens hidden_size` | `image_token_mask`; `image_features` | [vision](../documentation/components/vision.md) | The terms these tables use are defined in the [Glossary](glossary.md). # Glossary The terms this site uses, each in the one sense the pages give it, with the page that explains it. nnsight's own terms (envoy, invoke, occurrence, source) are defined only as far as nnterp leans on them, and [nnsight.net](https://nnsight.net) has the rest. **Alias**: a second attribute that nnsight's `rename=` installs on an envoy, so `model.layers[5]` and GPT-2's `model.transformer.h[5]` are one object and the native name keeps working. Every standard name is one ([Names](../guide/names.md)). **Attention sink**: a per-head column that joins the softmax and is dropped afterwards (GPT-OSS, MiMo-V2-Flash, DeepSeek-V4), so the pattern's rows sum to less than one ([Attention](../guide/values/attention.md)). **Availability**: whether this checkpoint has a value. `support()` reports it before any trace, and a read or write of a value the checkpoint lacks raises `Unavailable` with the same reason. The descriptor's `unavailable=`, a reason string or a predicate of the envoy, decides it per envoy ([Availability](../guide/values/availability.md)). **Axis**: one named dimension of a layout, such as `batch`, `seq` or `num_heads`. An axis that is a fact about the model is named as the size that holds it, and `nnterp.axis_size(envoy, axis)` resolves it ([Layouts](../guide/values/layouts.md#axes-named-as-sizes)). **Block**: one decoder layer, `model.layers[i]`, a `Layer` envoy. The residual stream enters it as `layer_input`, each sublayer adds its contribution, and it leaves as `layer_output` ([Residual stream](../guide/values/residual-stream.md)). **Checkpoint**: one Hub repository, such as `openai-community/gpt2`. Each family's test pins a tiny random checkpoint ([Testing](../developing/testing.md)). **Contribution**: what a sublayer adds to the residual stream, `attention_output` and `mlp_output`, defined by the contribution identity `layer_input + attention_output + mlp_output == layer_output`. On a sandwich or post-norm block it is the post-norm's output, and on BLOOM the tensor entering the residual add ([Residual stream](../guide/values/residual-stream.md)). **Decode step**: one forward call of `model.generate` after the prompt call, producing one token. Under `tracer.iter`, step 0 is the prompt call and every later step a decode step ([Generation](../guide/generation.md)). **Default family**: `nnterp.families.default`, the best-effort family that a `model_type` with no family loads with, and a warning. It serves `layer_output`, the root values, the sizes and the attention interior, and refuses the contributions and `project_on_vocab` ([Loading](../guide/loading.md#a-model_type-with-no-family)). **Derived value**: a `DerivedEProperty`, computed at read time from other served values and read-only: `next_token_probs`, `states`, a Mamba-1 mixer's `betas` and `decays`. Assigning one raises `AttributeError` ([Custom values](../developing/custom-values.md)). **Dispatch**: `dispatch=True` loads the weights at construction. Without it the model is built on the meta device, with no weights, and loads at the first trace, while `support()` and the sizes read the config either way ([Loading](../guide/loading.md#arguments)). **Drill**: instrumenting a function that the forward calls, through `.source`, so the operations inside it can be read. It works inside a trace only, and an `EProperty` whose path has two `source` segments drills before every read ([Finding source ops](../developing/finding-source-ops.md#inside-a-call)). **Eager attention**: transformers' `attn_implementation="eager"`, the Python attention forward the attention interior is read inside. The default is the checkpoint's, usually `sdpa`, and the interior is then unavailable ([Attention](../guide/values/attention.md#eager-attention)). **Encyclopedia**: one page per supported family, built from the family module: the block drawn with the standard names, every value and where it is read, the `support()` report per public checkpoint, and the family's quirks ([encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/)). **Envoy**: nnsight's proxy for one module, reached by attribute path, with `.input`, `.output` and `.source` ([nnsight.net](https://nnsight.net/features/6_modules/)). nnterp's `Layer`, `Attention`, `Mlp` and the other hosts are `Envoy` subclasses installed through `envoys=`. **EProperty**: the descriptor every value is: nnsight's `eproperty` with a key that is a path from the host (`"output"`, `"../ln_2.output"`, `"source..inputs"`) and an `unavailable=` reason ([Descriptors](../developing/internals/descriptors.md)). **Expert**: one of a mixture's routed feed-forward networks, `num_experts` per block ([Mixture of experts](../guide/architectures/mixture-of-experts.md)). **Family**: the module under `nnterp/families/` for one `model_type`, such as `nnterp.families.gpt2`, declaring `RENAME`, its `Layer`, `Attention` and `Mlp` subclasses and `ENVOYS`. `model.family` is the one a checkpoint resolved to ([Adding a family](../developing/adding-a-family.md)). **Forward order**: the order in which the forward reaches the values' locations: the inputs, then each block's interior before its outputs, block 3 before block 5, `logits` last. `model.order(i)` gives it for block `i` ([Attention](../guide/values/attention.md#read-order)). **Getter**: the function an `EProperty` decorates, the preprocess: what it returns is what a read serves ([Overriding a value](../developing/overriding-values.md#a-value-that-has-to-be-a-copy)). **Grouped-query attention**: fewer key/value heads than query heads. `attention_keys` and `attention_values` are read before `repeat_kv`, so their head axis is `num_kv_heads` wide ([Attention](../guide/values/attention.md#grouped-query-attention)). **Head outputs**: `attention_head_outputs`, each head's mix of the values before the output projection, `[batch, seq, num_heads, head_dim]` ([Attention](../guide/values/attention.md)). **Host**: the envoy a value is declared on and its path is walked from: `Layer` for `layer_output`, `Attention` for the pattern, the root for `logits` ([Custom values](../developing/custom-values.md)). **Hybrid**: a family whose blocks are of two kinds, `linear_attn` on most blocks and `self_attn` on a few (Qwen3-Next, Qwen3.5, OLMo-Hybrid, Kimi-Linear), or Nemotron-H's and Bamba's mixer, attention and MLP blocks. `support()` reads per block there ([Availability](../guide/values/availability.md#three-forms)). **Image features**: `vision.image_features`, the tensor the wrapper scatters into the token embeddings at the image tokens, `[image_tokens, hidden_size]`, so that `layers[0].input[vision.image_token_mask] == vision.image_features` ([Vision-language](../guide/architectures/vision-language.md)). **Image token mask**: `vision.image_token_mask`, `input_ids == image_token_id` as a `[batch, seq]` bool, read off the inputs and so first in a trace ([Vision-language](../guide/architectures/vision-language.md)). **Interface**: transformers' shared attention call, `attention_interface_1` (`nnterp.components.INTERFACE`), which under eager attention is `eager_attention_forward`. The base `Attention` reads the interior on and inside it, and a family with its own arithmetic relocates the names ([transformers compatibility](../developing/transformers-compatibility.md#the-shared-interface)). **Invoke**: nnsight's `tracer.invoke(prompt)`, one prompt batch inside a trace ([nnsight.net](https://nnsight.net/features/8_batching/)). A mixture's routing values serve each invoke its own rows, on vLLM an invoke is one prompt, and a trace with images carries one image-carrying invoke ([Mixture of experts](../guide/architectures/mixture-of-experts.md#ablating-and-rerouting)). **Kernel**: the function a recurrent mixer's forward calls to run the sequence: `torch_chunk_gated_delta_rule` on a prompt and `torch_recurrent_gated_delta_rule` on a decode step for a gated DeltaNet, the scan and the state update for Mamba. Everything but `attention_output` is read at that call ([Gated DeltaNet](../guide/architectures/delta-net.md#kernels-and-route_kernels)). **Latent attention**: DeepSeek-V2's and V3's attention, where queries and keys are `qk_head_dim` wide and values `head_dim`, the one case where the two sizes differ ([Attention](../guide/values/attention.md)). **Layout**: the one axis layout a value has on every family, a named `jaxtyping` type such as `Residual` or `Pattern`, read as `value.layout` with its axes as `value.dims` ([Layouts](../guide/values/layouts.md)). **Mixer**: the sublayer that mixes information across positions, `self_attn` or `linear_attn`. On the architecture pages it is short for recurrent mixer ([Cross-family sweep](../tutorials/cross-family-sweep.md)). **Module boundary**: a value read where a module takes or returns a tensor, between the block's stream values: `attention_norm_output`, `layer_mid`, `mlp_norm_output` and `mlp_input` ([Residual stream](../guide/values/residual-stream.md#between-the-sublayers)). **Name**: an entry in the vocabulary, the one set of module names every family answers to: `embed_tokens`, `layers[i]`, `layers[i].self_attn`, `layers[i].mlp`, `norm`, `lm_head`, with `linear_attn` on a hybrid and `vision` and `projector` on a wrapper. A name says which module; a value says which tensor ([Names](../guide/names.md)). **Native name**: the module's name in the checkpoint's own architecture, `transformer.h[i].attn` on GPT-2. It keeps working, `envoy.path` is always the native path, and above the host an `EProperty` path takes native names only ([Names](../guide/names.md)). **Occurrence**: nnsight's count of how many times a run has reached one location. `tracer.iter[t]` binds a read to occurrence `t`, which is how `state` is walked token by token ([Recurrent mixers](../developing/internals/recurrent-mixers.md#the-state-after-every-token)). **Operation**: one call or binding inside a module's forward as nnsight's `.source` labels it, `{callable}_{occurrence}`: `nn_functional_softmax_0`, `attention_interface_1`. A `source.` path names one ([Finding source ops](../developing/finding-source-ops.md)). **Parallel block**: a block where one norm feeds both sublayers and `x + attn(norm(x)) + mlp(norm(x))` is summed at the end (GPT-J, Phi, Falcon, Pythia). It has no `layer_mid`, and the contribution identity holds to rounding ([Availability](../guide/values/availability.md#the-reasons)). **Pattern**: `attention_probabilities`, the distribution the values are mixed with, after the softmax and its dropout, `[batch, num_heads, query, key]`. `attention_scores` is the same tensor before the softmax ([Attention](../guide/values/attention.md)). **Pinned read**: a read inside `tracer.iter[t]`, bound to occurrence `t` of its location. A read outside any loop is relaxed and takes the occurrence in flight ([Recurrent mixers](../developing/internals/recurrent-mixers.md#the-state-after-every-token)). **Postprocess**: the hook that runs on an assigned value, and what it returns is written to the location. The base `layer_output` puts a tensor back into the block's tuple with it ([Overriding a value](../developing/overriding-values.md#a-value-that-has-to-be-a-copy)). **Premix**: `attention_premix`, `o_proj.input` split by head: every head's output as the output projection takes it, served under any `attn_implementation` ([Residual stream](../guide/values/residual-stream.md#between-the-sublayers)). **Projector**: `model.projector`, the last module before the scatter on a wrapper. `projector.input` is what the tower hands over, which is not `tower_output` on every host ([Vision-language](../guide/architectures/vision-language.md#the-projector)). **Read order**: the order of the reads within one trace, which has to follow forward order. A read out of order raises nnsight's `OutOfOrderError` in a plain trace and binds the next step's occurrence under `tracer.iter` ([Generation](../guide/generation.md#a-read-out-of-order-inside-the-loop)). **Reason**: the string `support()` gives for a value a checkpoint lacks and `Unavailable` carries, from the descriptor's `unavailable=`: a fixed string, or a predicate of the envoy evaluated on the instance ([Availability](../guide/values/availability.md#the-reasons)). **Recurrent mixer**: a sublayer that mixes positions through a recurrent state rather than a pattern: a gated DeltaNet (`LinearAttention`), a Mamba-1 (`SelectiveScan`) or a Mamba-2 (`StateSpace`) mixer, all `RecurrentMixer` subclasses, served as `layers[i].linear_attn` under attention's names ([Gated DeltaNet](../guide/architectures/delta-net.md)). **Registry**: `nnterp.families`: `lookup(model_type)` returns what `register` put in `REGISTRY`, else `nnterp.families.` imported on first use, else the default family ([Registering](../developing/registering.md#the-lookup-order)). **Residual stream**: the `[batch, seq, hidden_size]` tensor a block passes to the next, `layer_input` entering and `layer_output` leaving ([Residual stream](../guide/values/residual-stream.md)). **Router**: a mixture's gate, `moe.router`, whose `router_logits` are scored to pick each token's experts ([Mixture of experts](../guide/architectures/mixture-of-experts.md)). **Routing**: on a mixture, which experts a token goes to, the pair `expert_weights` and `expert_indices`, `[batch, seq, top_k]` ([Mixture of experts](../guide/architectures/mixture-of-experts.md)). The [Recurrent mixers](../developing/internals/recurrent-mixers.md#pure-torch-kernels-and-routing) page also says routing for `route_kernels`, the process-wide rebinding of a family's kernels to transformers' pure-torch functions. **Sandwich block**: a block that norms a sublayer's output before adding it, `x + post_attention_layernorm(attn(input_layernorm(x)))` (Gemma-2, Gemma-3). The contribution is the post-norm's output, as it is on a post-norm block (OLMo-2, OLMo-3, EXAONE-4), which has no pre-norm ([Overriding a value](../developing/overriding-values.md#a-sibling-modules-output)). **Scatter**: the step in a wrapper's forward that writes the projected image features into the token embeddings at the image tokens, a `masked_scatter`. `image_features` is read there ([Vision](../developing/internals/vision.md#where-image_features-is-read)). **Size**: a read-only attribute read off the config: the root's `num_layers`, `hidden_size`, `vocab_size`, `num_heads`, `num_kv_heads`, `head_dim`, `qk_head_dim` and `intermediate_size`, each a `StandardizedProperty` that a family can define instead. A block's own sizes are on its `self_attn` and `mlp`, read off the module ([The whole model](../guide/values/whole-model.md#sizes)). **Slot**: one of a token's `top_k` routing choices, the third axis of `expert_weights`, `expert_indices` and `expert_outputs` ([Mixture of experts](../guide/architectures/mixture-of-experts.md)). **Softcap**: Gemma-2's `final_logit_softcapping`, `cap * tanh(logits / cap)` after `lm_head`. `logits` and `project_on_vocab` include it, and `lm_head_output` does not ([The whole model](../guide/values/whole-model.md#logits-and-lm_head_output)). **Source**: nnsight's `.source`, the operations of a module's forward ([nnsight.net](https://nnsight.net/features/11_source/)). A source-located value has a `source` segment in its path and is read inside the forward, before the module's own output ([Finding source ops](../developing/finding-source-ops.md)). **State**: a recurrent mixer's memory: `state_input` entering a call, `state_output` leaving it, `state` after one token and `states` after every token, `[batch, heads, key_dim, value_dim]` on a gated DeltaNet or Mamba-2 and `[batch, channels, state_dim]` on Mamba-1 ([Gated DeltaNet](../guide/architectures/delta-net.md#the-state-after-every-token)). **Sublayer**: the attention (or recurrent mixer) and the MLP of a block, each adding one contribution. `self_attn.input` and `mlp.input` are what enters each, whatever norm produced it ([Names](../guide/names.md#the-norms)). **Support**: `model.support()`, every value with `None` or a per-block `{layer: reason}`; `model.support(layer=i)`, one block, flat; `envoy.support()`, one envoy. Computed from the config and the module tree, nothing runs ([Availability](../guide/values/availability.md#three-forms)). **Tower**: `model.vision`, the vision encoder of a wrapper, a `Vision` envoy with its own `layers`, sizes, `patch_embeddings` and `tower_output`. Its blocks' stream values have the `Patches` layout ([Vision-language](../guide/architectures/vision-language.md)). **Transform**: the hook that runs on the served view after an in-place edit and returns what the model receives when the block moves on. Falcon's `mlp_output`, a copy, carries edits back with one ([Overriding a value](../developing/overriding-values.md#a-value-that-has-to-be-a-copy)). **Tuple-returning block**: a block whose forward returns `(hidden_states, ...)` (GPT-J, GPT-Neo, BLOOM, MPT, Falcon), whose family sets `Layer.returns_tuple`. `layer_output` is the tensor either way, and `skip_layers` packs the replacement as the block would ([Residual stream](../guide/values/residual-stream.md#a-tensor-on-every-family)). **Unavailable**: `nnterp.Unavailable`, a `RuntimeError` raised at the read or write of a value the checkpoint lacks, before the model runs, with the reason `support()` gives. `unavailable("...")` is the marker a family assigns in place of a value it does not have ([Availability](../guide/values/availability.md#nnterpunavailable-at-the-read)). **UnsupportedFamily**: `nnterp.UnsupportedFamily`, a `ValueError` raised at load when no family covers the `model_type` and the default family cannot standardize the checkpoint ([Loading](../guide/loading.md#a-model_type-with-no-family)). **Value**: a standard value: a tensor served by an `EProperty` on a host and meaning the same thing on every family, such as `layer_output` or `attention_probabilities`. It is listed in `print(model)`, reported by `support()`, and read, edited in place or assigned the way `.output` is ([Residual stream](../guide/values/residual-stream.md#reading-editing-in-place-assigning)). **Wrapper**: the image-text-to-text class around a text model, such as `LlavaForConditionalGeneration`. Its family is the text model's, from `text_config.model_type`, and `model.vision` serves anything only when it is loaded with `task="image-text-to-text"` ([Vision-language](../guide/architectures/vision-language.md)). Every symbol these terms name is on [API at a glance](api-at-a-glance.md). # About nnterp nnterp gives every transformer family the same module names and the same standard values, on top of [nnsight](https://nnsight.net). It is developed by the [NDIF](https://ndif.us/) team and maintained at [github.com/ndif-team/nnterp](https://github.com/ndif-team/nnterp). The [encyclopedia](https://ndif-team.github.io/nnterp/encyclopedia/) is one page per supported family, built from the family modules themselves. The skills that teach an agent nnterp are in [ndif-team/skills](https://github.com/ndif-team/skills). nnterp 1.x was written by Clément Dumas and described in [nnterp: A Standardized Interface for Mechanistic Interpretability of Transformers](https://arxiv.org/abs/2511.14465). nnterp 2 is a rebuild on nnsight 0.8: the names are nnsight aliases, the values are properties on the envoys, and each family is one module that says how its own names map onto the standard ones. ## Citation ```bibtex @misc{dumas2025nnterp, title={nnterp: A Standardized Interface for Mechanistic Interpretability of Transformers}, author={Cl{\'e}ment Dumas}, year={2025}, eprint={2511.14465}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2511.14465}, } ``` ## Found a bug? Reduce it to a minimal example first. If the same code written with nnsight's `TransformersModel` on the native module names also fails, the bug is in nnsight: open it on the [nnsight tracker](https://github.com/ndif-team/nnsight/issues). Otherwise open an issue on the [nnterp tracker](https://github.com/ndif-team/nnterp/issues).