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

1"""GLM-4.5 MoE architecture adapter. 

2 

3Supports GLM-4.5/4.6/4.7 mixture-of-experts families (`Glm4MoeForCausalLM`). 

4 

5Key features: 

6- RMSNorm with partial pre-norm layout. 

7- RoPE-style rotary embeddings (partial RoPE supported by Hugging Face model logic). 

8- Q/K normalization blocks (`q_norm`, `k_norm`) and GQA / MQA handling. 

9- Sparse MoE block in `model.layers[i].mlp`, with optional dense-prefix layers. 

10- QKVO rearrangements for bridge-side attention hooks. 

11 

12Optional Parameters (may not exist in state_dict): 

13------------------------------------------------- 

14- blocks.{i}.mlp.gate - absent on dense-prefix layers before sparse MoE starts. 

15""" 

16 

17from typing import Any 

18 

19from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

20from transformer_lens.model_bridge.generalized_components import ( 

21 BlockBridge, 

22 EmbeddingBridge, 

23 LinearBridge, 

24 MoEBridge, 

25 MoERouterBridge, 

26 PositionEmbeddingsAttentionBridge, 

27 RMSNormalizationBridge, 

28 RotaryEmbeddingBridge, 

29 UnembeddingBridge, 

30) 

31 

32 

33class Glm4MoeArchitectureAdapter(ArchitectureAdapter): 

34 """Architecture adapter for GLM-4.5 / 4.6 / 4.7 MoE decoder models. 

35 

36 GLM-4x MoE families use RMSNorm, RoPE and sparse routing, with early 

37 dense-MLP layers in some checkpoints. The dense layers are represented by 

38 a present-but-slightly-thinner `mlp` sub-module where routing is absent. 

39 """ 

40 

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

42 """Initialize the GLM-4 MoE architecture adapter.""" 

43 super().__init__(cfg) 

44 

45 self._set_rms_rotary_defaults() 

46 # Force eager attention for output_attentions / compatibility-path parity. 

47 self.cfg.attn_implementation = "eager" 

48 # GLM-4 defaults do not prepend BOS in current tiny checkpoints. 

49 self.cfg.default_prepend_bos = False 

50 

51 # QKVO rearrangements; MoE experts and gate are passed through unchanged. 

52 self.weight_processing_conversions = { 

53 **self._qkvo_weight_conversions(), 

54 } 

55 

56 self.component_mapping = { 

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

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

59 "blocks": BlockBridge( 

60 name="model.layers", 

61 submodules={ 

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

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

64 "attn": PositionEmbeddingsAttentionBridge( 

65 name="self_attn", 

66 config=self.cfg, 

67 submodules={ 

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

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

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

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

72 # Present only when use_qk_norm=True (config default False). 

73 "q_norm": RMSNormalizationBridge( 

74 name="q_norm", config=self.cfg, optional=True 

75 ), 

76 "k_norm": RMSNormalizationBridge( 

77 name="k_norm", config=self.cfg, optional=True 

78 ), 

79 }, 

80 requires_attention_mask=True, 

81 requires_position_embeddings=True, 

82 ), 

83 # Dense prefix layers (idx < first_k_dense_replace) expose 

84 # `mlp` with no router; the dense_* projections bind there 

85 # and carry the gated-MLP neuron hooks. 

86 "mlp": MoEBridge( 

87 name="mlp", 

88 config=self.cfg, 

89 sparse_required=("gate",), 

90 submodules={ 

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

92 # Dense-layer projections (present only on the 

93 # dense layers of this interleaved stack); their 

94 # presence is what makes MoEBridge bind gated-MLP 

95 # neuron hooks there (#1645). 

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

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

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

99 }, 

100 ), 

101 }, 

102 ), 

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

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

105 }