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

1"""OLMo 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 NormalizationBridge, 

11 PositionEmbeddingsAttentionBridge, 

12 RotaryEmbeddingBridge, 

13 UnembeddingBridge, 

14) 

15 

16 

17class OlmoArchitectureAdapter(ArchitectureAdapter): 

18 """Architecture adapter for OLMo (v1) models. 

19 

20 OLMo v1 uses a pre-norm architecture with a custom non-learnable LayerNorm 

21 (fixed weight=1, bias=0), rotary position embeddings (RoPE), and gated MLP 

22 (SwiGLU). Key differences from later OLMo variants: 

23 

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

25 - Non-learnable LayerNorm: Weight and bias are not trainable parameters. 

26 Delegating to HF's native forward via NormalizationBridge handles this correctly. 

27 - No Q/K normalization in attention. 

28 - Optional QKV clipping (applied out-of-place by the reconstructed 

29 attention forward when config.clip_qkv is set). 

30 

31 Optional Parameters (may not exist in state_dict): 

32 ------------------------------------------------- 

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

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

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

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

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

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

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

40 """ 

41 

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

43 """Initialize the OLMo architecture adapter.""" 

44 super().__init__(cfg) 

45 

46 # Set config variables for weight processing 

47 self.cfg.normalization_type = "LN" 

48 self.cfg.positional_embedding_type = "rotary" 

49 self.cfg.final_rms = False 

50 self.cfg.gated_mlp = True 

51 self.cfg.attn_only = False 

52 self.cfg.uses_rms_norm = False 

53 # Force eager attention for numerical consistency with benchmark reference 

54 self.cfg.attn_implementation = "eager" 

55 

56 n_kv_heads = ( 

57 self.cfg.n_key_value_heads 

58 if self.cfg.n_key_value_heads is not None 

59 else self.cfg.n_heads 

60 ) 

61 

62 self.weight_processing_conversions = { 

63 **self._qkvo_weight_conversions(), 

64 } 

65 

66 # Component mapping — PRE-NORM architecture: 

67 # ln1 = input_layernorm (applied BEFORE attention) 

68 # ln2 = post_attention_layernorm (applied BEFORE MLP) 

69 self.component_mapping = { 

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

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

72 "blocks": BlockBridge( 

73 name="model.layers", 

74 submodules={ 

75 "ln1": NormalizationBridge( 

76 name="input_layernorm", 

77 config=self.cfg, 

78 use_native_layernorm_autograd=True, 

79 ), 

80 "ln2": NormalizationBridge( 

81 name="post_attention_layernorm", 

82 config=self.cfg, 

83 use_native_layernorm_autograd=True, 

84 ), 

85 "attn": PositionEmbeddingsAttentionBridge( 

86 name="self_attn", 

87 config=self.cfg, 

88 submodules={ 

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

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

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

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

93 }, 

94 requires_attention_mask=True, 

95 requires_position_embeddings=True, 

96 ), 

97 "mlp": self._gated_mlp(), 

98 }, 

99 ), 

100 "ln_final": NormalizationBridge( 

101 name="model.norm", 

102 config=self.cfg, 

103 use_native_layernorm_autograd=True, 

104 ), 

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

106 }