nnterp · encyclopedia olmoe · nnterp.families.olmoe

family olmoe

OLMoE

Llama's pre-norm block with queries and keys RMS-normed across all heads, and a mixture of 64 experts on every block that weights a token's 8 experts by their softmax entries without renormalizing them, so the 8 weights sum to less than one.

Query/key normsMixture of expertsUnnormalized routing

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 OlmoeForCausalLMlayoutwhat 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] OlmoeDecoderLayerlayoutwhat 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 OlmoeAttentionlayoutwhat 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 OlmoeSparseMoeBlocklayoutwhat 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
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output

Printout

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

OlmoeForCausalLM(
  (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 OlmoeDecoderLayer(
      (self_attn): OlmoeAttention(
        (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): OlmoeRMSNorm((2048,), eps=1e-05)
        (k_norm): OlmoeRMSNorm((2048,), eps=1e-05)
        (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): OlmoeSparseMoeBlock(
        (router/gate): OlmoeTopKRouter()
        (experts): OlmoeExperts(
          (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
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): OlmoeRMSNorm((2048,), eps=1e-05)
      (post_attention_layernorm): OlmoeRMSNorm((2048,), eps=1e-05)
      (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): OlmoeRMSNorm((2048,), eps=1e-05)
  (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_size1024
architecturesOlmoeForCausalLM
num_hidden_layers16
hidden_size2048
intermediate_size1024
num_attention_heads16
num_key_value_heads16
vocab_size50304
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
num_experts64
num_experts_per_tok8
norm_topk_probFalse
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))       # q_norm, k_norm inside
out = h + mlp(post_attention_layernorm(h))    # 8 of 64 experts per token

Llama's pre-norm block and Llama's names: post_attention_layernorm is the MLP's input norm here, where on OLMo-2 the same name is a norm after the attention. Nothing norms a sublayer's output, so attention_output is o_proj's output and mlp_output the mixture's, and scaling either module's output scales what the block adds. The identity is the plain sum, exact in bfloat16:

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)

Queries and keys are normed across all heads

q_norm and k_norm are RMSNorms over the whole projection, 2048 wide on every checkpoint (16 heads of 128), applied before the heads are split and before the rotary embedding; attention_queries and attention_keys are read after both. One RMS covers every head, so an edit to one head's slice of q_proj.output reaches the others: tripling head 0's slice at block 8 changes head 1's queries by 17% on 0924. 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 checkpoint has 16 key/value heads for 16 query heads, so there is no grouping, and clip_qkv, the optional clamp on queries, keys and values, is off (null) on all of them.

Load with eager for the attention interior only

The attention interior needs attn_implementation="eager". The default sdpa load computes the same function: in float32 the two loads' logits differ by at most 2e-5 on 0924. There is no softcap, sink or window. 0924's weights are stored in bfloat16 (13.8 GB). 0125's base checkpoint is stored in float32 (torch_dtype in its config), so a default load of it takes about 28 GB. A float32 copy of 0924 is on its fp32 branch.

The router's weights are the softmax's top 8, not renormalized

The router takes a softmax over all 64 logits in float32, keeps the 8 largest, and casts them back to the model's dtype; norm_topk_prob is false on every checkpoint, so nothing rescales the 8 to sum to one. expert_weights is exactly those 8 probabilities, slot 0 the largest, and there is no shared expert: routed_output is mlp_output.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()       # [batch, seq, num_experts], before the softmax
    w = moe.expert_weights.save()           # [batch, seq, top_k]
    idx = moe.expert_indices.save()

top = logits.float().softmax(-1).topk(moe.top_k, dim=-1)
torch.testing.assert_close(top.values.to(w.dtype), w)    # the softmax's own entries
assert torch.equal(top.indices, idx)                     # slots in descending weight

On 3,072 tokens of English documentation and Python, a token's 8 weights sum to between 0.19 and 0.97. The median sum is 0.30 to 0.37 on blocks 0 to 8 and rises to 0.55 on block 15, so on 14 of the 16 blocks the median token's 8 chosen experts hold less than half its softmax mass. The median routing entropy is 3.45 to 3.99 nats, against 4.16 for a uniform router. Each expert's output is scaled by its own probability, so mlp_output's size moves with how much mass the router puts on the 8 it chose. A logit written at router_logits changes both the choice and the weights: forcing 8 experts to +10 and the rest to −10 gives each about 1/8.

Ablating an expert touches only the tokens routed to it

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the token's other 7 weights as they were. At the ablated block mlp_output changes on exactly the tokens that chose e, and an expert no token chose has an effect of exactly zero:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    idx = moe.expert_indices.save()
    clean = moe.mlp_output.save()
with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)
    ablated = moe.mlp_output.save()

changed = (ablated != clean).any(-1)                       # [batch, seq]
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The equality held at every expert of blocks 4, 8 and 12 on 0924 in float32, and at block 8 in bfloat16. On "The Eiffel Tower is in the city of" (p(Paris) = 0.87) the largest single-expert effect on log p(Paris) is 0.04, 0.06 and 0.07 nats at blocks 4, 8 and 12; zeroing block 8's whole mlp_output moves it by 0.02. Rerouting a slot by writing expert_indices keeps the old slot's weight and changes the expert's group sizes, so other tokens move by rounding (6e-5 in bfloat16).

Expert usage is uneven and routing changes token to token

On the same 3,072 tokens, every block's busiest expert takes 4.5 to 6.6 times its even share (1/64 of the slots), and its least-used one under a fifth of it; blocks 0 to 4 used all 64 experts, and deeper blocks left up to 4 unused on this text. Adjacent tokens share 2.4 of their 8 experts on average at block 0, and 3.8 to 4.3 at blocks 4, 8 and 15.

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. lm_head and embed_tokens are separate weights. The tokenizer is GPT-NeoX's: it prepends nothing, so model.input_ids is the prompt's tokens alone, and <|endoftext|> (id 50279) is the end of text. The SFT, DPO and Instruct chat templates start with bos_token, id 50279. On 0125-SFT, -DPO and -Instruct the tokenizer swaps two names: id 50279 decodes as |||IP_ADDRESS||| and the literal string <|endoftext|> encodes to id 0, so write the id, not the string, when you build a prompt by hand.

Intermediate checkpoints

allenai/OLMoE-1B-7B-0924 keeps 244 pretraining checkpoints as Hub branches, from step5000-tokens20B to step1220000-tokens5117B; main is annealed from step1200000-tokens5033B. allenai/OLMoE-1B-7B-0125 keeps three stage-2 runs, five checkpoints each, as branches named stage2-ingredient<run>-step<N>-tokens<M>B. Each loads with revision=:

model = StandardizedTransformer("allenai/OLMoE-1B-7B-0924", revision="step10000-tokens41B")

The OLMoE paper (arXiv 2409.02060) analyses these checkpoints' routing: router saturation over training, expert co-activation, and domain and vocabulary specialization.

From the family module

04 · nnterp/families/olmoe.py

OLMoE (OlmoeForCausalLM).

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. The attention norms its queries and keys (q_norm, k_norm) across all heads before the rotary embedding, and optionally clamps them (clip_qkv); both happen before the interface call, so the queries and keys are read after them. Every MLP is a sparse mixture of experts that returns the routed hidden states as a bare tensor; its experts are intermediate_size wide.