Skip to content

add t5gemma2 adapter and tests#1495

Open
zinichakraborty wants to merge 3 commits into
TransformerLensOrg:devfrom
zinichakraborty:feat/t5gemma2-adapter
Open

add t5gemma2 adapter and tests#1495
zinichakraborty wants to merge 3 commits into
TransformerLensOrg:devfrom
zinichakraborty:feat/t5gemma2-adapter

Conversation

@zinichakraborty

@zinichakraborty zinichakraborty commented Jul 9, 2026

Copy link
Copy Markdown

Description

Adds a TransformerBridge architecture adapter for T5Gemma2 (T5Gemma2ForConditionalGeneration), Google's multimodal encoder-decoder follow-up to T5Gemma. The adapter bridges the text path only (the SigLIP vision tower and multi-modal projector are intentionally left unmapped).

Key architectural differences from the existing T5Gemma adapter, each reflecting a real change in the HF model:

  • Merged decoder attention: T5Gemma2 fuses self- and cross-attention into a single T5Gemma2MergedAttention with shared q/k/v/o projections and one softmax over concatenated decoder+encoder keys. A new T5Gemma2MergedAttentionBridge delegates the math to the native module and exposes both pattern slices — hook_pattern (self) and a new hook_cross_pattern (cross). The decoder block has no separate cross_attn module, no cross-attention layernorms, and a single hook_resid_mid.
  • Encoder text stack lives under model.encoder.text_model.*.
  • Gemma-style per-head QK-norm (q_norm/k_norm) in every attention module, mapped and folded with the +1 RMSNorm offset.
  • Native attention delegation (rather than the manual path) because per-layer sliding/full windows, dual RoPE, and the merged-softmax logic cannot be reproduced manually.
  • Unembed maps to lm_head.out_proj (T5Gemma2LMHead wrapper).

Registered in all four sites (factory, sources, registry, seq2seq classification) and added google/t5gemma-2-270m-270m to the model registry.

Verification (verify_models on google/t5gemma-2-270m-270m):
Phase 1 = 100% (all 475 components), Phase 2 = 100% (logit + loss parity),
Phase 4 = 87 — status VERIFIED. Bridge-vs-HF logit max diff is exactly 0.0.

Tests added:

  • tests/unit/model_bridge/supported_architectures/test_t5gemma2_adapter.py: structural adapter tests (config flags, component mapping, key translation, conversion-table alignment, registration); no weight loads.
  • tests/integration/model_bridge/test_t5gemma2_adapter.py: 17 tests on a tiny random T5Gemma2 built from config (no Hub download, CI-safe): boot, exact HF logit parity, HF delegation identity, hook shapes, merged-softmax semantics (self + cross pattern rows jointly sum to 1), per-head QK-norm.

Local runs: make test-pr tiers all pass (unit, docstring, acceptance, integration), uv run mypy . clean, make format applied.

Fixes #1454

Type of change

  • New feature (non-breaking change which adds functionality)

Screenshots

N/A

Checklist:

  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes
  • I have not rewritten tests relating to key interfaces which would affect backward compatibility

@zinichakraborty zinichakraborty changed the base branch from main to dev July 9, 2026 01:25
# Conflicts:
#	transformer_lens/tools/model_registry/data/supported_models.json
#	transformer_lens/tools/model_registry/data/verification_history.json
# Conflicts:
#	transformer_lens/tools/model_registry/data/verification_history.json
@zinichakraborty

Copy link
Copy Markdown
Author

@jlarson4 ready for review whenever! I followed the T5Gemma adapter implementation from #1433. One design note: I added a dedicated t5gemma2_merged_attention component since T5Gemma2 fuses self- and cross-attention into a single module, but happy to restructure if you'd prefer extending an existing attention bridge.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support T5Gemma2 (Gemma 3) Architecture

1 participant