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

1"""GLM architecture adapter. 

2 

3Z.ai's dense GLM-4 family (``GlmForCausalLM``: glm-4-9b-*-hf, glm-edge): 

4llama-shaped GQA decoder with GLM's adjacent-pair interleaved RoPE at a 

5partial rotary factor, attention biases, and a Phi3-style combined 

6``gate_up_proj`` MLP. 

7""" 

8 

9from typing import Any 

10 

11from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

12from transformer_lens.model_bridge.generalized_components import ( 

13 BlockBridge, 

14 EmbeddingBridge, 

15 JointGateUpMLPBridge, 

16 LinearBridge, 

17 PositionEmbeddingsAttentionBridge, 

18 RMSNormalizationBridge, 

19 RotaryEmbeddingBridge, 

20 UnembeddingBridge, 

21) 

22from transformer_lens.model_bridge.supported_architectures.phi3 import ( 

23 Phi3ArchitectureAdapter, 

24) 

25 

26 

27class GlmArchitectureAdapter(ArchitectureAdapter): 

28 """Architecture adapter for GlmForCausalLM models.""" 

29 

30 _testing_eager = "config" 

31 

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

33 """Initialize the GLM architecture adapter.""" 

34 super().__init__(cfg) 

35 

36 self._set_rms_rotary_defaults() 

37 # GLM rotates adjacent element pairs (interleaved RoPE), like ERNIE. 

38 self.cfg.rotary_adjacent_pairs = True 

39 # GLM tokenizers carry no BOS token. 

40 self.cfg.default_prepend_bos = False 

41 

42 # Joint gate_up_proj cannot be folded by the standard LN machinery. 

43 self.supports_fold_ln = False 

44 # GLM has attention biases on q/k/v; the bias reshapes must use the 

45 # kv-head count or compat mode's value-bias fold crashes on GQA. 

46 self.weight_processing_conversions = { 

47 **self._qkvo_weight_conversions(include_biases=True), 

48 } 

49 

50 self.component_mapping = { 

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

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

53 "blocks": BlockBridge( 

54 name="model.layers", 

55 submodules={ 

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

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

58 **self._block_extra_norms(), 

59 "attn": PositionEmbeddingsAttentionBridge( 

60 name="self_attn", 

61 config=self.cfg, 

62 submodules={ 

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

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

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

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

67 }, 

68 requires_attention_mask=True, 

69 requires_position_embeddings=True, 

70 ), 

71 # GlmMLP chunks gate_up_proj output in half — same layout 

72 # Phi3 uses, so its splitter applies unchanged. 

73 "mlp": JointGateUpMLPBridge( 

74 name="mlp", 

75 config=self.cfg, 

76 split_gate_up_matrix=Phi3ArchitectureAdapter._split_gate_up, 

77 submodules={ 

78 "out": LinearBridge(name="down_proj"), 

79 }, 

80 ), 

81 }, 

82 ), 

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

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

85 } 

86 

87 def _block_extra_norms(self): 

88 """Sandwich-norms seam; GLM-4-0414 adds post_self_attn/post_mlp norms.""" 

89 return {}