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

1"""GLM-4-0414 architecture adapter. 

2 

3Z.ai's GLM-4-0414 family (``Glm4ForCausalLM``: GLM-4-32B-0414, GLM-Z1): 

4the dense GLM decoder (adjacent-pair partial RoPE, joint ``gate_up_proj`` 

5MLP) with Gemma-2-style sandwich norms — extra RMS norms applied to the 

6attention and MLP outputs before their residual adds. 

7""" 

8 

9from transformer_lens.model_bridge.generalized_components import RMSNormalizationBridge 

10from transformer_lens.model_bridge.supported_architectures.glm import ( 

11 GlmArchitectureAdapter, 

12) 

13 

14 

15class Glm4ArchitectureAdapter(GlmArchitectureAdapter): 

16 """Architecture adapter for Glm4ForCausalLM models.""" 

17 

18 def _block_extra_norms(self): 

19 """Sandwich norms applied before the residual adds.""" 

20 return { 

21 "ln1_post": RMSNormalizationBridge(name="post_self_attn_layernorm", config=self.cfg), 

22 "ln2_post": RMSNormalizationBridge(name="post_mlp_layernorm", config=self.cfg), 

23 }