transformer_lens.model_bridge.supported_architectures.olmo module

OLMo architecture adapter.

class transformer_lens.model_bridge.supported_architectures.olmo.OlmoArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for OLMo (v1) models.

OLMo v1 uses a pre-norm architecture with a custom non-learnable LayerNorm (fixed weight=1, bias=0), rotary position embeddings (RoPE), and gated MLP (SwiGLU). Key differences from later OLMo variants:

  • Pre-norm: LayerNorm is applied BEFORE attention and BEFORE MLP.

  • Non-learnable LayerNorm: Weight and bias are not trainable parameters. Delegating to HF’s native forward via NormalizationBridge handles this correctly.

  • No Q/K normalization in attention.

  • Optional QKV clipping (applied out-of-place by the reconstructed attention forward when config.clip_qkv is set).

Optional Parameters (may not exist in state_dict):

  • blocks.{i}.attn.b_Q - No bias on query projection

  • blocks.{i}.attn.b_K - No bias on key projection

  • blocks.{i}.attn.b_V - No bias on value projection

  • blocks.{i}.attn.b_O - No bias on output projection

  • blocks.{i}.mlp.b_in - No bias on MLP up_proj

  • blocks.{i}.mlp.b_gate - No bias on MLP gate_proj

  • blocks.{i}.mlp.b_out - No bias on MLP down_proj

__init__(cfg: Any) → None

Initialize the OLMo architecture adapter.