Coverage for transformer_lens/model_bridge/supported_architectures/olmo2.py: 100%

1"""OLMo 2 architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

6from transformer_lens.model_bridge.generalized_components import ( 

7 BlockBridge, 

8 EmbeddingBridge, 

9 LinearBridge, 

10 PositionEmbeddingsAttentionBridge, 

11 RMSNormalizationBridge, 

12 RotaryEmbeddingBridge, 

13 UnembeddingBridge, 

14) 

15 

16 

17class Olmo2ArchitectureAdapter(ArchitectureAdapter): 

18 """Architecture adapter for OLMo 2 models. 

19 

20 OLMo 2 uses a post-norm architecture with RMSNorm, Q/K normalization in attention, 

21 rotary position embeddings (RoPE), and gated MLP (SwiGLU). Key differences from 

22 pre-norm models like Llama: 

23 

24 - Post-norm: RMSNorm is applied AFTER attention and AFTER MLP, not before. 

25 ln1 maps to post_attention_layernorm, ln2 maps to post_feedforward_layernorm. 

26 - Q/K normalization: Per-head RMSNorm applied to queries and keys after projection. 

27 - No biases on any projections. 

28 

29 Optional Parameters (may not exist in state_dict): 

30 ------------------------------------------------- 

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

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

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

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

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

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

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

38 - blocks.{i}.ln1.b - RMSNorm has no bias 

39 - blocks.{i}.ln2.b - RMSNorm has no bias 

40 - ln_final.b - RMSNorm has no bias 

41 """ 

42 

43 # Attention bridge seam; EXAONE-4 swaps in its NoPE-gating variant. 

44 _attention_bridge_cls = PositionEmbeddingsAttentionBridge 

45 

46 def __init__(self, cfg: Any) -> None: 

47 """Initialize the OLMo 2 architecture adapter.""" 

48 super().__init__(cfg) 

49 

50 self._set_rms_rotary_defaults() 

51 # OLMo-2 uses post-norm (RMSNorm AFTER attention/MLP), so layer norm 

52 # folding into QKV/MLP weights is incorrect — the norms apply to the 

53 # output, not the input. Same pattern as BERT and Phi-3. 

54 self.supports_fold_ln = False 

55 # Force eager attention for numerical consistency with benchmark reference. 

56 # PositionEmbeddingsAttentionBridge delegates to native HF attention, so 

57 # both bridge and reference must use the same implementation. 

58 self.cfg.attn_implementation = "eager" 

59 

60 self.weight_processing_conversions = { 

61 **self._qkvo_weight_conversions(), 

62 } 

63 

64 # Component mapping — POST-NORM architecture: 

65 # ln1 = post_attention_layernorm (applied AFTER attention) 

66 # ln2 = post_feedforward_layernorm (applied AFTER MLP) 

67 self.component_mapping = { 

68 "embed": EmbeddingBridge(name="model.embed_tokens"), 

69 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

70 "blocks": BlockBridge( 

71 name="model.layers", 

72 submodules={ 

73 "ln1": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

74 "ln2": RMSNormalizationBridge( 

75 name="post_feedforward_layernorm", config=self.cfg 

76 ), 

77 "attn": self._attention_bridge_cls( 

78 name="self_attn", 

79 config=self.cfg, 

80 submodules={ 

81 "q": LinearBridge(name="q_proj"), 

82 "k": LinearBridge(name="k_proj"), 

83 "v": LinearBridge(name="v_proj"), 

84 "o": LinearBridge(name="o_proj"), 

85 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

86 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

87 }, 

88 requires_attention_mask=True, 

89 requires_position_embeddings=True, 

90 ), 

91 "mlp": self._build_mlp_bridge(), 

92 }, 

93 # Post-norm overrides: ln1/ln2 are applied AFTER attention/MLP and 

94 # BEFORE the residual add, so the residual mid-point is mlp.hook_in 

95 # and the additive contributions are the norm outputs, not the raw 

96 # module outputs (attn.hook_out / mlp.hook_out stay raw). 

97 hook_alias_overrides={ 

98 "hook_resid_mid": "mlp.hook_in", 

99 "hook_attn_out": "ln1.hook_out", 

100 "hook_mlp_out": "ln2.hook_out", 

101 }, 

102 # No pre-MLP norm: the MLP consumes the mid-residual directly, so 

103 # the hook_mlp_in capture must sit on the MLP, not on ln2. 

104 mlp_reads_resid_directly=True, 

105 ), 

106 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

107 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), 

108 } 

109 

110 def _build_mlp_bridge(self): 

111 """MLP bridge seam; FlexOlmo swaps in the MoE variant.""" 

112 return self._gated_mlp()