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

1"""OLMoE (Mixture of Experts) 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 MoEBridge, 

11 MoERouterBridge, 

12 PositionEmbeddingsAttentionBridge, 

13 RMSNormalizationBridge, 

14 RotaryEmbeddingBridge, 

15 UnembeddingBridge, 

16) 

17 

18 

19class OlmoeArchitectureAdapter(ArchitectureAdapter): 

20 """Architecture adapter for OLMoE (Mixture of Experts) models. 

21 

22 OLMoE uses a pre-norm architecture with RMSNorm, Q/K normalization in attention, 

23 rotary position embeddings (RoPE), and sparse Mixture of Experts MLP. Key features: 

24 

25 - Pre-norm: RMSNorm applied BEFORE attention and BEFORE MLP. 

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

27 - Sparse MoE: 64 experts with top-8 routing (configurable). 

28 - Batched expert parameters: gate_up_proj [num_experts, 2*d_mlp, d_model] and 

29 down_proj [num_experts, d_model, d_mlp] as single tensors, not a ModuleList. 

30 - Optional QKV clipping (handled by HF's native attention forward). 

31 - No biases on any projections. 

32 

33 Optional Parameters (may not exist in state_dict): 

34 ------------------------------------------------- 

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

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

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

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

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

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

41 - ln_final.b - RMSNorm has no bias 

42 """ 

43 

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

45 """Initialize the OLMoE architecture adapter.""" 

46 super().__init__(cfg) 

47 

48 self._set_rms_rotary_defaults(final_rms=False) 

49 # Force eager attention for numerical consistency with benchmark reference 

50 self.cfg.attn_implementation = "eager" 

51 

52 self.weight_processing_conversions = { 

53 **self._qkvo_weight_conversions(), 

54 } 

55 

56 # Component mapping — PRE-NORM architecture: 

57 # ln1 = input_layernorm (applied BEFORE attention) 

58 # ln2 = post_attention_layernorm (applied BEFORE MLP) 

59 # Deliberate mirror of qwen3_moe.py / minimax_m2.py: same wiring by structural 

60 # coincidence, not lineage (norm/router semantics differ per vendor), 

61 # so each file keeps its mapping inline and readable. 

62 self.component_mapping = { 

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

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

65 "blocks": BlockBridge( 

66 name="model.layers", 

67 submodules={ 

68 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

69 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

70 "attn": PositionEmbeddingsAttentionBridge( 

71 name="self_attn", 

72 config=self.cfg, 

73 submodules={ 

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

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

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

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

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

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

80 }, 

81 requires_attention_mask=True, 

82 requires_position_embeddings=True, 

83 ), 

84 # OLMoE uses batched expert parameters (gate_up_proj, down_proj 

85 # as 3D tensors) rather than a ModuleList of individual experts. 

86 # MoEBridge wraps the entire MLP module and delegates to HF's 

87 # native forward pass. 

88 "mlp": MoEBridge( 

89 name="mlp", 

90 config=self.cfg, 

91 submodules={ 

92 "gate": MoERouterBridge(name="gate"), 

93 }, 

94 ), 

95 }, 

96 ), 

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

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

99 }