Coverage for transformer_lens/model_bridge/supported_architectures/t5gemma2.py: 43%
47 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""T5Gemma2 architecture adapter (text-only).
3T5Gemma2ForConditionalGeneration is a multimodal encoder-decoder model. This
4adapter bridges the text path only:
5- Encoder text stack under model.encoder.text_model (the SigLIP vision_tower and
6 multi_modal_projector are intentionally left unmapped).
7- Decoder stack under model.decoder.
9Key differences from T5Gemma:
10- Encoder text lives at model.encoder.text_model.* (not model.encoder.*).
11- The decoder uses a single T5Gemma2MergedAttention that fuses self- and
12 cross-attention with shared q/k/v/o projections; there is no separate
13 cross-attention module and no cross-attention layernorms.
14- Both encoder and decoder attention add Gemma-style QK-norm (q_norm/k_norm).
15- Per-layer sliding/full attention with dual RoPE and per-head QK-norm are all
16 handled natively by HF — the bridge only routes inputs and fires hooks.
17"""
19from typing import Any
21from transformer_lens.conversion_utils.conversion_steps import (
22 RearrangeTensorConversion,
23 TransposeTensorConversion,
24)
25from transformer_lens.conversion_utils.param_processing_conversion import (
26 ParamProcessingConversion,
27)
28from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
29from transformer_lens.model_bridge.generalized_components import (
30 AttentionBridge,
31 BlockBridge,
32 EmbeddingBridge,
33 GatedMLPBridge,
34 LinearBridge,
35 RMSNormalizationBridge,
36 RotaryEmbeddingBridge,
37 UnembeddingBridge,
38)
39from transformer_lens.model_bridge.generalized_components.t5gemma2_decoder_block import (
40 T5Gemma2DecoderBlockBridge,
41)
42from transformer_lens.model_bridge.generalized_components.t5gemma2_merged_attention import (
43 T5Gemma2MergedAttentionBridge,
44)
45from transformer_lens.utilities.attn_implementation import force_eager_attention
48class T5Gemma2ArchitectureAdapter(ArchitectureAdapter):
49 """Architecture adapter for T5Gemma2ForConditionalGeneration (text-only).
51 Encoder: BlockBridge over model.encoder.text_model.layers (Gemma-style, QK-norm, no cross-attn)
52 Decoder: T5Gemma2DecoderBlockBridge over model.decoder.layers (merged self+cross attention)
53 """
55 def __init__(self, cfg: Any) -> None:
56 super().__init__(cfg)
58 self.supports_fold_ln = False
60 # Config flags used by bridge weight processing
61 self.cfg.normalization_type = "RMS"
62 self.cfg.positional_embedding_type = "rotary"
63 self.cfg.final_rms = True
64 self.cfg.gated_mlp = True
65 self.cfg.attn_only = False
66 # Gemma-family GELU; the nested enc/dec config defeats the auto-mapper,
67 # which would otherwise leave act_fn at the "relu" default.
68 self.cfg.act_fn = "gelu_pytorch_tanh"
69 self.cfg.uses_rms_norm = True
70 # T5Gemma2 uses Gemma-style (1.0 + weight) RMSNorm offset
71 self.cfg.rmsnorm_uses_offset = True
73 # n_heads/n_kv are decoder-effective; the builder surfaces the encoder
74 # text stack's own counts for unbalanced pairs.
75 n_heads = self.cfg.n_heads
76 n_kv = getattr(self.cfg, "n_key_value_heads", None) or n_heads
77 enc_heads = getattr(self.cfg, "encoder_attention_heads", None) or n_heads
78 enc_kv = getattr(self.cfg, "encoder_key_value_heads", None) or n_kv
80 self.weight_processing_conversions = {
81 # Encoder self-attention
82 "encoder_blocks.{i}.self_attn.q_proj.weight": ParamProcessingConversion(
83 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=enc_heads),
84 ),
85 "encoder_blocks.{i}.self_attn.k_proj.weight": ParamProcessingConversion(
86 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=enc_kv),
87 ),
88 "encoder_blocks.{i}.self_attn.v_proj.weight": ParamProcessingConversion(
89 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=enc_kv),
90 ),
91 "encoder_blocks.{i}.self_attn.o_proj.weight": ParamProcessingConversion(
92 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=enc_heads),
93 ),
94 # Norm weights (incl. q_norm/k_norm) carry no conversion entry:
95 # cfg.rmsnorm_uses_offset is the only offset signal — weight processing
96 # and the runtime norm read it directly and apply (1 + w) themselves.
97 # Encoder MLP (gated)
98 "encoder_blocks.{i}.mlp.gate_proj.weight": ParamProcessingConversion(
99 tensor_conversion=TransposeTensorConversion(),
100 ),
101 "encoder_blocks.{i}.mlp.up_proj.weight": ParamProcessingConversion(
102 tensor_conversion=TransposeTensorConversion(),
103 ),
104 "encoder_blocks.{i}.mlp.down_proj.weight": ParamProcessingConversion(
105 tensor_conversion=TransposeTensorConversion(),
106 ),
107 # Decoder merged attention (self + cross share these projections)
108 "decoder_blocks.{i}.self_attn.q_proj.weight": ParamProcessingConversion(
109 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_heads),
110 ),
111 "decoder_blocks.{i}.self_attn.k_proj.weight": ParamProcessingConversion(
112 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv),
113 ),
114 "decoder_blocks.{i}.self_attn.v_proj.weight": ParamProcessingConversion(
115 tensor_conversion=RearrangeTensorConversion("(n h) m -> n m h", n=n_kv),
116 ),
117 "decoder_blocks.{i}.self_attn.o_proj.weight": ParamProcessingConversion(
118 tensor_conversion=RearrangeTensorConversion("m (n h) -> n h m", n=n_heads),
119 ),
120 # Decoder MLP (gated)
121 "decoder_blocks.{i}.mlp.gate_proj.weight": ParamProcessingConversion(
122 tensor_conversion=TransposeTensorConversion(),
123 ),
124 "decoder_blocks.{i}.mlp.up_proj.weight": ParamProcessingConversion(
125 tensor_conversion=TransposeTensorConversion(),
126 ),
127 "decoder_blocks.{i}.mlp.down_proj.weight": ParamProcessingConversion(
128 tensor_conversion=TransposeTensorConversion(),
129 ),
130 # Unembed
131 "unembed.weight": ParamProcessingConversion(
132 tensor_conversion=TransposeTensorConversion(),
133 ),
134 }
136 self.component_mapping = {
137 # Encoder embedding and positional (text stack lives under text_model)
138 "encoder_embed": EmbeddingBridge(name="model.encoder.text_model.embed_tokens"),
139 "encoder_rotary_emb": RotaryEmbeddingBridge(name="model.encoder.text_model.rotary_emb"),
140 # Encoder layers - Gemma-style BlockBridge (pre/post norms, QK-norm attention, gated MLP)
141 "encoder_blocks": BlockBridge(
142 name="model.encoder.text_model.layers",
143 config=self.cfg,
144 submodules={
145 "ln1": RMSNormalizationBridge(name="pre_self_attn_layernorm", config=self.cfg),
146 "ln1_post": RMSNormalizationBridge(
147 name="post_self_attn_layernorm", config=self.cfg
148 ),
149 # Native delegation: the encoder uses per-layer sliding/full
150 # bidirectional windows carried by HF's per-layer mask, which the
151 # manual attention path does not apply (it drifts materially past
152 # the sliding_window length). Delegating keeps sliding correct.
153 "attn": AttentionBridge(
154 name="self_attn",
155 config=self.cfg,
156 # HF's T5Gemma2SelfAttention unpacks position_embeddings
157 # unconditionally, so component testing must supply it.
158 requires_position_embeddings=True,
159 submodules={
160 "q": LinearBridge(name="q_proj"),
161 "k": LinearBridge(name="k_proj"),
162 "v": LinearBridge(name="v_proj"),
163 "o": LinearBridge(name="o_proj"),
164 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg),
165 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg),
166 },
167 ),
168 "ln2": RMSNormalizationBridge(
169 name="pre_feedforward_layernorm", config=self.cfg
170 ),
171 "ln2_post": RMSNormalizationBridge(
172 name="post_feedforward_layernorm", config=self.cfg
173 ),
174 "mlp": GatedMLPBridge(
175 name="mlp",
176 config=self.cfg,
177 submodules={
178 "gate": LinearBridge(name="gate_proj"),
179 "in": LinearBridge(name="up_proj"),
180 "out": LinearBridge(name="down_proj"),
181 },
182 ),
183 },
184 ),
185 # Encoder final norm
186 "encoder_ln_final": RMSNormalizationBridge(
187 name="model.encoder.text_model.norm", config=self.cfg
188 ),
189 # Decoder embedding and positional
190 "decoder_embed": EmbeddingBridge(name="model.decoder.embed_tokens"),
191 "decoder_rotary_emb": RotaryEmbeddingBridge(name="model.decoder.rotary_emb"),
192 # Decoder layers — T5Gemma2DecoderBlockBridge (merged self+cross attention)
193 "decoder_blocks": T5Gemma2DecoderBlockBridge(
194 name="model.decoder.layers",
195 config=self.cfg,
196 submodules={
197 "ln1": RMSNormalizationBridge(name="pre_self_attn_layernorm", config=self.cfg),
198 "ln1_post": RMSNormalizationBridge(
199 name="post_self_attn_layernorm", config=self.cfg
200 ),
201 # Delegates to the native T5Gemma2MergedAttention (self+cross with
202 # shared q/k/v/o); the merged/cross logic, QK-norm, RoPE, and scaling
203 # cannot be reimplemented by the manual attention path. Exposes the
204 # self pattern (hook_pattern) and cross pattern (hook_cross_pattern).
205 "self_attn": T5Gemma2MergedAttentionBridge(
206 name="self_attn",
207 config=self.cfg,
208 submodules={
209 "q": LinearBridge(name="q_proj"),
210 "k": LinearBridge(name="k_proj"),
211 "v": LinearBridge(name="v_proj"),
212 "o": LinearBridge(name="o_proj"),
213 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg),
214 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg),
215 },
216 ),
217 "ln2": RMSNormalizationBridge(
218 name="pre_feedforward_layernorm", config=self.cfg
219 ),
220 "ln2_post": RMSNormalizationBridge(
221 name="post_feedforward_layernorm", config=self.cfg
222 ),
223 "mlp": GatedMLPBridge(
224 name="mlp",
225 config=self.cfg,
226 submodules={
227 "gate": LinearBridge(name="gate_proj"),
228 "in": LinearBridge(name="up_proj"),
229 "out": LinearBridge(name="down_proj"),
230 },
231 ),
232 },
233 ),
234 # Decoder final norm
235 "decoder_ln_final": RMSNormalizationBridge(name="model.decoder.norm", config=self.cfg),
236 # lm_head is T5Gemma2LMHead; the weight lives on its inner out_proj Linear
237 "unembed": UnembeddingBridge(name="lm_head.out_proj", config=self.cfg),
238 }
240 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None:
241 """Set up rotary embedding references for T5Gemma2 component testing.
243 Both the encoder text stack and the decoder carry their own rotary_emb. We
244 set the reference on all PositionEmbeddingsAttentionBridge instances so that
245 component-level forward calls can compute RoPE correctly, force eager
246 attention (so patterns are hookable), and enable native layernorm autograd
247 on QK-norm so the manual encoder path matches HF exactly.
248 """
249 encoder_rotary = hf_model.model.encoder.text_model.rotary_emb
251 force_eager_attention(hf_model)
253 # QK-norm must delegate to HF's exact RMSNorm autograd to avoid manual drift.
254 def _enable_qk_native_autograd(layers: Any) -> None:
255 for layer in layers:
256 attn = getattr(layer, "self_attn", None)
257 if attn is None:
258 continue
259 if hasattr(attn, "q_norm"):
260 attn.q_norm.use_native_layernorm_autograd = True
261 if hasattr(attn, "k_norm"):
262 attn.k_norm.use_native_layernorm_autograd = True
264 _enable_qk_native_autograd(hf_model.model.encoder.text_model.layers)
265 _enable_qk_native_autograd(hf_model.model.decoder.layers)
267 if bridge_model is not None:
268 for block in getattr(bridge_model, "encoder_blocks", []):
269 if hasattr(block, "attn") and hasattr(block.attn, "set_rotary_emb"):
270 block.attn.set_rotary_emb(encoder_rotary)
271 # Decoder self_attn delegates to native (which owns its RoPE), so it
272 # has no set_rotary_emb; nothing to wire.
274 enc_attn = self.get_generalized_component("encoder_blocks.0.attn")
275 if hasattr(enc_attn, "set_rotary_emb"):
276 enc_attn.set_rotary_emb(encoder_rotary)