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

1"""Qwen3MoE (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 Qwen3MoeArchitectureAdapter(ArchitectureAdapter): 

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

21 

22 Qwen3MoE is a sparse MoE decoder-only Transformer, structurally close to OLMoE. 

23 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: 128 experts with top-8 routing (public 30B-A3B checkpoints). 

28 - Batched expert parameters: gate_up_proj and down_proj as single 3D tensors, 

29 not a ModuleList. 

30 - final_rms=True (Qwen3-style; OLMoE uses False). 

31 - No biases on any projections. 

32 - GQA: n_key_value_heads < n_heads in all public checkpoints. 

33 

34 Only the all-MoE configuration is supported (decoder_sparse_step=1, 

35 mlp_only_layers=[]). Models with dense fallback layers cannot be wrapped 

36 because MoEBridge does not handle the dense Qwen3MoeMLP path. 

37 

38 Optional Parameters (may not exist in state_dict): 

39 ------------------------------------------------- 

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

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

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

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

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

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

46 - ln_final.b - RMSNorm has no bias 

47 """ 

48 

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

50 """Initialize the Qwen3MoE architecture adapter.""" 

51 super().__init__(cfg) 

52 

53 self._set_rms_rotary_defaults() 

54 # Force eager attention for output_attentions hook support 

55 self.cfg.attn_implementation = "eager" 

56 self.cfg.default_prepend_bos = False # Qwen3 family convention 

57 

58 # QKVO rearrangements; MoE expert and gate weights pass through unchanged 

59 self.weight_processing_conversions = { 

60 **self._qkvo_weight_conversions(), 

61 } 

62 

63 # Component mapping — PRE-NORM architecture: 

64 # ln1 = input_layernorm (applied BEFORE attention) 

65 # ln2 = post_attention_layernorm (applied BEFORE MLP) 

66 # Deliberate mirror of olmoe.py / minimax_m2.py: same wiring by structural 

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

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

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": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

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

77 "attn": PositionEmbeddingsAttentionBridge( 

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 # Qwen3MoeSparseMoeBlock stores experts as batched 3D tensors 

92 # rather than a ModuleList. MoEBridge wraps the entire block and 

93 # delegates to HF's native forward — same pattern as OLMoE. 

94 "mlp": MoEBridge( 

95 name="mlp", 

96 config=self.cfg, 

97 sparse_required=("gate",), 

98 submodules={ 

99 # Dense fallback layers (mlp_only_layers / 

100 # decoder_sparse_step) have no router; their 

101 # projections bind gated-MLP neuron hooks (#1645). 

102 "gate": MoERouterBridge(name="gate", optional=True), 

103 "dense_gate": LinearBridge(name="gate_proj", optional=True), 

104 "dense_in": LinearBridge(name="up_proj", optional=True), 

105 "dense_out": LinearBridge(name="down_proj", optional=True), 

106 }, 

107 ), 

108 }, 

109 ), 

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

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

112 }