olmo2

olmo2

OLMo 2 (Olmo2ForCausalLM).

Llama's names with the norms moved after the sublayers, 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 2 has only the post-norms: no input_layernorm exists to alias.

RENAME module-attribute

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

ENVOYS module-attribute

ENVOYS = {Olmo2DecoderLayer: Layer, Olmo2Attention: Attention, Olmo2MLP: Mlp}