nnterp · encyclopedia olmo3 · nnterp.families.olmo3

family olmo3

OLMo 3OLMo 3.1

OLMo-2's post-norm block: attention and MLP read the raw stream and the stream receives the post-norms' outputs; three blocks in four attend over a 4096-token window with the plain rotary, and the fourth attends over the whole prefix with a YaRN-scaled rotary.

Post-norms onlyQuery/key normsSliding window

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block; this family's blocks differ by layer_types.

embed_tokenstoken_embeddings
→
layers× 32
→
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 Olmo3ForCausalLMlayoutwhat 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] Olmo3DecoderLayerlayoutwhat 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 Olmo3Attentionlayoutwhat 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 Olmo3MLPlayoutwhat 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-3-1025-7B: the native modules, standard names and every standard value with its layout.

Olmo3ForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=100278, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100278, 4096, padding_idx=100277)
  (layers): ModuleList(
    (0-31): 32 x Olmo3DecoderLayer(
      (self_attn): Olmo3Attention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (q_norm): Olmo3RMSNorm((4096,), eps=1e-06)
        (k_norm): Olmo3RMSNorm((4096,), 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): Olmo3MLP(
        (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)
        (up_proj): Linear(in_features=4096, out_features=11008, bias=False)
        (down_proj): Linear(in_features=11008, out_features=4096, 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): Olmo3RMSNorm((4096,), eps=1e-06)
      (post_feedforward_layernorm): Olmo3RMSNorm((4096,), 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): Olmo3RMSNorm((4096,), 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_layers32
model.hidden_size4096
model.num_heads32
model.num_kv_heads32
model.head_dim128
model.qk_head_dim128
model.vocab_size100278
model.intermediate_size11008
architecturesOlmo3ForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size11008
num_attention_heads32
num_key_value_heads32
vocab_size100278
max_position_embeddings65536
sliding_window4096
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
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: here it follows the attention, and nothing norms the MLP's input. There is no input_layernorm. self_attn.input is layers[i].input itself and mlp.input is layers[i].input + attention_output, both unnormalized, so a probe trained on a sublayer's input sees the stream at its own scale.

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

Each post-norm divides by the RMS of what it reads, so a uniform scale of mlp.output or self_attn.output[0] cancels in the contribution except through eps (1e-6). Zeroing a module's output still removes its contribution, since the norm of zero is zero. Partial ablation and steering of a sublayer's effect go through the contribution:

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

Direct logit attribution and any per-sublayer decomposition use the post-norms' outputs as the terms.

Three sliding-window blocks for every full one

config.layer_types is sliding_attention on three blocks and full_attention on the fourth, repeating: blocks 3, 7, 11, ... are full, 8 of 32 on 7B and 16 of 64 on 32B, and the last block is full. On a sliding block a query attends to its own position and the 4095 before it (sliding_window 4096); on a full block, to the whole prefix. On a prompt shorter than the window the two masks coincide, and the two kinds of block are told apart by their rotary, below, not by their patterns. Read a block's kind outside the trace:

kinds = [model.layers[i].self_attn._module.sliding_window for i in range(model.num_layers)]
# 4096 on a sliding block, None on a full one

The rotary differs by block kind

The config keeps one rotary per kind in rope_parameters. Sliding blocks use the plain rotary, rope_theta 500000. Full blocks use YaRN, factor 8 over original_max_position_embeddings 8192: the lowest rotary frequencies are divided by 8, the highest are kept, and attention_factor 1.2079 multiplies the cosines and sines. attention_queries and attention_keys are read after the rotary, so on a full block each query and key head has 1.2079 times the norm it had leaving q_norm or k_norm, and the scores are 1.459 (1.2079²) times what the rotation without the factor would give. Both kinds apply the same head_dim ** -0.5 query scale. Every checkpoint's config sets these rotary parameters.

Queries and keys are normed across all heads

q_norm and k_norm are RMSNorms over the whole projection, num_heads * head_dim (4096 on 7B, 5120 on 32B) and num_kv_heads * head_dim, applied before the heads are split and before the rotary. One RMS covers every head, so an edit to one head's slice of q_proj.output reaches the others. 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

7B has 32 key/value heads for 32 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

The attention interior needs attn_implementation="eager"; the default load runs sdpa and reports the six values unavailable. There is no softcap and no sink.

The readout and the tokenizer

logits is lm_head.output, with no cap or scale after it, and project_on_vocab on the last block's layer_output equals logits. The final norm is a plain RMSNorm whose gain is model.norm._module.weight, eps 1e-6. lm_head and embed_tokens are separate weights on every checkpoint. The base tokenizer prepends nothing, so model.input_ids is the prompt's tokens alone; <|endoftext|> (id 100257) is its end-of-text token.

Intermediate checkpoints

The base repositories keep their training checkpoints as Hub branches named stage1-step<N>, stage2-step<N> and stage3-step<N>. Olmo-3-1025-7B has 1,421 stage-1 branches, from stage1-step0 to stage1-step1413814, 52 stage-2 and 13 stage-3. Olmo-3-1125-32B has 655 stage-1 branches, two stage-2 runs (stage2-ingredient1-step<N>, stage2-ingredient2-step<N>) and their merge, and the stage-3 branches. Each loads with revision=:

model = StandardizedTransformer("allenai/Olmo-3-1025-7B", revision="stage1-step1000")

From the family module

04 · nnterp/families/olmo3.py

OLMo 3 (Olmo3ForCausalLM).

Llama's names plus a post-norm block: every sublayer is normed after, not before, 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 3 has only the post-norms, and mixes sliding-window and full attention layers: no input_layernorm exists to alias.