Coverage for transformer_lens/model_bridge/supported_architectures/t5gemma.py: 47%

1"""T5Gemma architecture adapter. 

2 

3T5GemmaForConditionalGeneration is an encoder-decoder model combining: 

4- Gemma-style RoPE, GQA, gated MLP, and RMSNorm with offset (+1.0) 

5- Encoder-decoder cross-attention in the decoder stack 

6- Nested config: encoder/decoder dims live in cfg.encoder / cfg.decoder 

7 

8Key differences from plain T5: 

9- Uses model.encoder.layers / model.decoder.layers (not .block) 

10- No relative position bias; uses RoPE instead 

11- All norms are Gemma-style (weight + 1.0) 

12- lm_head is T5GemmaLMHead wrapping out_proj (no .weight at the top level) 

13""" 

14 

15from typing import Any 

16 

17from transformer_lens.conversion_utils.conversion_steps import ( 

18 RearrangeTensorConversion, 

19 TransposeTensorConversion, 

20) 

21from transformer_lens.conversion_utils.param_processing_conversion import ( 

22 ParamProcessingConversion, 

23) 

24from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

25from transformer_lens.model_bridge.generalized_components import ( 

26 AttentionBridge, 

27 BlockBridge, 

28 EmbeddingBridge, 

29 LinearBridge, 

30 PositionEmbeddingsAttentionBridge, 

31 RMSNormalizationBridge, 

32 RotaryEmbeddingBridge, 

33 UnembeddingBridge, 

34) 

35from transformer_lens.model_bridge.generalized_components.t5gemma_decoder_block import ( 

36 T5GemmaDecoderBlockBridge, 

37) 

38from transformer_lens.utilities.attn_implementation import force_eager_attention 

39 

40 

41class T5GemmaArchitectureAdapter(ArchitectureAdapter): 

42 """Architecture adapter for T5GemmaForConditionalGeneration. 

43 

44 Encoder: BlockBridge over model.encoder.layers (Gemma-style, no cross-attn) 

45 Decoder: T5GemmaDecoderBlockBridge over model.decoder.layers (adds cross-attn hooks) 

46 """ 

47 

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

49 super().__init__(cfg) 

50 

51 self.supports_fold_ln = False 

52 

53 # Config flags used by bridge weight processing 

54 self._set_rms_rotary_defaults() 

55 # Gemma-family GELU; the nested enc/dec config defeats the auto-mapper, 

56 # which would otherwise leave act_fn at the "relu" default. 

57 self.cfg.act_fn = "gelu_pytorch_tanh" 

58 # T5Gemma uses Gemma-style (1.0 + weight) RMSNorm offset 

59 self.cfg.rmsnorm_uses_offset = True 

60 

61 # n_heads/n_kv are decoder-effective; unbalanced pairs (t5gemma-9b-2b) 

62 # set different encoder counts, surfaced by the builder. 

63 n_heads = self.cfg.n_heads 

64 n_kv = getattr(self.cfg, "n_key_value_heads", None) or n_heads 

65 enc_heads = getattr(self.cfg, "encoder_attention_heads", None) or n_heads 

66 enc_kv = getattr(self.cfg, "encoder_key_value_heads", None) or n_kv 

67 

68 self.weight_processing_conversions = { 

69 # Encoder self-attention 

70 "encoder_blocks.{i}.self_attn.q_proj.weight": ParamProcessingConversion( 

71 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=enc_heads), 

72 ), 

73 "encoder_blocks.{i}.self_attn.k_proj.weight": ParamProcessingConversion( 

74 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=enc_kv), 

75 ), 

76 "encoder_blocks.{i}.self_attn.v_proj.weight": ParamProcessingConversion( 

77 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=enc_kv), 

78 ), 

79 "encoder_blocks.{i}.self_attn.o_proj.weight": ParamProcessingConversion( 

80 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=enc_heads), 

81 ), 

82 # Norm weights carry no conversion entry: cfg.rmsnorm_uses_offset is the 

83 # only offset signal — weight processing and the runtime norm read it 

84 # directly and apply (1 + w) themselves. 

85 # Encoder MLP (gated) 

86 "encoder_blocks.{i}.mlp.gate_proj.weight": ParamProcessingConversion( 

87 tensor_conversion=TransposeTensorConversion(), 

88 ), 

89 "encoder_blocks.{i}.mlp.up_proj.weight": ParamProcessingConversion( 

90 tensor_conversion=TransposeTensorConversion(), 

91 ), 

92 "encoder_blocks.{i}.mlp.down_proj.weight": ParamProcessingConversion( 

93 tensor_conversion=TransposeTensorConversion(), 

94 ), 

95 # Decoder self-attention 

96 "decoder_blocks.{i}.self_attn.q_proj.weight": ParamProcessingConversion( 

97 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_heads), 

98 ), 

99 "decoder_blocks.{i}.self_attn.k_proj.weight": ParamProcessingConversion( 

100 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv), 

101 ), 

102 "decoder_blocks.{i}.self_attn.v_proj.weight": ParamProcessingConversion( 

103 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv), 

104 ), 

105 "decoder_blocks.{i}.self_attn.o_proj.weight": ParamProcessingConversion( 

106 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=n_heads), 

107 ), 

108 # Decoder cross-attention 

109 "decoder_blocks.{i}.cross_attn.q_proj.weight": ParamProcessingConversion( 

110 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_heads), 

111 ), 

112 "decoder_blocks.{i}.cross_attn.k_proj.weight": ParamProcessingConversion( 

113 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv), 

114 ), 

115 "decoder_blocks.{i}.cross_attn.v_proj.weight": ParamProcessingConversion( 

116 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv), 

