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

1"""Gemma 4 architecture adapter. 

2 

3Bridges the text path of ``Gemma4ForConditionalGeneration`` 

4(``model.language_model`` + ``lm_head``) and the vision pipeline. For the standard 

5variants (E2B / E4B / 31B / 26B-A4B) the vision encoder (``model.vision_tower``) and 

6projector (``model.embed_vision``) are both bridged, enabling Phase 7 multimodal testing. 

7 

8The same adapter also covers ``Gemma4UnifiedForConditionalGeneration`` (the 

9encoder-free 12B variant, transformers >= 5.10): its text decoder is a strict 

10structural subset — same module paths, no PLE and no MoE, both optional here. 

11It is still multimodal but has no ``vision_tower`` — ``model.embed_vision`` is the 

12full vision pipeline (raw-patch projection), mapped as the projector only. 

13 

14Per-layer structure is heterogeneous across the family, so all math is deferred to HF 

15and submodules are decomposed only for hooks (parity-safe delegation): 

16 

17- **KV sharing** (E2B/E4B): the last ``num_kv_shared_layers`` layers reuse earlier KV 

18 states and drop their own ``k_proj`` / ``v_proj`` / ``k_norm`` / ``v_norm``. 

19- **K==V attention** (31B / 26B-A4B): global-attention layers share key and value 

20 weights (``attention_k_eq_v``) and have no ``v_proj``. 

21- **Per-Layer Embeddings** (E2B/E4B): each layer mixes in a per-layer input via 

22 ``per_layer_input_gate`` / ``per_layer_projection`` / ``post_per_layer_input_norm``. 

23- **MoE** (26B-A4B): layers add a ``router`` + batched ``experts`` block in parallel 

24 with the dense MLP, sandwiched by three extra norms. 

25 

26Unlike Gemma 1-3, ``Gemma4RMSNorm`` multiplies by ``weight`` directly — there is no 

27``(1.0 + weight)`` offset. 

28""" 

29 

30from typing import Any 

31 

32from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

33from transformer_lens.model_bridge.generalized_components import ( 

34 AttentionBridge, 

35 DelegatedAttentionBlockBridge, 

36 EmbeddingBridge, 

37 LinearBridge, 

38 RotaryEmbeddingBridge, 

39 UnembeddingBridge, 

40) 

41from transformer_lens.model_bridge.generalized_components.base import ( 

42 GeneralizedComponent, 

43) 

44 

45 

46class Gemma4ArchitectureAdapter(ArchitectureAdapter): 

47 """Adapter for Gemma 4 (`Gemma4ForConditionalGeneration` — multimodal, or 

48 `Gemma4UnifiedForConditionalGeneration` — text-only 12B).""" 

49 

50 _testing_lm_attr = "model.language_model" 

51 _testing_wire_rotary = False 

52 

53 # Phase 3 (processed/compatibility mode) folds LN into a single residual stream, 

54 # which the PLE residual mix, per-layer `layer_scalar` buffers, and the MoE branch 

55 # can't represent. Phases 1 (HF parity), 2 (hooks), and 4 (text quality) apply. 

56 applicable_phases: list[int] = [1, 2, 4] 

57 

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

59 super().__init__(cfg) 

60 

61 # Both variants are multimodal (take pixel_values). The difference: 

62 # - Gemma4ForConditionalGeneration: vision_tower (encoder) + embed_vision (projector) 

63 # - Gemma4UnifiedForConditionalGeneration (12B): embed_vision only — encoder-free 

64 # embedder that does raw-patch projection without an attention-based vision encoder. 

65 arch = getattr(cfg, "architecture", "") or "" 

66 self._is_unified = "Gemma4Unified" in arch 

67 self.cfg.is_multimodal = True 

68 

69 self._extract_vision_dims(cfg) 

70 if hasattr(cfg, "vision_config"): 

71 self.cfg.mm_tokens_per_image = getattr(cfg, "vision_soft_tokens_per_image", 256) 

72 

73 self.cfg.gated_mlp = True 

74 self.cfg.uses_rms_norm = True 

75 self.cfg.normalization_type = "RMS" 

76 # Gemma4RMSNorm scales by weight directly — no (1 + weight) offset, unlike Gemma 1-3. 

