flex_olmo

flex_olmo

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.

RENAME module-attribute

RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', 'gate': 'router'}

ENVOYS module-attribute

ENVOYS = {FlexOlmoDecoderLayer: Layer, FlexOlmoAttention: Attention, FlexOlmoSparseMoeBlock: Mlp}

Mlp

Mlp(*args: Any, **kwargs: Any)

Bases: Moe

FlexOlmo's mixture of experts: what reaches the residual stream is the post-feedforward norm's output.

mlp_output

mlp_output(value) -> Residual