117 ), 

118 "decoder_blocks.{i}.cross_attn.o_proj.weight": ParamProcessingConversion( 

119 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=n_heads), 

120 ), 

121 # Decoder MLP (gated) 

122 "decoder_blocks.{i}.mlp.gate_proj.weight": ParamProcessingConversion( 

123 tensor_conversion=TransposeTensorConversion(), 

124 ), 

125 "decoder_blocks.{i}.mlp.up_proj.weight": ParamProcessingConversion( 

126 tensor_conversion=TransposeTensorConversion(), 

127 ), 

128 "decoder_blocks.{i}.mlp.down_proj.weight": ParamProcessingConversion( 

129 tensor_conversion=TransposeTensorConversion(), 

130 ), 

131 # Unembed 

132 "unembed.weight": ParamProcessingConversion( 

133 tensor_conversion=TransposeTensorConversion(), 

134 ), 

135 } 

136 

137 self.component_mapping = { 

138 # Encoder embedding and positional 

139 "encoder_embed": EmbeddingBridge(name="model.encoder.embed_tokens"), 

140 "encoder_rotary_emb": RotaryEmbeddingBridge(name="model.encoder.rotary_emb"), 

141 # Encoder layers - Gemma-style BlockBridge (pre/post norms, RoPE attention, gated MLP) 

142 "encoder_blocks": BlockBridge( 

143 name="model.encoder.layers", 

144 config=self.cfg, 

145 submodules={ 

146 "ln1": RMSNormalizationBridge(name="pre_self_attn_layernorm", config=self.cfg), 

147 "ln1_post": RMSNormalizationBridge( 

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

149 ), 

150 "attn": PositionEmbeddingsAttentionBridge( 

151 name="self_attn", 

152 config=self.cfg, 

153 submodules={ 

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

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

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

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

158 }, 

159 requires_attention_mask=True, 

160 requires_position_embeddings=True, 

161 is_causal=False, # T5Gemma encoder is bidirectional 

162 ), 

163 "ln2": RMSNormalizationBridge( 

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

165 ), 

166 "ln2_post": RMSNormalizationBridge( 

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

168 ), 

169 "mlp": self._gated_mlp(), 

170 }, 

171 ), 

172 # Encoder final norm 

173 "encoder_ln_final": RMSNormalizationBridge(name="model.encoder.norm", config=self.cfg), 

174 # Decoder embedding and positional 

175 "decoder_embed": EmbeddingBridge(name="model.decoder.embed_tokens"), 

176 "decoder_rotary_emb": RotaryEmbeddingBridge(name="model.decoder.rotary_emb"), 

177 # Decoder layers — T5GemmaDecoderBlockBridge (adds cross-attn + two mid hooks) 

178 "decoder_blocks": T5GemmaDecoderBlockBridge( 

179 name="model.decoder.layers", 

180 config=self.cfg, 

181 submodules={ 

182 # Self-attention norms 

183 "ln1": RMSNormalizationBridge(name="pre_self_attn_layernorm", config=self.cfg), 

184 "ln1_post": RMSNormalizationBridge( 

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

186 ), 

187 "self_attn": PositionEmbeddingsAttentionBridge( 

188 name="self_attn", 

189 config=self.cfg, 

190 submodules={ 

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

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

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

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

195 }, 

196 requires_attention_mask=True, 

197 requires_position_embeddings=True, 

198 ), 

199 # Cross-attention norms 

200 "ln2": RMSNormalizationBridge(name="pre_cross_attn_layernorm", config=self.cfg), 

201 "ln2_post": RMSNormalizationBridge( 

202 name="post_cross_attn_layernorm", config=self.cfg 

203 ), 

204 "cross_attn": AttentionBridge( 

205 name="cross_attn", 

206 config=self.cfg, 

207 submodules={ 

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

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

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

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

212 }, 

213 is_cross_attention=True, 

214 ), 

215 # MLP norms 

216 "ln3": RMSNormalizationBridge( 

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

218 ), 

219 "ln3_post": RMSNormalizationBridge( 

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

221 ), 

222 "mlp": self._gated_mlp(), 

223 }, 

224 ), 

225 # Decoder final norm 

226 "decoder_ln_final": RMSNormalizationBridge(name="model.decoder.norm", config=self.cfg), 

227 # lm_head is T5GemmaLMHead; the weight lives on its inner out_proj Linear 

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

229 } 

230 

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

232 """Set up rotary embedding references for T5Gemma component testing. 

233 

234 Both the encoder and decoder carry their own rotary_emb. We set the 

235 reference on all PositionEmbeddingsAttentionBridge instances so that 

236 component-level forward calls can compute RoPE correctly. 

237 """ 

238 encoder_rotary = hf_model.model.encoder.rotary_emb 

239 decoder_rotary = hf_model.model.decoder.rotary_emb 

240 

241 force_eager_attention(hf_model) 

242 

243 if bridge_model is not None: 

244 for block in getattr(bridge_model, "encoder_blocks", []): 

245 if hasattr(block, "attn"): 

246 block.attn.set_rotary_emb(encoder_rotary) 

247 for block in getattr(bridge_model, "decoder_blocks", []): 

248 if hasattr(block, "self_attn"): 

249 block.self_attn.set_rotary_emb(decoder_rotary) 

250 

251 enc_attn = self.get_generalized_component("encoder_blocks.0.attn") 

252 enc_attn.set_rotary_emb(encoder_rotary) 

253 dec_self_attn = self.get_generalized_component("decoder_blocks.0.self_attn") 

254 dec_self_attn.set_rotary_emb(decoder_rotary)