nnterp · encyclopedia olmo · nnterp.families.olmo

family olmo

OLMoOLMo 0424OLMo 0724

Llama's pre-norm block with LayerNorm in place of RMSNorm, and every norm without weight or bias, so a sublayer reads the stream standardized to mean 0 and variance 1, and the 0424 and 0724 base checkpoints clamp queries, keys and values to ±8 inside the attention.

LayerNorm

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 OlmoForCausalLMlayoutwhat 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] OlmoDecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
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 OlmoAttentionlayoutwhat 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 module's own 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 OlmoMLPlayoutwhat 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 module's own output

Printout

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

OlmoForCausalLM(
  (lm_head): Linear(in_features=2048, out_features=50304, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50304, 2048, padding_idx=1)
  (layers): ModuleList(
    (0-15): 16 x OlmoDecoderLayer(
      (self_attn): OlmoAttention(
        (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)
        (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
      )
      (mlp): OlmoMLP(
        (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_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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): OlmoLayerNorm()
      (post_attention_layernorm): OlmoLayerNorm()
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (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): OlmoLayerNorm()
  (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_size50304
model.intermediate_size8192
architecturesOlmoForCausalLM
num_hidden_layers16
hidden_size2048
intermediate_size8192
num_attention_heads16
num_key_value_heads16
vocab_size50304
max_position_embeddings2048
hidden_actsilu
tie_word_embeddingsTrue
dtypefloat32
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

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

Llama's block and Llama's names, with LayerNorm where Llama has RMSNorm. post_attention_layernorm is named for its place on the stream, after the attention's add: it is the MLP's input norm. Nothing norms a sublayer's output.

The norms have no weight and no bias

OlmoLayerNorm is F.layer_norm over the hidden axis in float32, eps 1e-5, with no weight and no bias, cast back to the input dtype. The module holds no parameters: list(model.norm._module.parameters()) is empty, and so is every block's input_layernorm and post_attention_layernorm. What a sublayer reads is the stream standardized per position, mean 0 and variance 1:

import torch.nn.functional as F

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn_in = model.layers[1].self_attn.input.save()

torch.testing.assert_close(attn_in, F.layer_norm(x, x.shape[-1:], eps=1e-5))   # float32

Each norm subtracts the mean over the hidden axis, so a shift along the all-ones direction never reaches a sublayer or the logits. On 1B-0724 in float32, adding 10 to every coordinate of layers[8].input, more than 20 times the stream's RMS there, moves the logits by about 1e-4. The final norm has no gain either, so project_on_vocab(h) is lm_head(layer_norm(h)), and a direction's logit effect is read off lm_head alone.

The contributions are the modules' outputs

attention_output is self_attn.output[0] and mlp_output is mlp.output, added to the stream with nothing in between. The identity is the plain sum, exact in float32:

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)

Scaling a module's output scales its contribution. The inputs are the norms' outputs: self_attn.input is input_layernorm.output and mlp.input is post_attention_layernorm.output.

clip_qkv clamps queries, keys and values on 0424 and 0724

clip_qkv is 8.0 on OLMo-7B-0424-hf, OLMo-1B-0724-hf and OLMo-7B-0724-hf, and null on the other checkpoints, the 0424 and 0724 SFT and Instruct ones included. Where it is set, the attention clamps the outputs of q_proj, k_proj and v_proj to [-8, 8], before the heads are split and before the rotary embedding. attention_queries, attention_keys and attention_values are read after the clamp. The rotary mixes pairs of query and key dimensions, so queries and keys can reach 8·√2.

On 1B-0724 every block clamps some query or key entries: up to about 0.1% of query entries and 0.3% of key entries per block, with raw values up to 13.6. Value entries reach the bound only on block 1.

The clamp is in place on the projection's output tensor. A saved q_proj.output therefore holds the clamped values once the trace ends, and an edit written to q_proj.output is clamped after it: clone to keep the raw projection, and edit at attention_queries, which no clamp follows.

attn = model.layers[1].self_attn
with model.trace(prompt):
    raw = attn.q_proj.output.clone().save()     # before the clamp
    clamped = attn.q_proj.output.save()         # the same tensor, clamped in place

Every checkpoint has as many key/value heads as query heads, so there is no grouping.

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-0724 in float32 the two loads' logits differ by at most 1.5e-5, with the same top token at every position. There is no softcap, sink or window; the rotary is the plain one, rope_theta 10000.

The readout and the tokenizer

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits. OLMo-1B-hf ties lm_head to embed_tokens (one parameter), so an edit to embed_tokens.weight is an edit to the unembedding there; every other checkpoint has its own lm_head. The tokenizer is GPT-NeoX's: it prepends nothing, so model.input_ids is the prompt's tokens alone, and <|endoftext|> (id 50279) ends a text. The Instruct chat templates start with eos_token.

Intermediate checkpoints

The base repositories keep their pretraining checkpoints as Hub branches: 351 on OLMo-1B-hf, from step1000-tokens4B to step738020-tokens3094B; 1,446 on OLMo-1B-0724-hf, from step0-tokens0B to step1454000-tokens3048B; 832 on OLMo-7B-0424-hf and 819 on OLMo-7B-0724-hf. Each loads with revision=:

model = StandardizedTransformer("allenai/OLMo-1B-hf", revision="step10000-tokens41B")

From the family module

04 · nnterp/families/olmo.py

OLMo 1 (OlmoForCausalLM).

Llama's tree and Llama's block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, the residual added in the block, attention through the shared eager forward.