nnterp · encyclopedia olmo2 · nnterp.families.olmo2

family olmo2

OLMo 2

Llama's tree with the norms moved after the sublayers: attention and MLP read the raw residual stream, what reaches the stream is a post-norm's output, and the queries and keys are RMS-normed across all heads before the rotary embedding.

Post-norms onlyQuery/key norms

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 16
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model Olmo2ForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] Olmo2DecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
⚠︎attention_norm_output scalar Unavailable: a post-norm block: the attention reads the residual stream itself, layer_input

nowhere: the family does not serve it
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module mlp's input
⚠︎mlp_norm_output scalar Unavailable: a post-norm block: the MLP reads the residual stream itself, layer_mid

nowhere: the family does not serve it
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn Olmo2Attentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream: the post-attention norm's output

the sibling module post_attention_layernorm's output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp Olmo2MLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream: the post-feedforward norm's output

the sibling module post_feedforward_layernorm's output

Printout

What print(model) shows on allenai/OLMo-2-0425-1B: the native modules, standard names and every standard value with its layout.

Olmo2ForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 2048, padding_idx=100277)
  (layers): ModuleList(
    (0-15): 16 x Olmo2DecoderLayer(
      (self_attn): Olmo2Attention(
        (q_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (k_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (v_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (o_proj): Linear(in_features=2048, out_features=2048, bias=False)
        (q_norm): Olmo2RMSNorm((2048,), eps=1e-06)
        (k_norm): Olmo2RMSNorm((2048,), eps=1e-06)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: the post-attention norm's output
        (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_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
      )
      (mlp): Olmo2MLP(
        (gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (up_proj): Linear(in_features=2048, out_features=8192, bias=False)
        (down_proj): Linear(in_features=8192, out_features=2048, bias=False)
        (act_fn): SiLUActivation()
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream: the post-feedforward norm's output
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
      )
      (post_attention_layernorm): Olmo2RMSNorm((2048,), eps=1e-06)
      (post_feedforward_layernorm): Olmo2RMSNorm((2048,), eps=1e-06)
      (attention_norm_output): Unavailable: a post-norm block: the attention reads the residual stream itself, layer_input
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output): Unavailable: a post-norm block: the MLP reads the residual stream itself, layer_mid
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): Olmo2RMSNorm((2048,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers16
model.hidden_size2048
model.num_heads16
model.num_kv_heads16
model.head_dim128
model.qk_head_dim128
model.vocab_size100352
model.intermediate_size8192
architecturesOlmo2ForCausalLM
num_hidden_layers16
hidden_size2048
intermediate_size8192
num_attention_heads16
num_key_value_heads16
vocab_size100352
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypefloat32
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = x + post_attention_layernorm(self_attn(x))
out = h + post_feedforward_layernorm(mlp(h))

Two RMSNorms on the block, both after a sublayer, none before. post_attention_layernorm has Llama's name and the opposite place: on Llama it is the norm before the MLP, here it follows the attention, and nothing norms the MLP's input. There is no input_layernorm, so code that hooks input_layernorm.output to read what the attention sees has nothing to attach to.

The sublayers read the raw stream

self_attn.input is layers[i].input itself and mlp.input is layers[i].input + attention_output, unnormalized. A probe or dictionary trained on a sublayer's input sees the stream at its own scale, which grows with depth.

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn_in = model.layers[1].self_attn.input.save()           # equals x
    attn = model.layers[1].self_attn.attention_output.save()
    mlp_in = model.layers[1].mlp.input.save()                  # equals x + attn

The attention's contribution depends only on the direction of its input. q_norm and k_norm remove the scale from the queries and keys, the value and output projections are linear with no bias, and post_attention_layernorm removes the scale again: doubling what the attention reads changes attention_output by about 1e-5 relative on 1B. A steering vector added to the stream reaches the attention only by turning the stream's direction. The MLP has no such invariance: doubling what it reads changes mlp_output by about 20 to 50% on 1B's blocks 2, 8 and 14.

The contributions are the post-norms' outputs

attention_output and mlp_output are what the block adds, the outputs of post_attention_layernorm and post_feedforward_layernorm, not of the modules. The identity holds exactly, in float32 and bfloat16 alike:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(x + attn + mlp, out)

RMSNorm is scale-invariant, so scale the contribution, not the module. On 1B, halving mlp.output at block 8 moves layer_output by at most 4e-4, while halving mlp_output moves it by 1.2; halving self_attn.output[0] moves it by less than 1e-4. Zeroing the module output does work, because the norm of zero is zero. Partial ablation and steering of a sublayer's effect go through mlp_output / attention_output:

with model.trace(prompt):
    model.layers[1].mlp.mlp_output[:] *= 0.5     # halves what the block adds

with model.trace(prompt):
    model.layers[1].mlp.output[:] *= 0.5         # close to a no-op on 1B

Direct logit attribution and any per-sublayer decomposition use the post-norm outputs as the terms. The post-norms' gains are small and set the size of each term: their means run from about 0.08 to 0.37 on the attention side and from about 0.15 to 0.62 on the MLP side across 1B's 16 blocks, largest in the last blocks.

Queries and keys are normed across all heads

q_norm and k_norm are RMSNorms over the whole projection, num_heads * head_dim (2048 on 1B) and num_kv_heads * head_dim, applied before the heads are split and before the rotary embedding. attention_queries and attention_keys are read after both. Because one RMS is shared by every head, an edit to one head's slice of q_proj.output reaches the other heads: tripling head 0's slice at block 8 of 1B changes head 1's queries by about 15%. Edit a single head at attention_queries or attention_keys, which no norm follows.

with model.trace(prompt):
    model.layers[1].self_attn.attention_queries[:, 0] = 0   # head 0 only

Every size but 32B has as many key/value heads as query heads; 32B has 40 query heads over 8 key/value heads, so there an edit to key/value head j reaches query heads 5j to 5j + 4.

Load with eager for the attention interior only

The attention interior needs attn_implementation="eager". The default sdpa load computes the same function: on 1B in float32 the two loads' logits differ by at most 5e-5. There is no softcap, sink or window.

The readout is lm_head after a plain RMSNorm

logits is lm_head.output, with no cap or scale after it, and project_on_vocab is lm_head(norm(h)), equal to logits on the last block's layer_output. The final norm's gain is model.norm._module.weight itself, between 0.92 and 5.11 on 1B (mean 2.26), and eps is 1e-6. lm_head and embed_tokens are separate weights on every size.

The tokenizer adds no BOS

The base tokenizer prepends nothing: model.input_ids is the prompt's tokens alone. <|endoftext|> (id 100257) is both bos_token and eos_token. The chat templates of the 1B, 7B and 13B instruct checkpoints start with it; 32B-Instruct's does not. token_embeddings is the unscaled lookup and equals layers[0].input.

Intermediate checkpoints and kin

Each base repository keeps its pretraining checkpoints as Hub branches, stage1-step<N>-tokens<M>B and stage2-ingredient<k>-step<N>-tokens<M>B, which load with revision=:

model = StandardizedTransformer(
    "allenai/OLMo-2-0425-1B", revision="stage1-step10000-tokens21B"
)

The older olmo family is a pre-norm Llama block with a weightless LayerNorm, olmo3 keeps this block and mixes sliding-window and full attention, and olmo_hybrid puts gated DeltaNet blocks between post-norm attention blocks; each is its own nnterp family.

From the family module

04 · nnterp/families/olmo2.py

OLMo 2 (Olmo2ForCausalLM).

Llama's names with the norms moved after the sublayers, x + post_attention_layernorm(attn(x)) and then + post_feedforward_layernorm(mlp(x)). What the block adds is the post-norm's output, not the module's, so the contributions point at the sibling norms. OLMo 2 has only the post-norms: no input_layernorm exists to alias.