77 self.cfg.rmsnorm_uses_offset = False 

78 self.cfg.positional_embedding_type = "rotary" 

79 self.cfg.attn_implementation = "eager" 

80 # PLE / layer_scalar / MoE residual topology isn't fold-safe. 

81 self.supports_fold_ln = False 

82 self.weight_processing_conversions: dict = {} 

83 

84 # Vision components. Gemma4ForConditionalGeneration has a separate vision 

85 # encoder (model.vision_tower) + projector (model.embed_vision). The 12B 

86 # unified variant is encoder-free — model.embed_vision is the full vision 

87 # pipeline (raw-patch projection), so it maps as the projector with no encoder. 

88 _vision_mapping: dict[str, Any] = { 

89 "vision_projector": GeneralizedComponent(name="model.embed_vision"), 

90 } 

91 if not self._is_unified: 

92 _vision_mapping = { 

93 "vision_encoder": GeneralizedComponent(name="model.vision_tower"), 

94 **_vision_mapping, 

95 } 

96 

97 self.component_mapping = { 

98 **_vision_mapping, 

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

100 # Single rotary module serving both layer types (full / sliding) via a 

101 # per-layer-type forward kwarg, with separate rope parameters per type. 

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

103 "blocks": DelegatedAttentionBlockBridge( 

104 name="model.language_model.layers", 

105 submodules={ 

106 # Sandwich norms: ln1/ln1_post around attention, ln2/ln2_post 

107 # around the MLP (same shape as Gemma 2/3). 

108 "ln1": GeneralizedComponent(name="input_layernorm"), 

109 "ln1_post": GeneralizedComponent(name="post_attention_layernorm"), 

110 "ln2": GeneralizedComponent(name="pre_feedforward_layernorm"), 

111 "ln2_post": GeneralizedComponent(name="post_feedforward_layernorm"), 

112 # PLE residual mix — present only when hidden_size_per_layer_input > 0 

113 # (E2B/E4B; absent on 31B and 26B-A4B). 

114 "per_layer_input_gate": GeneralizedComponent( 

115 name="per_layer_input_gate", optional=True 

116 ), 

117 "per_layer_projection": GeneralizedComponent( 

118 name="per_layer_projection", optional=True 

119 ), 

120 "post_per_layer_input_norm": GeneralizedComponent( 

121 name="post_per_layer_input_norm", optional=True 

122 ), 

123 # MoE branch — present only when enable_moe_block (26B-A4B). 

124 "router": GeneralizedComponent(name="router", optional=True), 

125 "experts": GeneralizedComponent(name="experts", optional=True), 

126 "pre_feedforward_layernorm_2": GeneralizedComponent( 

127 name="pre_feedforward_layernorm_2", optional=True 

128 ), 

129 "post_feedforward_layernorm_1": GeneralizedComponent( 

130 name="post_feedforward_layernorm_1", optional=True 

131 ), 

132 "post_feedforward_layernorm_2": GeneralizedComponent( 

133 name="post_feedforward_layernorm_2", optional=True 

134 ), 

135 # AttentionBridge for the per-head hook surface a bare 

136 # component lacks. Delegated semantics: pattern AND scores 

137 # both fire HF's post-softmax weights (no pre-softmax tensor 

138 # exists, unlike gemma1/2/3); writes to them are inert. 

139 "attn": AttentionBridge( 

140 name="self_attn", 

141 config=self.cfg, 

142 maintain_native_attention=True, 

143 submodules={ 

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

145 # KV-shared layers (E2B/E4B) drop k/v projections and norms; 

146 # K==V layers (31B / 26B-A4B global attention) drop v_proj. 

147 "k": LinearBridge(name="k_proj", optional=True), 

148 "v": LinearBridge(name="v_proj", optional=True), 

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

150 "q_norm": GeneralizedComponent(name="q_norm"), 

151 "k_norm": GeneralizedComponent(name="k_norm", optional=True), 

152 "v_norm": GeneralizedComponent(name="v_norm", optional=True), 

153 }, 

154 ), 

155 "mlp": self._gated_mlp(), 

156 }, 

157 ), 

158 "ln_final": GeneralizedComponent(name="model.language_model.norm"), 

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

160 }