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

1"""GLM-4 MoE Lite architecture adapter. 

2 

3Supports the GLM-4.7-Flash family (`Glm4MoeLiteForCausalLM`): DeepSeek-style 

4Multi-head Latent Attention (LoRA-compressed Q and KV, nope/rope split heads, 

5interleaved partial RoPE) combined with GLM's sparse MoE — sigmoid router with 

6e_score_correction_bias, batched routed experts, one shared expert — and a 

7per-layer dense/sparse MLP mix declared in ``config.mlp_layer_types``. 

8""" 

9 

10from transformer_lens.model_bridge.generalized_components import MoERouterBridge 

11from transformer_lens.model_bridge.supported_architectures.deepseek_v2 import ( 

12 DeepSeekMLAFamilyArchitectureAdapter, 

13) 

14 

15 

16class Glm4MoeLiteArchitectureAdapter(DeepSeekMLAFamilyArchitectureAdapter): 

17 """GLM-4.7-Flash (Glm4MoeLiteForCausalLM) adapter: DeepSeek-V2 MLA + GLM-4-MoE 

18 routing (dense/sparse per mlp_layer_types).""" 

19 

20 _testing_eager = None 

21 

22 # Public GLM-4.7 checkpoints set q_lora_rank — two-stage LoRA Q compression; 

23 # direct q_proj kept optional for hypothetical uncompressed variants. 

24 q_lora_optional = True 

25 # Verified against zai-org/GLM-4.7-Flash: tokenizer has no BOS token. 

26 prepend_bos = False 

27 

28 def _build_router(self) -> MoERouterBridge: 

29 """Sigmoid + e_score_correction_bias router; absent on dense layers.""" 

30 return MoERouterBridge(name="gate", optional=True)