View as Markdown

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.<model_type> 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, 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:

"""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 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). 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:

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:

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 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:

# 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

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:

"""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 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 is how to find the one it does.