Coverage for transformer_lens/model_bridge/supported_architectures/gemma3_multimodal.py: 92%

1"""Gemma3 Multimodal architecture adapter. 

2 

3This adapter supports Gemma3ForConditionalGeneration, the vision-language 

4variant of Gemma 3 used by models like MedGemma. 

5""" 

6 

7from typing import Any 

8 

9from transformer_lens.conversion_utils.conversion_steps import TransposeTensorConversion 

10from transformer_lens.conversion_utils.param_processing_conversion import ( 

11 ParamProcessingConversion, 

12) 

13from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

14from transformer_lens.model_bridge.generalized_components import ( 

15 BlockBridge, 

16 EmbeddingBridge, 

17 LinearBridge, 

18 RMSNormalizationBridge, 

19 RotaryEmbeddingBridge, 

20 SiglipVisionEncoderBridge, 

21 UnembeddingBridge, 

22 VisionProjectionBridge, 

23) 

24from transformer_lens.model_bridge.generalized_components.position_embeddings_attention import ( 

25 PositionEmbeddingsAttentionBridge, 

26) 

27from transformer_lens.model_bridge.supported_architectures.gemma3 import ( 

28 _enable_native_qk_norm_autograd, 

29) 

30 

31 

32class Gemma3MultimodalArchitectureAdapter(ArchitectureAdapter): 

33 """Architecture adapter for Gemma3 multimodal models (Gemma3ForConditionalGeneration). 

34 

35 This adapter handles vision-language models like Gemma 3 4B/12B/27B and MedGemma. 

36 The model structure is: 

37 - model.vision_tower: SigLIP vision encoder 

38 - model.multi_modal_projector: Projects vision embeddings to language space 

39 - model.language_model: Gemma3TextModel (same as text-only Gemma 3) 

40 - lm_head: Output projection 

41 

42 The language model component follows the same patterns as Gemma3ArchitectureAdapter. 

43 """ 

44 

45 _testing_lm_attr = "model.language_model" 

46 

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

48 """Initialize the Gemma3 multimodal architecture adapter.""" 

49 super().__init__(cfg) 

50 

51 self.cfg.is_multimodal = True 

52 

53 # Language model configuration (same as text-only Gemma 3) 

54 self.cfg.gated_mlp = True 

55 self.cfg.uses_rms_norm = True 

56 self.cfg.normalization_type = "RMS" 

57 # Gemma models use (1.0 + weight) in RMSNorm instead of just weight. 

58 # Without this, fold_ln sets identity to 1.0 instead of 0.0, causing 2x scaling. 

59 self.cfg.rmsnorm_uses_offset = True 

60 self.cfg.positional_embedding_type = "rotary" 

61 self.cfg.attn_implementation = "eager" 

62 

63 # Store vision-related config 

64 self._extract_vision_dims(cfg) 

65 

66 # Store multimodal projection config 

67 self.cfg.mm_tokens_per_image = getattr(cfg, "mm_tokens_per_image", 256) 

68 

69 # Weight processing conversions for the language model 

70 # Note: The language model weights are under "model.language_model.*" 

71 self.weight_processing_conversions = { 

72 # Q/K/V weight conversions for language model 

73 **self._qkvo_weight_conversions(), 

74 # Norm weights (incl. q_norm/k_norm) carry no conversion entry: 

75 # cfg.rmsnorm_uses_offset is the only offset signal — weight processing 

76 # and the runtime norm read it directly and apply (1 + w) themselves. 

77 # MLP weight conversions 

78 "blocks.{i}.mlp.gate.weight": ParamProcessingConversion( 

79 tensor_conversion=TransposeTensorConversion(), 

80 ), 

81 "blocks.{i}.mlp.in.weight": ParamProcessingConversion( 

82 tensor_conversion=TransposeTensorConversion(), 

83 ), 

84 "blocks.{i}.mlp.out.weight": ParamProcessingConversion( 

85 tensor_conversion=TransposeTensorConversion(), 

86 ), 

87 # Unembed weight conversion 

88 "unembed.weight": ParamProcessingConversion( 

89 tensor_conversion=TransposeTensorConversion(), 

90 ), 

91 } 

92 

93 # Component mapping for the full multimodal model 

94 # Note: We use distinct TL names (vision_encoder, vision_projector) to avoid 

95 # conflicting with HF model attribute names (vision_tower, multi_modal_projector) 

96 self.component_mapping = { 

97 # Vision components 

98 "vision_encoder": SiglipVisionEncoderBridge(name="model.vision_tower", config=self.cfg), 

99 "vision_projector": VisionProjectionBridge(name="model.multi_modal_projector"), 

100 # Language model components (under model.language_model) 

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

102 "rotary_emb": RotaryEmbeddingBridge(name="model.language_model.rotary_emb"), 

103 "blocks": BlockBridge( 

104 name="model.language_model.layers", 

105 submodules={ 

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

107 "ln1_post": RMSNormalizationBridge( 

108 name="post_attention_layernorm", config=self.cfg 

109 ), 

110 "ln2": RMSNormalizationBridge( 

111 name="pre_feedforward_layernorm", config=self.cfg 

112 ), 

113 "ln2_post": RMSNormalizationBridge( 

114 name="post_feedforward_layernorm", config=self.cfg 

115 ), 

116 "attn": PositionEmbeddingsAttentionBridge( 

117 name="self_attn", 

118 config=self.cfg, 

119 submodules={ 

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

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

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

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

124 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

125 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

126 }, 

127 ), 

128 "mlp": self._gated_mlp(), 

129 }, 

130 ), 

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

132 "unembed": UnembeddingBridge(name="lm_head"), 

133 } 

134 

135 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: 

136 """Wire rotary + eager, then enable native autograd on the Q/K norms.""" 

137 super().setup_component_testing(hf_model, bridge_model) 

138 _enable_native_qk_norm_autograd(bridge_model)