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

1"""GLM-4V / GLM-4.1V architecture adapter. 

2 

3Z.ai's GLM-4V line (``Glm4vForConditionalGeneration``, GLM-4.1V-9B 

4Thinking): a GLM vision tower at ``model.visual`` (RMS-normed blocks, 

5learned position embeddings + 2D rotary, patch merger + conv 

6downsample) feeding a GLM-4-0414-layout text decoder at 

7``model.language_model`` — sandwich norms and the joint ``gate_up_proj`` 

8MLP. Text attention uses mRoPE (three position streams), so it stays 

9HF-native; the tower is delegated opaquely with the merger as the 

10projector. 

11""" 

12 

13from typing import Any 

14 

15from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

16from transformer_lens.model_bridge.generalized_components import ( 

17 AttentionBridge, 

18 BlockBridge, 

19 EmbeddingBridge, 

20 JointGateUpMLPBridge, 

21 LinearBridge, 

22 RMSNormalizationBridge, 

23 UnembeddingBridge, 

24 VisionProjectionBridge, 

25) 

26from transformer_lens.model_bridge.generalized_components.base import ( 

27 GeneralizedComponent, 

28) 

29from transformer_lens.model_bridge.supported_architectures.phi3 import ( 

30 Phi3ArchitectureAdapter, 

31) 

32 

33 

34class Glm4vArchitectureAdapter(ArchitectureAdapter): 

35 """Architecture adapter for Glm4vForConditionalGeneration models.""" 

36 

37 required_libraries: list[str] = ["torchvision"] 

38 required_libraries_group: str = "multimodal" 

39 

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

41 """Initialize the GLM-4V architecture adapter.""" 

42 super().__init__(cfg) 

43 

44 self.cfg.is_multimodal = True 

45 self._set_rms_rotary_defaults() 

46 self.cfg.attn_implementation = "eager" 

47 # GLM tokenizers carry no BOS token. 

48 self.cfg.default_prepend_bos = False 

49 

50 self._extract_vision_dims(cfg) 

51 

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

53 self.supports_fold_ln = False 

54 # GLM attention carries QKV biases. 

55 self.weight_processing_conversions = { 

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

57 } 

58 

59 self.component_mapping = { 

60 "vision_encoder": GeneralizedComponent(name="model.visual"), 

61 "vision_projector": VisionProjectionBridge(name="model.visual.merger"), 

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

63 "blocks": BlockBridge( 

64 name="model.language_model.layers", 

65 submodules={ 

66 # GLM-4-0414 sandwich layout: post_attention_layernorm is 

67 # the pre-MLP norm; the sandwich norms sit on sublayer 

68 # outputs before their residual adds. 

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

70 "ln1_post": RMSNormalizationBridge( 

71 name="post_self_attn_layernorm", config=self.cfg 

72 ), 

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

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

75 # mRoPE (3-section multimodal rotary) lives in HF's forward. 

76 "attn": AttentionBridge( 

77 name="self_attn", 

78 config=self.cfg, 

79 submodules={ 

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

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

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

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

84 }, 

85 maintain_native_attention=True, 

86 requires_attention_mask=True, 

87 ), 

88 "mlp": JointGateUpMLPBridge( 

89 name="mlp", 

90 config=self.cfg, 

91 split_gate_up_matrix=Phi3ArchitectureAdapter._split_gate_up, 

92 submodules={ 

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

94 }, 

95 ), 

96 }, 

97 ), 

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

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

100 }