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 Glm4MoeRouterBridge(MoERouterBridge): 

34 """Tuple-preserving router bridge for ``Glm4MoeTopkRouter``.""" 

35 

36 

37class Glm4MoeArchitectureAdapter(ArchitectureAdapter): 

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

39 

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

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

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

43 """ 

44 

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

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

47 super().__init__(cfg) 

48 

49 self._set_rms_rotary_defaults() 

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

51 self.cfg.attn_implementation = "eager" 

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

53 self.cfg.default_prepend_bos = False 

54 

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

56 self.weight_processing_conversions = { 

57 **self._qkvo_weight_conversions(), 

58 } 

59 

60 self.component_mapping = { 

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

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

63 "blocks": BlockBridge( 

64 name="model.layers", 

65 submodules={ 

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

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

68 "attn": PositionEmbeddingsAttentionBridge( 

69 name="self_attn", 

70 config=self.cfg, 

71 submodules={ 

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

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

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

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

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

77 "q_norm": RMSNormalizationBridge( 

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

79 ), 

80 "k_norm": RMSNormalizationBridge( 

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

82 ), 

83 }, 

84 requires_attention_mask=True, 

85 requires_position_embeddings=True, 

86 ), 

87 # Dense prefix layers (idx < first_k_dense_replace) expose 

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

89 # and carry the gated-MLP neuron hooks. 

90 "mlp": MoEBridge( 

91 name="mlp", 

92 config=self.cfg, 

93 sparse_required=("gate",), 

94 submodules={ 

95 "gate": Glm4MoeRouterBridge(name="gate", optional=True), 

96 # Dense-layer projections (present only on the 

97 # dense layers of this interleaved stack); their 

98 # presence is what makes MoEBridge bind gated-MLP 

99 # neuron hooks there (#1645). 

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

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

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

103 }, 

104 ), 

105 }, 

106 ), 

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

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

109 }