nnterp · encyclopedia flex_olmo · nnterp.families.flex_olmo

family flex_olmo

FlexOlmo

OLMo-2's post-norm block with a mixture of experts for the MLP: the released checkpoints route every token to every expert, weighted by a softmax over all of them, and the stream receives the post-norm of the weighted sum.

Post-norms onlyQuery/key normsMixture of experts

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× 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 FlexOlmoForCausalLMlayoutwhat 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] FlexOlmoDecoderLayerlayoutwhat 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 FlexOlmoAttentionlayoutwhat 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 FlexOlmoSparseMoeBlocklayoutwhat 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
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/FlexOlmo-7x7B-1T: the native modules, standard names and every standard value with its layout.

FlexOlmoForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=100352, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(100352, 4096, padding_idx=100277)
  (layers): ModuleList(
    (0-31): 32 x FlexOlmoDecoderLayer(
      (self_attn): FlexOlmoAttention(
        (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): FlexOlmoRMSNorm((4096,), eps=1e-06)
        (k_norm): FlexOlmoRMSNorm((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): FlexOlmoSparseMoeBlock(
        (router/gate): FlexOlmoTopKRouter()
        (experts): FlexOlmoExperts(
          (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
        (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
      )
      (post_attention_layernorm): FlexOlmoRMSNorm((4096,), eps=1e-06)
      (post_feedforward_layernorm): FlexOlmoRMSNorm((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): FlexOlmoRMSNorm((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_size100352
model.intermediate_size11008
architecturesFlexOlmoForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size11008
num_attention_heads32
num_key_value_heads32
vocab_size100352
max_position_embeddings4096
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypefloat32
num_experts7
num_experts_per_tok7
norm_topk_probFalse
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))     # mlp: router, experts, weighted sum

OLMo-2's block with the MLP replaced by a mixture on every block. Two RMSNorms, both after a sublayer, none before: self_attn.input is layers[i].input and mlp.input is layers[i].input + attention_output, both unnormalized. post_attention_layernorm follows the attention, where on Llama the same name is the MLP's input norm.

The contributions are the post-norms' outputs

attention_output is post_attention_layernorm's output and mlp_output is post_feedforward_layernorm's. The mixture's own output is routed_output, which equals mlp.output (there is no shared expert), and mlp_output is the post-norm of it. 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()
    routed = model.layers[1].mlp.routed_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)

The post-norm divides by the RMS of what it reads, so a uniform scale of routed_output, or of every expert_weights entry of a token together, cancels in mlp_output except through eps (1e-6). What the routing decides is the mix of the experts' directions, not the size of the term. Scale mlp_output to scale what the block adds.

Every expert runs on every token

num_experts_per_tok equals num_experts on every released checkpoint: 7 of 7 on FlexOlmo-7x7B-1T and -RT, 2 of 2 on each Flex-<domain>-2x7B-1T. The router takes a softmax over all logits in float32 and keeps the top top_k; norm_topk_prob is false, but with every expert kept the weights are the whole softmax and sum to one. expert_indices is then every expert, ordered by weight, slot 0 the largest; which slot an expert takes changes from token to token. The pinned tiny checkpoint keeps 2 of 8, so there the weights sum to less than one.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()       # [batch, seq, num_experts]
    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)

The experts are trained on separate data

Per the model cards, FlexOlmo-7x7B-1T combines a public-mix expert, trained on 1T tokens, with six experts branched from it and trained on 50B tokens each of news, math, code, academic text, creative writing and Reddit; FlexOlmo-7x7B-1T-RT has the same configuration. Each Flex-<domain>-2x7B-1T pairs the public-mix expert with one domain expert. The configs do not say which index holds which expert. Each expert is a SwiGLU MLP intermediate_size (11008) wide.

Removing an expert: the weight or the logit

Zeroing expert e's weight where expert_indices == e removes its term and leaves the others' weights as they were, so a token's weights then sum to 1 - w_e. With every expert on every token, that changes every token. Writing -inf at router_logits for e instead gives e a weight of exactly 0 and renormalizes the softmax over the rest, as if the router had no logit for e:

moe, e = model.layers[1].mlp, 3
with model.trace(prompt):
    moe.router_logits[..., e] = float("-inf")    # weights renormalized over the other experts
    w = moe.expert_weights.save()

with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)   # not renormalized
    w_zeroed = moe.expert_weights.save()

With every expert kept, the two give routed_outputs that differ by the factor 1 / (1 - w_e) per token, which the post-norm divides out: the two edits give the same mlp_output up to eps. On a checkpoint that keeps fewer than all experts (the pinned tiny one) they differ, since -inf lets another expert into the top-k.

Queries and keys are normed across all heads

q_norm and k_norm are RMSNorms over the whole 4096-wide projection, applied before the heads are split and before the rotary, and attention_queries and attention_keys are read after both. An edit to one head's slice of q_proj.output reaches the other heads through the shared RMS; edit a single head at attention_queries. Every checkpoint has 32 key/value heads for 32 query heads.

Loading

The attention interior needs attn_implementation="eager"; expert_outputs needs the default experts_implementation. There is no softcap, sink or window; the rotary is the plain one, rope_theta 500000. The weights are stored in float32: 133 GB for 7x7B, 47 GB for each 2x7B.

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 dolma2: it prepends nothing, so model.input_ids is the prompt's tokens alone, and <|endoftext|> (id 100257) is both bos_token and eos_token.

From the family module

04 · nnterp/families/flex_olmo.py

FlexOlmo (FlexOlmoForCausalLM).

OLMo-2's post-norm block with a mixture of experts for the MLP: 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. Only the post-norms exist: no input_layernorm to alias, and the MLP takes the residual stream after the attention add. The mixture (FlexOlmoSparseMoeBlock) is on every block and returns the routed experts' sum as one tensor; its experts are intermediate_size wide.