From 4c004ae137773d373adb9e88b714f6268599c8d9 Mon Sep 17 00:00:00 2001 From: Jaonary Rabarisoa Date: Mon, 26 Jan 2026 08:56:41 +0100 Subject: [PATCH 1/5] init t5gemma2 model --- bonsai/models/t5gemma2/README.md | 1 + bonsai/models/t5gemma2/__init__.py | 0 bonsai/models/t5gemma2/tests/__init__.py | 0 3 files changed, 1 insertion(+) create mode 100644 bonsai/models/t5gemma2/README.md create mode 100644 bonsai/models/t5gemma2/__init__.py create mode 100644 bonsai/models/t5gemma2/tests/__init__.py diff --git a/bonsai/models/t5gemma2/README.md b/bonsai/models/t5gemma2/README.md new file mode 100644 index 00000000..2c5d8e54 --- /dev/null +++ b/bonsai/models/t5gemma2/README.md @@ -0,0 +1 @@ +# T5Gemma2 in JAX (WIP) \ No newline at end of file diff --git a/bonsai/models/t5gemma2/__init__.py b/bonsai/models/t5gemma2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bonsai/models/t5gemma2/tests/__init__.py b/bonsai/models/t5gemma2/tests/__init__.py new file mode 100644 index 00000000..e69de29b From 7a91bfc898ad0c472eed90011b7e25997b71d19c Mon Sep 17 00:00:00 2001 From: Jaonary Rabarisoa Date: Mon, 26 Jan 2026 11:51:37 +0100 Subject: [PATCH 2/5] Add T5Gemma2 model implementation --- bonsai/models/t5gemma2/modeling.py | 2375 ++++++++++++++++++++++++++++ 1 file changed, 2375 insertions(+) create mode 100644 bonsai/models/t5gemma2/modeling.py diff --git a/bonsai/models/t5gemma2/modeling.py b/bonsai/models/t5gemma2/modeling.py new file mode 100644 index 00000000..0bc278d4 --- /dev/null +++ b/bonsai/models/t5gemma2/modeling.py @@ -0,0 +1,2375 @@ +# Copyright 2025 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import functools +from dataclasses import dataclass, field +from enum import Enum + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx +from flax.nnx.module import first_from +from jaxtyping import Array, Bool, Float, Int + + +# ============================================================================= +# Configuration Classes +# ============================================================================= + + +class AttentionType(Enum): + GLOBAL = "global" + LOCAL_SLIDING = "local_sliding" + + +# Attention pattern: 5 local sliding + 1 global +_ATTN_PATTERN = ( + AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, +) + + +def _make_layer_types(num_layers: int) -> tuple[AttentionType, ...]: + """Generate attention types for all layers using the standard pattern.""" + n = len(_ATTN_PATTERN) + return (_ATTN_PATTERN * (num_layers // n) + _ATTN_PATTERN[: num_layers % n]) + + +@dataclass(frozen=True) +class RoPEParameters: + rope_type: str = "default" + rope_theta: float = 10_000.0 + factor: float = 1.0 + + +# Default RoPE parameters used across all model sizes +_DEFAULT_ROPE_PARAMS: dict[str, RoPEParameters | None] = { + "full_attention": RoPEParameters(rope_type="linear", rope_theta=1_000_000.0, factor=8.0), + "sliding_attention": RoPEParameters(rope_type="default", rope_theta=10_000.0, factor=1.0), +} + + +@dataclass(frozen=True) +class T5Gemma2VisionConfig: + width: int = 1152 + image_size: int = 896 + patch_size: int = 14 + depth: int = 27 + mlp_dim: int = 4304 + num_heads: int = 16 + posemb: str = "learn" + dropout: float = 0.0 + + +@dataclass(frozen=True) +class T5Gemma2TextConfig: + num_hidden_layers: int + embed_dim: int + intermediate_size: int + num_attention_heads: int + num_key_value_heads: int + sliding_window: int + layer_types: tuple[AttentionType, ...] = () + vocab_size: int = 262144 + head_dim: int = 256 + rms_norm_eps: float = 1e-6 + rope_parameters: dict[str, RoPEParameters | None] = field(default_factory=lambda: _DEFAULT_ROPE_PARAMS) + pad_token_id: int = 0 + bos_token_id: int = 2 + eos_token_id: int = 1 + attn_logit_softcapping: float | None = None + + @functools.cached_property + def query_pre_attn_scalar(self) -> float: + return self.head_dim**-0.5 + + +@dataclass(frozen=True) +class T5Gemma2EncoderConfig: + text_config: T5Gemma2TextConfig + vision_config: T5Gemma2VisionConfig | None = None + mm_tokens_per_image: int = 256 + image_token_id: int = 256001 + pad_token_id: int = 0 + + +@dataclass(frozen=True) +class T5Gemma2DecoderConfig(T5Gemma2TextConfig): + """Decoder config - inherits all fields from T5Gemma2TextConfig.""" + + +@dataclass(frozen=True) +class T5Gemma2Config: + encoder: T5Gemma2EncoderConfig + decoder: T5Gemma2DecoderConfig + eoi_token_index: int = 256000 + pad_token_id: int = 0 + + @classmethod + def _from_params( + cls, + num_layers: int, + embed_dim: int, + intermediate_size: int, + num_attention_heads: int, + num_key_value_heads: int, + sliding_window: int, + with_vision: bool = True, + ) -> "T5Gemma2Config": + layer_types = _make_layer_types(num_layers) + text_cfg = T5Gemma2TextConfig( + num_hidden_layers=num_layers, + embed_dim=embed_dim, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + sliding_window=sliding_window, + layer_types=layer_types, + ) + vision_cfg = T5Gemma2VisionConfig() if with_vision else None + return cls( + encoder=T5Gemma2EncoderConfig(text_config=text_cfg, vision_config=vision_cfg), + decoder=T5Gemma2DecoderConfig( + num_hidden_layers=num_layers, + embed_dim=embed_dim, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + sliding_window=sliding_window, + layer_types=layer_types, + ), + ) + + @classmethod + def t5gemma2_270m_270m(cls, with_vision: bool = True) -> "T5Gemma2Config": + return cls._from_params(18, 640, 2048, 4, 1, 512, with_vision) + + @classmethod + def t5gemma2_1b_1b(cls, with_vision: bool = True) -> "T5Gemma2Config": + return cls._from_params(26, 1152, 6912, 4, 1, 512, with_vision) + + @classmethod + def t5gemma2_4b_4b(cls, with_vision: bool = True) -> "T5Gemma2Config": + return cls._from_params(34, 2560, 10240, 8, 4, 1024, with_vision) + + +# ============================================================================= +# Constants +# ============================================================================= + +# Large negative number for masking in attention +K_MASK = -2.3819763e38 + +# Special tokens +BOS_TOKEN = 2 +EOS_TOKEN = 1 +NEW_LINE_TOKEN = 108 +START_OF_IMAGE_TOKEN = 255999 +END_OF_IMAGE_TOKEN = 256000 +IMAGE_PLACEHOLDER_IN_PROMPT = "" +IMAGE_PLACEHOLDER_TOKEN = 256001 # Placeholder for image. Different from Gemma3. +NUM_PLACEHOLDER_TOKENS_PER_IMAGE = 256 + +# Default initializers +NORMAL_INIT = nnx.initializers.normal() +ZEROS_INIT = nnx.initializers.zeros_init() + + +# ============================================================================= +# Helper Functions +# ============================================================================= + + +def _get_rope_base_frequency( + text_config: T5Gemma2TextConfig, + attn_type: AttentionType, +) -> int: + """Get RoPE base frequency for a given attention type. + + Args: + text_config: Text configuration containing rope_parameters. + attn_type: The attention type (GLOBAL or LOCAL_SLIDING). + + Returns: + The base frequency for RoPE. + """ + # Map attention type to rope parameter key + if attn_type == AttentionType.GLOBAL: + key = "full_attention" + else: + key = "sliding_attention" + + rope_params = text_config.rope_parameters.get(key) + if rope_params is not None: + return int(rope_params.rope_theta) + return 10_000 # Default + + +def _get_rope_scale_factor( + text_config: T5Gemma2TextConfig, + attn_type: AttentionType, +) -> float: + """Get RoPE scale factor for a given attention type. + + For "linear" rope_type, this returns the factor from config. + For "default" rope_type, returns 1.0 (no scaling). + + Args: + text_config: Text configuration containing rope_parameters. + attn_type: The attention type (GLOBAL or LOCAL_SLIDING). + + Returns: + The scale factor for RoPE. + """ + # Map attention type to rope parameter key + if attn_type == AttentionType.GLOBAL: + key = "full_attention" + else: + key = "sliding_attention" + + rope_params = text_config.rope_parameters.get(key) + if rope_params is not None: + # "linear" rope_type uses factor, "default" uses 1.0 + if rope_params.rope_type == "linear": + return rope_params.factor + return 1.0 # Default (no scaling) + + +def apply_rope( + inputs: Float[Array, "B L N H"], + positions: Float[Array, "B L"], + *, + base_frequency: int, + scale_factor: float = 1.0, + rope_proportion: float = 1.0, +) -> Float[Array, "B L N H"]: + """Applies Rotary Position Embeddings (RoPE). + + Args: + inputs: Array of shape [B, L, N, H]. + positions: Array of shape [B, L]. + base_frequency: Base frequency used to compute rotations. + scale_factor: Scale factor for positional interpolation. + rope_proportion: Proportion of head dimension to apply RoPE to. + + Returns: + Array of shape [B, L, N, H]. + """ + head_dim = inputs.shape[-1] + rope_angles = int(rope_proportion * head_dim // 2) + nope_angles = head_dim // 2 - rope_angles + freq_exponents = (2.0 / head_dim) * jnp.arange(0, rope_angles, dtype=jnp.float32) + timescale = jnp.pad( + base_frequency**freq_exponents, + (0, nope_angles), + mode="constant", + constant_values=(0, jnp.inf), + ) + + sinusoid_inp = positions[..., jnp.newaxis] / timescale[jnp.newaxis, jnp.newaxis, :] + sinusoid_inp = sinusoid_inp[..., jnp.newaxis, :] + if scale_factor < 1.0: + raise ValueError(f"scale_factor must be >= 1.0, got {scale_factor}") + sinusoid_inp /= scale_factor + + sin = jnp.sin(sinusoid_inp) + cos = jnp.cos(sinusoid_inp) + + first_half, second_half = jnp.split(inputs, 2, axis=-1) + first_part = first_half * cos - second_half * sin + second_part = second_half * cos + first_half * sin + out = jnp.concatenate([first_part, second_part], axis=-1) + return out.astype(inputs.dtype) + + +# ============================================================================= +# Base Layers +# ============================================================================= + + +class T5Gemma2Einsum(nnx.Module): + """Parameterized einsum layer.""" + + def __init__( + self, + shape: tuple[int, ...], + *, + kernel_init: nnx.Initializer = NORMAL_INIT, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.w = nnx.Param(kernel_init(rngs.params(), shape, dtype)) + + def __call__(self, eqn: str, x: jax.Array) -> jax.Array: + return jnp.einsum(eqn, x, self.w.value) + + +class T5Gemma2RMSNorm(nnx.Module): + """RMS Normalization layer.""" + + def __init__( + self, + features: int, + *, + scale_init: nnx.Initializer = ZEROS_INIT, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.scale = nnx.Param(scale_init(rngs.params(), (features,), dtype)) + + def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: + var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) + + # Jax.lax.rsqrt is used because it returns different floats than + # jnp.reciprocal(jnp.sqrt(var + 1e-06)) + normed_inputs = x * jax.lax.rsqrt(var + 1e-06) + + # normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is + # a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to + # a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs. + scale = jnp.expand_dims(self.scale.value, axis=range(len(x.shape) - 1)) + return normed_inputs * (1 + scale) + + +class T5Gemma2FeedForward(nnx.Module): + """Feed-forward module with gated activation.""" + + def __init__( + self, + features: int, + hidden_dim: int, + *, + transpose_gating_einsum: bool, + kernel_init: nnx.Initializer = NORMAL_INIT, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.features = features + self.hidden_dim = hidden_dim + self.transpose_gating_einsum = transpose_gating_einsum + + if transpose_gating_einsum: + gating_shape = (2, hidden_dim, features) + else: + gating_shape = (2, features, hidden_dim) + + self.gating_einsum = Gemma3Einsum( + shape=gating_shape, + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + self.linear = Gemma3Einsum( + shape=(hidden_dim, features), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: + """Applies feed-forward transformation. + + Args: + x: Input of shape [batch_size, seq_len, features]. + + Returns: + Output of shape [batch_size, seq_len, features]. + """ + eq = "...F,NHF->...NH" if self.transpose_gating_einsum else "...F,NFH->...NH" + gate = self.gating_einsum(eq, x) + activations = jax.nn.gelu(gate[..., 0, :]) * gate[..., 1, :] + return self.linear("...H,HF->...F", activations) + + +# ============================================================================= +# Attention Masks +# ============================================================================= + + +def make_bidirectional_mask( + input_mask: Bool[Array, "B L"], +) -> Bool[Array, "B 1 L L"]: + """Creates a bidirectional attention mask for encoder. + + Args: + input_mask: Boolean mask where True indicates valid tokens. + + Returns: + Attention mask of shape [B, 1, L, L]. + """ + # [B, L] -> [B, 1, 1, L] + mask = input_mask[:, None, None, :] + return mask + + +def make_causal_mask( + input_mask: Bool[Array, "B L"], +) -> Bool[Array, "B 1 L L"]: + """Creates a causal attention mask for decoder. + + Args: + input_mask: Boolean mask where True indicates valid tokens. + + Returns: + Causal attention mask of shape [B, 1, L, L]. + """ + seq_len = input_mask.shape[-1] + # Create causal mask: [L, L] + causal = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) + # Combine with input mask: [B, 1, L, L] + mask = input_mask[:, None, None, :] & causal[None, None, :, :] + return mask + + +def make_sliding_window_mask( + positions: Int[Array, "B L"], + sliding_window: int, + *, + bidirectional: bool = False, +) -> Bool[Array, "B 1 L L"]: + """Creates a sliding window mask. + + Args: + positions: Position indices of shape [B, L]. + sliding_window: Size of the sliding window. + bidirectional: If True, creates a symmetric bidirectional window where + the total window size equals `sliding_window` (split evenly on both + sides). If False, creates a causal window of size `sliding_window`. + + Returns: + Sliding window mask of shape [B, 1, L, L]. + """ + # [B, L, 1] and [B, 1, L] + q_pos = positions[:, :, None] + k_pos = positions[:, None, :] + + if bidirectional: + # Bidirectional: symmetric window centered on query position + # Total window size = sliding_window, split evenly on both sides + # Matches PyTorch: left_window = (sw + 1) // 2, right_window = sw // 2 + 1 + left_window = (sliding_window + 1) // 2 + right_window = sliding_window // 2 + 1 + dist = q_pos - k_pos + left_mask = (dist >= 0) & (dist < left_window) + right_mask = (dist < 0) & (-dist < right_window) + mask = left_mask | right_mask + else: + # Causal: can only attend to past tokens within window + dist = jnp.abs(q_pos - k_pos) + mask = dist < sliding_window + + return mask[:, None, :, :] + + +def make_sliding_window_causal_mask( + input_mask: Bool[Array, "B L"], + positions: Int[Array, "B L"], + sliding_window: int, +) -> Bool[Array, "B 1 L L"]: + """Creates a causal mask with sliding window for decoder. + + Combines causal mask (lower triangular) with sliding window constraint. + + Args: + input_mask: Boolean mask where True indicates valid tokens. + positions: Position indices of shape [B, L]. + sliding_window: Size of the sliding window. + + Returns: + Sliding window causal mask of shape [B, 1, L, L]. + """ + # Start with causal mask + causal_mask = make_causal_mask(input_mask) + + # Create sliding window mask + sliding_mask = make_sliding_window_mask(positions, sliding_window) + + # Combine: must satisfy both causal AND sliding window + return causal_mask & sliding_mask + + +def make_merged_attention_mask( + decoder_mask: Bool[Array, "B 1 L_dec L_dec"], + encoder_mask: Bool[Array, "B 1 1 L_enc"], +) -> Bool[Array, "B 1 L_dec L_combined"]: + """Creates merged attention mask for decoder's merged attention. + + Concatenates the decoder self-attention mask with the cross-attention + mask to form a single mask for the merged attention operation. + + Args: + decoder_mask: Causal mask for decoder self-attention [B, 1, L_dec, L_dec]. + encoder_mask: Mask for encoder hidden states [B, 1, 1, L_enc]. + + Returns: + Merged mask of shape [B, 1, L_dec, L_dec + L_enc]. + """ + batch_size, _, seq_len, _ = decoder_mask.shape + enc_len = encoder_mask.shape[-1] + + # Broadcast encoder mask to [B, 1, L_dec, L_enc] + cross_mask = jnp.broadcast_to(encoder_mask, (batch_size, 1, seq_len, enc_len)) + + # Concatenate along key dimension + return jnp.concatenate([decoder_mask, cross_mask], axis=-1) + + +def make_decode_mode_self_mask( + batch_size: int, + query_len: int, + cache_len: int, +) -> Bool[Array, "B 1 Q K"]: + """Creates self-attention mask for decode mode (with KV cache). + + During decode mode, the query attends to all previously cached keys. + Since we're generating tokens sequentially and causality is already + enforced by the order of generation, all cached positions are valid. + + Args: + batch_size: Batch size. + query_len: Number of query positions (typically 1 for decode mode). + cache_len: Total length of cached keys (cache_index + query_len). + + Returns: + Mask of shape [B, 1, query_len, cache_len] with all True. + """ + return jnp.ones((batch_size, 1, query_len, cache_len), dtype=jnp.bool_) + + +def make_decode_mode_sliding_mask( + batch_size: int, + query_pos: int, + cache_len: int, + sliding_window: int, +) -> Bool[Array, "B 1 1 K"]: + """Creates sliding window mask for decode mode (with KV cache). + + During decode mode with sliding window, the query can only attend to + positions within the sliding window range. + + Args: + batch_size: Batch size. + query_pos: Position of the current query (typically cache_index). + cache_len: Total length of cached keys. + sliding_window: Size of the sliding window. + + Returns: + Mask of shape [B, 1, 1, cache_len] with True for valid positions. + """ + # Key positions: 0, 1, 2, ..., cache_len-1 + k_positions = jnp.arange(cache_len) + # Within sliding window: |query_pos - k_pos| < sliding_window + in_window = jnp.abs(query_pos - k_positions) < sliding_window + # Also must be causal (k_pos <= query_pos) + is_causal = k_positions <= query_pos + mask = in_window & is_causal + return jnp.broadcast_to(mask[None, None, None, :], (batch_size, 1, 1, cache_len)) + + +# ============================================================================= +# Text Embeddings +# ============================================================================= + + +class T5Gemma2ScaledWordEmbedding(nnx.Module): + """Scaled word embedding with special EOI (end-of-image) token. + + The embeddings are scaled by the provided embed_scale, and the EOI token + has a separate learnable embedding. + """ + + def __init__( + self, + vocab_size: int, + embed_dim: int, + *, + padding_idx: int = 0, + embed_scale: float | None = None, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.vocab_size = vocab_size + self.embed_dim = embed_dim + self.padding_idx = padding_idx + # If embed_scale is not provided, default to sqrt(embed_dim) + if embed_scale is None: + embed_scale = embed_dim**0.5 + self.embed_scale = jnp.array(embed_scale, dtype=dtype) + + self.embedding = nnx.Param( + NORMAL_INIT(rngs.params(), (vocab_size, embed_dim), dtype), + ) + self.eoi_embedding = nnx.Param( + ZEROS_INIT(rngs.params(), (embed_dim,), dtype), + ) + + def __call__(self, input_ids: Int[Array, "B L"]) -> Float[Array, "B L D"]: + """Embed input tokens with scaling. + + Args: + input_ids: Input token IDs of shape [B, L]. + + Returns: + Embeddings of shape [B, L, hidden_size]. + """ + embeddings = self.embedding[input_ids] * self.embed_scale + + # Replace EOI token embeddings + eoi_mask = input_ids == END_OF_IMAGE_TOKEN + + return jnp.where(eoi_mask[..., None], self.eoi_embedding[...], embeddings) + + +# ============================================================================= +# Vision Components +# ============================================================================= + + +def _posemb_sincos_2d( + h: int, + w: int, + *, + width: int, + temperature: float = 10_000.0, + dtype: jnp.dtype = jnp.float32, +) -> Float[Array, "1 M D"]: + """Sinusoidal 2D position embeddings (MoCo v3 style).""" + y, x = jnp.mgrid[:h, :w] + + assert width % 4 == 0, "Width must be mult of 4 for sincos posemb" + omega = jnp.arange(width // 4) / (width // 4 - 1) + omega = 1.0 / (temperature**omega) + y = jnp.einsum("m,d->md", y.flatten(), omega) + x = jnp.einsum("m,d->md", x.flatten(), omega) + pe = jnp.concatenate([jnp.sin(x), jnp.cos(x), jnp.sin(y), jnp.cos(y)], axis=1) + return jnp.asarray(pe, dtype)[None, :, :] + + +class T5Gemma2VisionMLP(nnx.Module): + """MLP for Vision Transformer.""" + + def __init__( + self, + *, + width: int, + mlp_dim: int, + dropout: float = 0.0, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.dense1 = nnx.Linear( + in_features=width, + out_features=mlp_dim, + dtype=dtype, + kernel_init=nnx.initializers.xavier_uniform(), + bias_init=nnx.initializers.normal(stddev=1e-6), + rngs=rngs, + ) + self.dense2 = nnx.Linear( + in_features=mlp_dim, + out_features=width, + dtype=dtype, + kernel_init=nnx.initializers.xavier_uniform(), + bias_init=nnx.initializers.normal(stddev=1e-6), + rngs=rngs, + ) + self.dropout = nnx.Dropout(rate=dropout, rngs=rngs) + + def __call__( + self, + x: Float[Array, "B L D"], + *, + deterministic: bool | None = None, + ) -> Float[Array, "B L D"]: + x = self.dense1(x) + x = nnx.gelu(x) + x = self.dropout(x, deterministic=deterministic) + return self.dense2(x) + + +class T5Gemma2VisionEncoderBlock(nnx.Module): + """Transformer encoder block for Vision Transformer.""" + + def __init__( + self, + *, + width: int, + mlp_dim: int, + num_heads: int = 12, + dropout: float = 0.0, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.layer_norm1 = nnx.LayerNorm( + num_features=width, + scale_init=nnx.initializers.ones_init(), + bias_init=nnx.initializers.zeros_init(), + rngs=rngs, + ) + self.layer_norm2 = nnx.LayerNorm( + num_features=width, + scale_init=nnx.initializers.ones_init(), + bias_init=nnx.initializers.zeros_init(), + rngs=rngs, + ) + + self.dropout1 = nnx.Dropout(rate=dropout, rngs=rngs) + self.dropout2 = nnx.Dropout(rate=dropout, rngs=rngs) + + self.mha = nnx.MultiHeadAttention( + num_heads=num_heads, + in_features=width, + dtype=dtype, + decode=False, + kernel_init=nnx.initializers.xavier_uniform(), + out_kernel_init=nnx.initializers.xavier_uniform(), + bias_init=nnx.initializers.zeros_init(), + out_bias_init=nnx.initializers.zeros_init(), + rngs=rngs, + ) + + self.mlp = Gemma3VisionMLP( + width=width, + mlp_dim=mlp_dim, + dropout=dropout, + dtype=dtype, + rngs=rngs, + ) + + def __call__( + self, + x: Float[Array, "B L D"], + *, + deterministic: bool | None = None, + ) -> Float[Array, "B L D"]: + y = self.layer_norm1(x) + y = self.mha(y, deterministic=deterministic) + y = self.dropout1(y, deterministic=deterministic) + x = x + y + + y = self.layer_norm2(x) + y = self.mlp(y, deterministic=deterministic) + y = self.dropout2(y, deterministic=deterministic) + return x + y + + +class T5Gemma2VisionEncoder(nnx.Module): + """Vision Transformer Encoder.""" + + def __init__( + self, + *, + width: int, + depth: int, + mlp_dim: int, + num_heads: int = 12, + dropout: float = 0.0, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.depth = depth + + self.blocks = nnx.List( + [ + T5Gemma2VisionEncoderBlock( + width=width, + mlp_dim=mlp_dim, + num_heads=num_heads, + dropout=dropout, + dtype=dtype, + rngs=rngs, + ) + for _ in range(depth) + ] + ) + + self.encoder_norm = nnx.LayerNorm( + num_features=width, + scale_init=nnx.initializers.ones_init(), + bias_init=nnx.initializers.zeros_init(), + rngs=rngs, + ) + + def __call__( + self, + x: jax.Array, + *, + deterministic: bool | None = None, + ) -> jax.Array: + for i in range(self.depth): + x = self.blocks[i](x, deterministic=deterministic) + return self.encoder_norm(x) + + +class T5Gemma2VisionExit(nnx.Module): + """Vision exit layer - spatially pools soft tokens to output length.""" + + def __init__(self, output_length: int = 256, *, rngs: nnx.Rngs): + self.output_length = output_length + + def __call__(self, x: jax.Array) -> jax.Array: + cur_length = x.shape[1] + if cur_length == self.output_length: + return x + + cur_width = int(cur_length**0.5) + assert cur_width**2 == cur_length + output_width = int(self.output_length**0.5) + assert output_width**2 == self.output_length, f"Cannot pool {x.shape=} to {self.output_length}=!" + + batch_size = x.shape[0] + embed_dim = x.shape[-1] + x = jnp.reshape(x, (batch_size, cur_width, cur_width, embed_dim)) + assert not cur_width % output_width, f"{cur_width=} {output_width=}" + window = cur_width // output_width + window_shape = (window, window) + x = nnx.avg_pool(x, window_shape=window_shape, strides=window_shape) + batch_size, height, width, embed_dim = x.shape + return jnp.reshape(x, (batch_size, height * width, embed_dim)) + + +class T5Gemma2VisionSoftTokenizer(nnx.Module): + """Vision soft tokenizer (ViT trained with SigLiP). + + Transforms images into soft tokens that can be embedded into the + text embedding space. + """ + + def __init__( + self, + config: VisionEncoderConfig, + *, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.config = config + self.rngs = rngs + + self.embedding = nnx.Conv( + in_features=3, + out_features=config.width, + kernel_size=(config.patch_size, config.patch_size), + strides=(config.patch_size, config.patch_size), + padding="VALID", + kernel_init=nnx.initializers.xavier_uniform(), + bias_init=nnx.initializers.zeros_init(), + dtype=dtype, + rngs=rngs, + ) + + self.pos_embedding = self._get_posemb( + config.posemb, + seqshape=( + config.image_size // config.patch_size, + config.image_size // config.patch_size, + ), + width=config.width, + dtype=dtype, + ) + + self.dropout = nnx.Dropout(rate=config.dropout, rngs=rngs) + + self.transformer = Gemma3VisionEncoder( + width=config.width, + depth=config.depth, + mlp_dim=config.mlp_dim, + num_heads=config.num_heads, + dropout=config.dropout, + dtype=dtype, + rngs=rngs, + ) + + self.vision_exit = Gemma3VisionExit(output_length=256, rngs=rngs) + + def __call__( + self, + images: Float[Array, "B N H W C"], + *, + deterministic: bool | None = None, + ) -> Float[Array, "B N P D"]: + if len(images.shape) == 4: + images = images[:, None, :] + b, n, h, w, c = images.shape + x = jnp.reshape(images, [b * n, h, w, c]) + + x = self.embedding(x) + bn, h, w, c = x.shape + x = jnp.reshape(x, [bn, h * w, c]) + + x = x + self.pos_embedding.value + x = self.dropout(x, deterministic=deterministic) + x = self.transformer(x, deterministic=deterministic) + x = self.vision_exit(x) + + bn, s, d = x.shape + return jnp.reshape(x, [b, n, s, d]) + + def _get_posemb( + self, + typ: str, + *, + seqshape: tuple[int, int], + width: int, + dtype: jnp.dtype = jnp.float32, + ) -> nnx.Param: + """Returns the position embedding.""" + if typ == "learn": + shape = (1, seqshape[0] * seqshape[1], width) + initializer = nnx.initializers.normal(stddev=1 / (width**0.5)) + return nnx.Param(initializer(self.rngs.params(), shape, dtype)) + if typ == "sincos2d": + return nnx.Param( + _posemb_sincos_2d(*seqshape, width=width, dtype=dtype), + ) + raise ValueError(f"Unknown posemb type: {typ}") + + +class T5Gemma2VisionSoftTokensEmbedder(nnx.Module): + """Embeds vision soft tokens into the text embedding space.""" + + def __init__( + self, + embed_dim: int, + *, + soft_tokens_dim: int, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.mm_soft_embedding_norm = T5Gemma2RMSNorm( + soft_tokens_dim, + scale_init=nnx.initializers.zeros_init(), + dtype=dtype, + rngs=rngs, + ) + self.mm_input_projection = T5Gemma2Einsum( + (soft_tokens_dim, embed_dim), + kernel_init=nnx.initializers.normal(), + dtype=dtype, + rngs=rngs, + ) + + def __call__(self, x: Float[Array, "B N P Dv"]) -> Float[Array, "B N P De"]: + x = self.mm_soft_embedding_norm(x) + return self.mm_input_projection("...tm,md->...td", x) + + +class T5Gemma2VisionEmbedder(nnx.Module): + """Vision embedder - tokenizes images and embeds into text space.""" + + def __init__( + self, + *, + vision_config: VisionEncoderConfig, + embed_dim: int, + freeze_params: bool = True, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.embed_dim = embed_dim + self.freeze_params = freeze_params + + self.soft_tokenizer = T5Gemma2VisionSoftTokenizer( + vision_config, + dtype=dtype, + rngs=rngs, + ) + self.soft_tokens_embedder = T5Gemma2VisionSoftTokensEmbedder( + embed_dim, + soft_tokens_dim=vision_config.width, + dtype=dtype, + rngs=rngs, + ) + + def __call__( + self, + images: Float[Array, "B N H W C"], + *, + deterministic: bool | None = None, + ) -> Float[Array, "B N P De"]: + soft_tokens = self.soft_tokenizer(images, deterministic=deterministic) + + if self.freeze_params: + soft_tokens = jax.lax.stop_gradient(soft_tokens) + + return self.soft_tokens_embedder(soft_tokens) + + +# ============================================================================= +# Multimodal Utilities +# ============================================================================= +def _get_new_text_positions(*, offset_on: np.ndarray, offset_by: int) -> np.ndarray: + """Create the positions of the new tokens.""" + offset = np.cumsum(offset_on, axis=-1) * offset_by + new_positions = np.arange(offset_on.shape[-1]) + offset + new_positions -= offset_by * offset_on + return new_positions + + +def _insert_sequence( + tokens: np.ndarray, + *, + at: int, + sequence: list[int], + max_num_images: int, +) -> np.ndarray: + """Inserts a sequence of tokens at all occurrences of a specific token.""" + original_dim = tokens.ndim + if original_dim == 1: + tokens = tokens[None, :] + + batch_size, length = tokens.shape + mm_tokens_to_insert = np.array(sequence) + offset_by = len(mm_tokens_to_insert) - 1 + length_with_mm = length + max_num_images * offset_by + mm_start = tokens == at + + new_tokens = np.zeros((batch_size, length_with_mm), dtype=np.int64) + new_text_pos = _get_new_text_positions(offset_on=mm_start, offset_by=offset_by) + np.put_along_axis(new_tokens, new_text_pos, tokens, axis=1) + + batch_indices_to_zero, _ = np.where(mm_start) + new_pos_to_zero = new_text_pos[mm_start] + if batch_indices_to_zero.size > 0: + new_tokens[batch_indices_to_zero, new_pos_to_zero] = 0 + + batch_indices, seq_indices = np.nonzero(mm_start) + + if batch_indices.size > 0: + intra_batch_img_idx = np.cumsum(mm_start, axis=1)[mm_start] - 1 + final_img_start_pos = seq_indices + intra_batch_img_idx * offset_by + indices_to_insert = final_img_start_pos[:, None] + np.arange( + len(mm_tokens_to_insert), + ) + new_tokens[batch_indices[:, None], indices_to_insert] = mm_tokens_to_insert + + if original_dim == 1: + new_tokens = np.squeeze(new_tokens) + return new_tokens + + +def add_extra_tokens_for_images( + tokens: np.ndarray | list, + *, + max_num_images: int = 1, +) -> np.ndarray: + """Add extra image tokens to text tokens.""" + + mm_tokens = [ + NEW_LINE_TOKEN, + START_OF_IMAGE_TOKEN, + *[IMAGE_PLACEHOLDER_TOKEN] * NUM_PLACEHOLDER_TOKENS_PER_IMAGE, + END_OF_IMAGE_TOKEN, + NEW_LINE_TOKEN, + ] + if not isinstance(tokens, np.ndarray): + tokens = np.asarray(tokens) + + return _insert_sequence( + at=START_OF_IMAGE_TOKEN, + sequence=mm_tokens, + tokens=tokens, + max_num_images=max_num_images, + ) + + +def merge_mm_embeddings( + text_embeddings: jnp.ndarray, + vision_embeddings: jnp.ndarray, + mask: jnp.ndarray, +) -> jnp.ndarray: + """Merge text and multimodal (vision) embeddings.""" + return jax.vmap(_merge_mm_embeddings_inner, in_axes=(0, 0, 0))( + text_embeddings, + vision_embeddings, + mask, + ) + + +def _merge_mm_embeddings_inner( + text_embeddings: jnp.ndarray, + vision_embeddings: jnp.ndarray, + mask: jnp.ndarray, +) -> jnp.ndarray: + """Merge embeddings without batch dimension.""" + num_images, num_toks_per_image, d = vision_embeddings.shape + vision_embeddings = jnp.reshape( + vision_embeddings, + (num_images * num_toks_per_image, d), + ) + + target_pos = jnp.nonzero(mask, size=len(vision_embeddings)) + first_pos = text_embeddings[0] + merged = text_embeddings.at[target_pos, :].set(vision_embeddings) + return merged.at[0].set(first_pos) + + +# ============================================================================= +# Encoder Components +# ============================================================================= + + +class T5Gemma2EncoderAttention(nnx.Module): + """Bidirectional self-attention for T5Gemma2 encoder. + + This is separate from Gemma3Attention because the encoder needs + bidirectional sliding window attention, not causal sliding window. + """ + + def __init__( + self, + num_q_heads: int, + num_kv_heads: int, + features: int, + head_dim: int, + attn_type: AttentionType, + query_pre_attn_scalar: float, + *, + rope_base_frequency: int = 10_000, + rope_scale_factor: float = 1.0, + attn_logits_soft_cap: float | None = None, + sliding_window_size: int | None = None, + use_qk_norm: bool = False, + kernel_init: nnx.Initializer = NORMAL_INIT, + scale_init: nnx.Initializer = ZEROS_INIT, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.num_q_heads = num_q_heads + self.num_kv_heads = num_kv_heads + self.features = features + self.head_dim = head_dim + self.attn_type = attn_type + self.query_pre_attn_scalar = query_pre_attn_scalar + self.rope_base_frequency = rope_base_frequency + self.rope_scale_factor = rope_scale_factor + self.attn_logits_soft_cap = attn_logits_soft_cap + self.sliding_window_size = sliding_window_size + self.use_qk_norm = use_qk_norm + + self.attn_vec_einsum = T5Gemma2Einsum( + shape=(num_q_heads, head_dim, features), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + # Check if we can use combined QKV projection + if num_kv_heads == num_q_heads: + self.qkv_einsum = T5Gemma2Einsum( + shape=(3, num_q_heads, features, head_dim), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + self.q_einsum = None + self.kv_einsum = None + else: + self.qkv_einsum = None + self.q_einsum = T5Gemma2Einsum( + shape=(num_q_heads, features, head_dim), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + self.kv_einsum = T5Gemma2Einsum( + shape=(2, num_kv_heads, features, head_dim), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + if use_qk_norm: + self._query_norm = T5Gemma2RMSNorm( + head_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + self._key_norm = T5Gemma2RMSNorm( + head_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + @property + def use_qkv_einsum(self) -> bool: + return self.qkv_einsum is not None + + @property + def use_gqa(self) -> bool: + return self.num_kv_heads != self.num_q_heads and self.num_kv_heads > 1 + + def __call__( + self, + x: Float[Array, "B L D"], + segment_pos: Float[Array, "B L"], + attn_mask: Float[Array, "B L L"], + ) -> Float[Array, "B L D"]: + """Applies bidirectional multi-head attention. + + Args: + x: Input sequence of shape [batch_size, seq_len, embed_dim]. + segment_pos: Absolute positions of shape [batch_size, seq_len]. + attn_mask: Attention mask of shape [batch_size, seq_len, seq_len]. + + Returns: + Output sequence of shape [batch_size, seq_len, embed_dim]. + """ + # Project Q, K, V + if self.use_qkv_einsum: + query_proj, key_proj, value_proj = self.qkv_einsum("BTD,SNDH->SBTNH", x) + else: + query_proj = self.q_einsum("BTD,NDH->BTNH", x) + key_proj, value_proj = self.kv_einsum("BSD,CKDH->CBSKH", x) + + # Apply QK normalization + if self.use_qk_norm: + query_proj = self._query_norm(query_proj) + key_proj = self._key_norm(key_proj) + + # Apply RoPE + query_proj = apply_rope( + query_proj, + segment_pos, + base_frequency=self.rope_base_frequency, + scale_factor=self.rope_scale_factor, + ) + query_scaled = query_proj * self.query_pre_attn_scalar + + key_proj = apply_rope( + key_proj, + segment_pos, + base_frequency=self.rope_base_frequency, + scale_factor=self.rope_scale_factor, + ) + + # Compute attention logits + if self.use_gqa: + b, t, kg, h = query_scaled.shape + query_scaled = query_scaled.reshape( + (b, t, self.num_kv_heads, int(kg / self.num_kv_heads), h), + ) + logits = jnp.einsum("BTKGH,BSKH->BTKGS", query_scaled, key_proj) + b, t, k, g, s = logits.shape + logits = logits.reshape((b, t, k * g, s)) + else: + logits = jnp.einsum("BTNH,BSNH->BTNS", query_scaled, key_proj) + + # Apply soft capping + if self.attn_logits_soft_cap is not None: + logits = jnp.tanh(logits / self.attn_logits_soft_cap) + logits = logits * self.attn_logits_soft_cap + + # Apply bidirectional sliding window mask for LOCAL_SLIDING layers + if self.attn_type == AttentionType.LOCAL_SLIDING: + if self.sliding_window_size is None: + raise ValueError( + "sliding_window_size must be set for LOCAL_SLIDING attention", + ) + # Use make_sliding_window_mask with bidirectional=True, squeeze to [B, L, L] + sliding_mask = make_sliding_window_mask( + segment_pos, + self.sliding_window_size, + bidirectional=True, + )[:, 0, :, :] + attn_mask = attn_mask * sliding_mask + + # Apply mask and softmax + padded_logits = jnp.where(jnp.expand_dims(attn_mask, -2), logits, K_MASK) + probs = jax.nn.softmax(padded_logits, axis=-1).astype(key_proj.dtype) + + # Compute attention output + if self.use_gqa: + b, t, kg, h = probs.shape + probs = probs.reshape( + (b, t, self.num_kv_heads, int(kg / self.num_kv_heads), h), + ) + encoded = jnp.einsum("BTKGS,BSKH->BTKGH", probs, value_proj) + b, t, k, g, h = encoded.shape + encoded = encoded.reshape((b, t, k * g, h)) + else: + encoded = jnp.einsum("BTNS,BSNH->BTNH", probs, value_proj) + + return self.attn_vec_einsum("BTNH,NHD->BTD", encoded) + + +class T5Gemma2EncoderBlock(nnx.Module): + """Encoder transformer block with bidirectional attention. + + Uses T5Gemma2EncoderAttention which handles bidirectional sliding window + correctly (unlike Gemma3Attention which uses causal sliding window). + """ + + def __init__( + self, + num_q_heads: int, + num_kv_heads: int, + embed_dim: int, + head_dim: int, + hidden_dim: int, + *, + use_post_attn_norm: bool, + use_post_ffw_norm: bool, + attn_type: AttentionType, + query_pre_attn_scalar: float, + transpose_gating_einsum: bool, + rope_base_frequency: int = 10_000, + rope_scale_factor: float = 1.0, + attn_logits_soft_cap: float | None = None, + sliding_window_size: int | None = None, + use_qk_norm: bool = False, + kernel_init: nnx.Initializer = NORMAL_INIT, + scale_init: nnx.Initializer = ZEROS_INIT, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.use_post_attn_norm = use_post_attn_norm + self.use_post_ffw_norm = use_post_ffw_norm + + self.pre_attention_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + self.attn = T5Gemma2EncoderAttention( + num_q_heads=num_q_heads, + features=embed_dim, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + attn_type=attn_type, + query_pre_attn_scalar=query_pre_attn_scalar, + rope_base_frequency=rope_base_frequency, + rope_scale_factor=rope_scale_factor, + attn_logits_soft_cap=attn_logits_soft_cap, + sliding_window_size=sliding_window_size, + use_qk_norm=use_qk_norm, + scale_init=scale_init, + dtype=dtype, + rngs=rngs, + ) + + if use_post_attn_norm: + self.post_attention_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + self.pre_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + self.mlp = T5Gemma2FeedForward( + features=embed_dim, + hidden_dim=hidden_dim, + transpose_gating_einsum=transpose_gating_einsum, + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + if use_post_ffw_norm: + self.post_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + def __call__( + self, + x: Float[Array, "B L D"], + segment_pos: Int[Array, "B L"], + attn_mask: Bool[Array, "B L L"], + ) -> Float[Array, "B L D"]: + """Apply encoder layer. + + Args: + x: Input hidden states [B, L, D]. + segment_pos: Position indices [B, L]. + attn_mask: Bidirectional attention mask [B, L, L]. + + Returns: + Output hidden states [B, L, D]. + """ + # Attention block + inputs_normalized = self.pre_attention_norm(x) + attn_output = self.attn(inputs_normalized, segment_pos, attn_mask) + + if self.use_post_attn_norm: + attn_output = self.post_attention_norm(attn_output) + + attn_output = attn_output + x + + # Feed-forward block + outputs = self.pre_ffw_norm(attn_output) + outputs = self.mlp(outputs) + + if self.use_post_ffw_norm: + outputs = self.post_ffw_norm(outputs) + + return outputs + attn_output + + +class T5Gemma2Encoder(nnx.Module): + """T5Gemma2 Encoder with optional vision support. + + Processes text (and optional images) using bidirectional self-attention. + """ + + def __init__( + self, + config: T5Gemma2EncoderConfig, + *, + embedder: nnx.Module | None = None, + dtype: jnp.dtype = jnp.float32, + dtype_mm: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.config = config + text_config = config.text_config + + # Text embeddings + self.embedder = embedder or T5Gemma2ScaledWordEmbedding( + vocab_size=text_config.vocab_size, + embed_dim=text_config.embed_dim, + padding_idx=text_config.pad_token_id, + dtype=dtype, + rngs=rngs, + ) + + # Vision encoder (optional) - uses Gemma3 vision component + self.vision_embedder = nnx.data(None) + if config.vision_config is not None: + self.vision_embedder = T5Gemma2VisionEmbedder( + vision_config=config.vision_config, + embed_dim=text_config.embed_dim, + freeze_params=True, + dtype=dtype_mm, + rngs=rngs, + ) + + # Determine attention types per layer + attention_types = text_config.layer_types + if not attention_types: + attention_types = tuple(AttentionType.GLOBAL for _ in range(text_config.num_hidden_layers)) + + # Encoder layers (use Gemma3Block via alias) + self.blocks = nnx.List( + [ + T5Gemma2EncoderBlock( + num_q_heads=text_config.num_attention_heads, + num_kv_heads=text_config.num_key_value_heads, + embed_dim=text_config.embed_dim, + head_dim=text_config.head_dim, + hidden_dim=text_config.intermediate_size, + use_post_attn_norm=True, + use_post_ffw_norm=True, + attn_type=attention_types[i], + query_pre_attn_scalar=text_config.query_pre_attn_scalar, + transpose_gating_einsum=True, + rope_base_frequency=_get_rope_base_frequency( + text_config, + attention_types[i], + ), + rope_scale_factor=_get_rope_scale_factor( + text_config, + attention_types[i], + ), + attn_logits_soft_cap=text_config.attn_logit_softcapping, + sliding_window_size=( + text_config.sliding_window if attention_types[i] == AttentionType.LOCAL_SLIDING else None + ), + use_qk_norm=True, + dtype=dtype, + rngs=rngs, + ) + for i in range(text_config.num_hidden_layers) + ] + ) + + # Final normalization + self.norm = T5Gemma2RMSNorm(text_config.embed_dim, dtype=dtype, rngs=rngs) + + def __call__( + self, + tokens: Int[Array, "B L"], + attention_mask: Bool[Array, "B L"] | None = None, + position_ids: Int[Array, "B L"] | None = None, + images: Float[Array, "B N H W C"] | None = None, + *, + deterministic: bool = True, + ) -> Float[Array, "B L D"]: + """Forward pass of the encoder. + + Args: + input_ids: Input token IDs [B, L]. + attention_mask: Attention mask [B, L]. + position_ids: Position indices [B, L]. + pixel_values: Optional input images [B, N, H, W, C] where N is images per batch. + deterministic: Whether to run in deterministic mode. + + Returns: + Encoder hidden states [B, L, D]. + """ + batch_size, seq_len = tokens.shape + + # Create position IDs if not provided + if position_ids is None: + position_ids = jnp.arange(seq_len)[None, :].repeat(batch_size, axis=0) + + # Create attention mask if not provided + if attention_mask is None: + attention_mask = jnp.ones((batch_size, seq_len), dtype=jnp.bool_) + + # Embed tokens + x = self.embedder(tokens) + + # Process images and merge with text embeddings + if images is not None and self.vision_embedder is not None: + image_embeddings = self.vision_embedder(images, deterministic=deterministic) + x = merge_mm_embeddings( + x, + image_embeddings, + tokens == IMAGE_PLACEHOLDER_TOKEN, + ) + + # Create bidirectional attention mask for encoder [B, L, L] + # Each query can attend to all valid (non-padded) key positions + attn_mask = make_bidirectional_mask(attention_mask) + # Squeeze to [B, L, L] if it's [B, 1, L, L] + if attn_mask.ndim == 4: + attn_mask = attn_mask[:, 0, :, :] + + # Ensure mask is [B, L, L] by broadcasting if needed + if attn_mask.shape[-2] == 1: + # [B, 1, L] -> [B, L, L] + attn_mask = jnp.broadcast_to(attn_mask, (batch_size, seq_len, seq_len)) + + # Apply encoder layers + # T5Gemma2EncoderBlock handles bidirectional sliding window internally + for block in self.blocks: + x = block(x, position_ids, attn_mask) + + # Final normalization + return self.norm(x) + + +# ============================================================================= +# Decoder Components +# ============================================================================= + + +class T5Gemma2MergedAttention(nnx.Module): + """Merged self-attention and cross-attention for decoder. + + Combines self-attention and cross-attention in a single operation: + - Query comes from decoder hidden states + - Keys/Values are concatenation of [self_kv, cross_kv] + - Mask is concatenation of [causal_mask, encoder_mask] + + This is more efficient than separate attention operations. + + Uses nnx.Cache for autoregressive decoding. Call init_cache() before + using decode=True. + """ + + def __init__( + self, + num_q_heads: int, + num_kv_heads: int, + features: int, + head_dim: int, + query_pre_attn_scalar: float, + *, + rope_base_frequency: int = 10_000, + rope_scale_factor: float = 1.0, + attn_logits_soft_cap: float | None = None, + use_qk_norm: bool = True, + kernel_init: nnx.Initializer = NORMAL_INIT, + scale_init: nnx.Initializer = ZEROS_INIT, + dtype: jnp.dtype = jnp.float32, + decode: bool | None = None, + rngs: nnx.Rngs, + ): + self.num_q_heads = num_q_heads + self.num_kv_heads = num_kv_heads + self.features = features + self.head_dim = head_dim + self.query_pre_attn_scalar = query_pre_attn_scalar + self.rope_base_frequency = rope_base_frequency + self.rope_scale_factor = rope_scale_factor + self.attn_logits_soft_cap = attn_logits_soft_cap + self.use_qk_norm = use_qk_norm + self.decode = decode + + self.num_kv_groups = num_q_heads // num_kv_heads + + # Projections + self.q_proj = T5Gemma2Einsum( + shape=(num_q_heads, features, head_dim), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + self.k_proj = T5Gemma2Einsum( + shape=(num_kv_heads, features, head_dim), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + self.v_proj = T5Gemma2Einsum( + shape=(num_kv_heads, features, head_dim), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + self.o_proj = T5Gemma2Einsum( + shape=(num_q_heads, head_dim, features), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + # QK normalization + if use_qk_norm: + self.q_norm = T5Gemma2RMSNorm( + head_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + self.k_norm = T5Gemma2RMSNorm( + head_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + # Cache for autoregressive decoding (nnx.Cache pattern like Gemma3) + # Self-attention cache + self.cached_self_key: nnx.Cache[Array] | None = nnx.data(None) + self.cached_self_value: nnx.Cache[Array] | None = nnx.data(None) + # Cross-attention cache (computed once from encoder outputs) + self.cached_cross_key: nnx.Cache[Array] | None = nnx.data(None) + self.cached_cross_value: nnx.Cache[Array] | None = nnx.data(None) + # Cache index for self-attention + self.cache_index: nnx.Cache[Array] | None = nnx.data(None) + + def init_cache( + self, + *, + batch_size: int, + max_decode_length: int, + encoder_seq_length: int, + dtype: jnp.dtype = jnp.float32, + ) -> None: + """Initialize KV caches for autoregressive decoding. + + Must be called before using decode=True. + + Args: + batch_size: Number of sequences in the batch. + max_decode_length: Maximum decoder sequence length. + encoder_seq_length: Encoder sequence length (for cross-attention). + dtype: Data type for cache arrays. + """ + # Self-attention cache shape: [B, max_decode_len, num_kv_heads, head_dim] + self_cache_shape = ( + batch_size, + max_decode_length, + self.num_kv_heads, + self.head_dim, + ) + self.cached_self_key = nnx.Cache(jnp.zeros(self_cache_shape, dtype)) + self.cached_self_value = nnx.Cache(jnp.zeros(self_cache_shape, dtype)) + + # Cross-attention cache shape: [B, enc_seq_len, num_kv_heads, head_dim] + cross_cache_shape = ( + batch_size, + encoder_seq_length, + self.num_kv_heads, + self.head_dim, + ) + self.cached_cross_key = nnx.Cache(jnp.zeros(cross_cache_shape, dtype)) + self.cached_cross_value = nnx.Cache(jnp.zeros(cross_cache_shape, dtype)) + + # Cache index tracks current position in self-attention cache + self.cache_index = nnx.Cache(jnp.array(0, dtype=jnp.int32)) + + def __call__( + self, + hidden_states: Float[Array, "B L_dec D"], + encoder_hidden_states: Float[Array, "B L_enc D"], + position_ids: Int[Array, "B L_dec"], + merged_attention_mask: Bool[Array, "B 1 L_dec L_combined"] | None = None, + *, + decode: bool | None = None, + ) -> Float[Array, "B L_dec D"]: + """Apply merged self-attention and cross-attention. + + Args: + hidden_states: Decoder hidden states [B, L_dec, D]. + encoder_hidden_states: Encoder outputs [B, L_enc, D]. + position_ids: Decoder position indices [B, L_dec]. + merged_attention_mask: Merged mask [B, 1, L_dec, L_combined]. + decode: Whether to use KV cache for autoregressive decoding. + + Returns: + Output hidden states [B, L_dec, D]. + """ + seq_len = hidden_states.shape[1] + + # Project decoder Q, K, V (self-attention) + query = self.q_proj("BTD,NDH->BTNH", hidden_states) + self_key = self.k_proj("BTD,NDH->BTNH", hidden_states) + self_value = self.v_proj("BTD,NDH->BTNH", hidden_states) + + # QK normalization for self-attention + if self.use_qk_norm: + query = self.q_norm(query) + self_key = self.k_norm(self_key) + + # Apply RoPE to query and self-attention key + query = apply_rope( + query, + position_ids, + base_frequency=self.rope_base_frequency, + scale_factor=self.rope_scale_factor, + ) + self_key = apply_rope( + self_key, + position_ids, + base_frequency=self.rope_base_frequency, + scale_factor=self.rope_scale_factor, + ) + + # Determine decode mode + decode = first_from( + decode, + self.decode, + error_msg="No decode argument provided to T5Gemma2MergedAttention", + ) + + # Handle caching for autoregressive decoding + if decode: + if ( + self.cached_self_key is None + or self.cached_self_value is None + or self.cached_cross_key is None + or self.cached_cross_value is None + or self.cache_index is None + ): + raise ValueError( + "Autoregressive cache not initialized. Call init_cache() first.", + ) + + cur_index = self.cache_index[...] + slice_indices = (0, cur_index, 0, 0) + + # Update self-attention cache using dynamic_update_slice + self_key_updated = jax.lax.dynamic_update_slice( + self.cached_self_key[...], + self_key, + slice_indices, + ) + self_value_updated = jax.lax.dynamic_update_slice( + self.cached_self_value[...], + self_value, + slice_indices, + ) + + self.cached_self_key[...] = self_key_updated + self.cached_self_value[...] = self_value_updated + + # Slice out the valid portion of self-attention cache + # Use dynamic_slice to get [0:cur_index+seq_len] + new_cache_len = cur_index + seq_len + self_key = jax.lax.dynamic_slice( + self_key_updated, + (0, 0, 0, 0), + ( + self_key_updated.shape[0], + new_cache_len, + self.num_kv_heads, + self.head_dim, + ), + ) + self_value = jax.lax.dynamic_slice( + self_value_updated, + (0, 0, 0, 0), + ( + self_value_updated.shape[0], + new_cache_len, + self.num_kv_heads, + self.head_dim, + ), + ) + + # Cross-attention: compute and cache on first call (when index is 0) + # After that, reuse cached values + def compute_cross_kv(): + cross_k = self.k_proj("BTD,NDH->BTNH", encoder_hidden_states) + cross_v = self.v_proj("BTD,NDH->BTNH", encoder_hidden_states) + if self.use_qk_norm: + cross_k = self.k_norm(cross_k) + return cross_k, cross_v + + # Use lax.cond to conditionally compute cross-attention KV + cross_key, cross_value = jax.lax.cond( + cur_index == 0, + compute_cross_kv, + lambda: (self.cached_cross_key[...], self.cached_cross_value[...]), + ) + + # Update cross cache (will be no-op after first call due to cond) + self.cached_cross_key[...] = cross_key + self.cached_cross_value[...] = cross_value + + # Update cache index + self.cache_index[...] = cur_index + seq_len + else: + # No caching - compute cross-attention KV directly + cross_key = self.k_proj("BTD,NDH->BTNH", encoder_hidden_states) + cross_value = self.v_proj("BTD,NDH->BTNH", encoder_hidden_states) + if self.use_qk_norm: + cross_key = self.k_norm(cross_key) + + # Scale query + query = query * self.query_pre_attn_scalar + + # Concatenate self and cross KV + key = jnp.concatenate([self_key, cross_key], axis=1) + value = jnp.concatenate([self_value, cross_value], axis=1) + + # Transpose to [B, N, T, H] for standard attention computation + query = jnp.transpose(query, (0, 2, 1, 3)) # [B, N, T, H] + key = jnp.transpose(key, (0, 2, 1, 3)) # [B, K, S, H] + value = jnp.transpose(value, (0, 2, 1, 3)) # [B, K, S, H] + + # Expand KV heads for GQA + if self.num_kv_groups > 1: + key = jnp.repeat(key, self.num_kv_groups, axis=1) + value = jnp.repeat(value, self.num_kv_groups, axis=1) + + # Compute attention scores: [B, N, T, S] + attn_weights = jnp.einsum("BNTH,BNSH->BNTS", query, key) + + # Apply softcapping + if self.attn_logits_soft_cap is not None: + attn_weights = jnp.tanh(attn_weights / self.attn_logits_soft_cap) + attn_weights = attn_weights * self.attn_logits_soft_cap + + # Apply attention mask + if merged_attention_mask is not None: + attn_weights = jnp.where(merged_attention_mask, attn_weights, K_MASK) + + # Softmax + attn_weights = jax.nn.softmax(attn_weights, axis=-1).astype(value.dtype) + + # Apply attention to values: [B, N, T, H] + attn_output = jnp.einsum("BNTS,BNSH->BNTH", attn_weights, value) + + # Transpose back to [B, T, N, H] + attn_output = jnp.transpose(attn_output, (0, 2, 1, 3)) + + # Output projection + output = self.o_proj("BTNH,NHD->BTD", attn_output) + + return output + + +class T5Gemma2DecoderBlock(nnx.Module): + """Decoder transformer block with merged self/cross attention.""" + + def __init__( + self, + num_q_heads: int, + num_kv_heads: int, + embed_dim: int, + head_dim: int, + hidden_dim: int, + *, + attn_type: AttentionType = AttentionType.GLOBAL, + use_post_attn_norm: bool = True, + use_post_ffw_norm: bool = True, + query_pre_attn_scalar: float, + rope_base_frequency: int = 10_000, + rope_scale_factor: float = 1.0, + attn_logits_soft_cap: float | None = None, + use_qk_norm: bool = True, + kernel_init: nnx.Initializer = NORMAL_INIT, + scale_init: nnx.Initializer = ZEROS_INIT, + decode: bool | None = None, + dtype: jnp.dtype = jnp.float32, + rngs: nnx.Rngs, + ): + self.attn_type = attn_type + self.use_post_attn_norm = use_post_attn_norm + self.use_post_ffw_norm = use_post_ffw_norm + + # Pre-attention norm + self.pre_attention_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + # Merged attention (self + cross) + self.attn = T5Gemma2MergedAttention( + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + features=embed_dim, + head_dim=head_dim, + query_pre_attn_scalar=query_pre_attn_scalar, + rope_base_frequency=rope_base_frequency, + rope_scale_factor=rope_scale_factor, + attn_logits_soft_cap=attn_logits_soft_cap, + use_qk_norm=use_qk_norm, + kernel_init=kernel_init, + scale_init=scale_init, + decode=decode, + dtype=dtype, + rngs=rngs, + ) + + # Post-attention norm + if use_post_attn_norm: + self.post_attention_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + # Pre-FFW norm + self.pre_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + # Feed-forward + self.mlp = T5Gemma2FeedForward( + features=embed_dim, + hidden_dim=hidden_dim, + transpose_gating_einsum=True, + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + # Post-FFW norm + if use_post_ffw_norm: + self.post_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + def __call__( + self, + x: Float[Array, "B L_dec D"], + encoder_hidden_states: Float[Array, "B L_enc D"], + segment_pos: Int[Array, "B L_dec"], + attn_mask: Bool[Array, "B 1 L_dec L_combined"] | None = None, + *, + decode: bool | None = None, + ) -> Float[Array, "B L_dec D"]: + """Apply decoder layer. + + Args: + x: Decoder hidden states [B, L_dec, D]. + encoder_hidden_states: Encoder outputs [B, L_enc, D]. + segment_pos: Decoder position indices [B, L_dec]. + attn_mask: Merged attention mask [B, 1, L_dec, L_combined]. + decode: Whether to use KV cache. + + Returns: + Output hidden states [B, L_dec, D]. + """ + # Merged attention block + inputs_normalized = self.pre_attention_norm(x) + attn_output = self.attn( + inputs_normalized, + encoder_hidden_states, + segment_pos, + attn_mask, + decode=decode, + ) + + if self.use_post_attn_norm: + attn_output = self.post_attention_norm(attn_output) + + attn_output += x + + # Feed-forward block + outputs = self.pre_ffw_norm(attn_output) + outputs = self.mlp(outputs) + + if self.use_post_ffw_norm: + outputs = self.post_ffw_norm(outputs) + + outputs += attn_output + + return outputs + + +class T5Gemma2Decoder(nnx.Module): + """T5Gemma2 Decoder with merged self/cross attention. + + Uses nnx.Cache for autoregressive decoding. Call init_cache() before + using decode=True. + """ + + def __init__( + self, + config: T5Gemma2DecoderConfig, + *, + embedder: nnx.Module | None = None, + dtype: jnp.dtype = jnp.float32, + decode: bool | None = None, + rngs: nnx.Rngs, + ): + self.config = config + self.dtype = dtype + + # Text embeddings + self.embedder = embedder or T5Gemma2ScaledWordEmbedding( + vocab_size=config.vocab_size, + embed_dim=config.embed_dim, + padding_idx=config.pad_token_id, + dtype=dtype, + rngs=rngs, + ) + + # Determine attention types per layer + attention_types = config.layer_types + if not attention_types: + attention_types = tuple(AttentionType.GLOBAL for _ in range(config.num_hidden_layers)) + + # Decoder layers + self.blocks = nnx.List( + [ + T5Gemma2DecoderBlock( + num_q_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + embed_dim=config.embed_dim, + head_dim=config.head_dim, + hidden_dim=config.intermediate_size, + attn_type=attention_types[i], + use_post_attn_norm=True, + use_post_ffw_norm=True, + query_pre_attn_scalar=config.query_pre_attn_scalar, + rope_base_frequency=_get_rope_base_frequency(config, attention_types[i]), + rope_scale_factor=_get_rope_scale_factor(config, attention_types[i]), + attn_logits_soft_cap=config.attn_logit_softcapping, + use_qk_norm=True, + decode=decode, + dtype=dtype, + rngs=rngs, + ) + for i in range(config.num_hidden_layers) + ] + ) + + # Final normalization + self.norm = T5Gemma2RMSNorm(config.embed_dim, dtype=dtype, rngs=rngs) + + def init_cache( + self, + *, + batch_size: int, + max_decode_length: int, + encoder_seq_length: int, + dtype: jnp.dtype | None = None, + ) -> None: + """Initialize KV caches for all decoder layers. + + Must be called before using decode=True. + + Args: + batch_size: Number of sequences in the batch. + max_decode_length: Maximum decoder sequence length. + encoder_seq_length: Encoder sequence length (for cross-attention). + dtype: Data type for cache arrays. Defaults to model dtype. + """ + if dtype is None: + dtype = self.dtype + + for block in self.blocks: + block.attn.init_cache( + batch_size=batch_size, + max_decode_length=max_decode_length, + encoder_seq_length=encoder_seq_length, + dtype=dtype, + ) + + def get_cache_index(self) -> int: + """Get current cache index from the first layer.""" + if len(self.blocks) > 0 and self.blocks[0].attn.cache_index is not None: + return int(self.blocks[0].attn.cache_index[...]) + return 0 + + def __call__( + self, + input_ids: Int[Array, "B L_dec"], + encoder_hidden_states: Float[Array, "B L_enc D"], + attention_mask: Bool[Array, "B L_dec"] | None = None, + encoder_attention_mask: Bool[Array, "B L_enc"] | None = None, + position_ids: Int[Array, "B L_dec"] | None = None, + *, + decode: bool | None = None, + ) -> Float[Array, "B L_dec D"]: + """Forward pass of the decoder. + + Args: + input_ids: Decoder input token IDs [B, L_dec]. + encoder_hidden_states: Encoder outputs [B, L_enc, D]. + attention_mask: Decoder attention mask [B, L_dec]. + encoder_attention_mask: Encoder attention mask [B, L_enc]. + position_ids: Decoder position indices [B, L_dec]. + decode: Whether to use KV cache. + + Returns: + Decoder hidden states [B, L_dec, D]. + """ + batch_size, dec_seq_len = input_ids.shape + enc_seq_len = encoder_hidden_states.shape[1] + + # Create position IDs if not provided + if position_ids is None: + cache_index = self.get_cache_index() + if decode and cache_index > 0: + # During autoregressive decoding, position is cache_index + position_ids = jnp.full( + (batch_size, dec_seq_len), + cache_index, + dtype=jnp.int32, + ) + else: + position_ids = jnp.arange(dec_seq_len)[None, :].repeat( + batch_size, + axis=0, + ) + + # Create attention masks if not provided + if attention_mask is None: + attention_mask = jnp.ones((batch_size, dec_seq_len), dtype=jnp.bool_) + if encoder_attention_mask is None: + encoder_attention_mask = jnp.ones( + (batch_size, enc_seq_len), + dtype=jnp.bool_, + ) + + # Create encoder mask for cross-attention (bidirectional) + encoder_mask = encoder_attention_mask[:, None, None, :] + + if decode: + # Decode mode: create masks accounting for cached sequence length + cache_index = self.get_cache_index() + cache_len = cache_index + dec_seq_len # Total self-attention key length + + # Full attention: new tokens can attend to all cached + encoder + full_decoder_mask = make_decode_mode_self_mask(batch_size, dec_seq_len, cache_len) + # Sliding attention: limited to sliding window + sliding_decoder_mask = make_decode_mode_sliding_mask( + batch_size, + cache_index, # Current query position + cache_len, + self.config.sliding_window, + ) + + # Merged masks with encoder (broadcast encoder mask for query length) + enc_mask_broadcast = jnp.broadcast_to(encoder_mask, (batch_size, 1, dec_seq_len, enc_seq_len)) + merged_mask_full = jnp.concatenate([full_decoder_mask, enc_mask_broadcast], axis=-1) + merged_mask_sliding = jnp.concatenate([sliding_decoder_mask, enc_mask_broadcast], axis=-1) + else: + # Non-decode mode: standard causal masks + # Full attention: standard causal mask + full_decoder_mask = make_causal_mask(attention_mask) + # Sliding attention: causal + sliding window + sliding_decoder_mask = make_sliding_window_causal_mask( + attention_mask, + position_ids, + self.config.sliding_window, + ) + + # Create merged masks for each attention type + merged_mask_full = make_merged_attention_mask(full_decoder_mask, encoder_mask) + merged_mask_sliding = make_merged_attention_mask(sliding_decoder_mask, encoder_mask) + + # Embed tokens + hidden_states = self.embedder(input_ids) + + # Apply decoder layers with per-layer masks + for block in self.blocks: + # Select mask based on layer's attention type + if block.attn_type == AttentionType.LOCAL_SLIDING: + block_mask = merged_mask_sliding + else: + block_mask = merged_mask_full + + hidden_states = block( + hidden_states, + encoder_hidden_states, + position_ids, + block_mask, + decode=decode, + ) + + # Final normalization + hidden_states = self.norm(hidden_states) + + return hidden_states + + +# ============================================================================= +# Full Model +# ============================================================================= + + +class T5Gemma2(nnx.Module): + """T5Gemma2 Encoder-Decoder Model. + + Combines the encoder and decoder into a full sequence-to-sequence model. + Uses nnx.Cache for autoregressive decoding. Call init_cache() before + using decode=True. + """ + + def __init__( + self, + config: T5Gemma2Config, + *, + dtype: jnp.dtype = jnp.float32, + dtype_mm: jnp.dtype = jnp.float32, + decode: bool | None = None, + rngs: nnx.Rngs, + ): + self.config = config + + # Use tied embedding for encoder and decoder + self.embedder = T5Gemma2ScaledWordEmbedding( + vocab_size=config.encoder.text_config.vocab_size, + embed_dim=config.encoder.text_config.embed_dim, + padding_idx=config.encoder.text_config.pad_token_id, + dtype=dtype, + rngs=rngs, + ) + + # Encoder + self.encoder = T5Gemma2Encoder( + config.encoder, + embedder=self.embedder, + dtype=dtype, + dtype_mm=dtype_mm, + rngs=rngs, + ) + + # Decoder + self.decoder = T5Gemma2Decoder( + config.decoder, + embedder=self.embedder, + dtype=dtype, + decode=decode, + rngs=rngs, + ) + + def init_cache( + self, + *, + batch_size: int, + max_decode_length: int, + encoder_seq_length: int, + dtype: jnp.dtype | None = None, + ) -> None: + """Initialize decoder KV caches for autoregressive decoding. + + Args: + batch_size: Number of sequences in the batch. + max_decode_length: Maximum decoder sequence length. + encoder_seq_length: Encoder sequence length. + dtype: Data type for cache arrays. + """ + self.decoder.init_cache( + batch_size=batch_size, + max_decode_length=max_decode_length, + encoder_seq_length=encoder_seq_length, + dtype=dtype, + ) + + def __call__( + self, + input_ids: Int[Array, "B L_enc"], + decoder_input_ids: Int[Array, "B L_dec"], + attention_mask: Bool[Array, "B L_enc"] | None = None, + decoder_attention_mask: Bool[Array, "B L_dec"] | None = None, + position_ids: Int[Array, "B L_enc"] | None = None, + decoder_position_ids: Int[Array, "B L_dec"] | None = None, + pixel_values: Float[Array, "B N H W C"] | None = None, + encoder_outputs: Float[Array, "B L_enc D"] | None = None, + *, + decode: bool | None = None, + deterministic: bool = True, + ) -> tuple[Float[Array, "B L_dec D"], Float[Array, "B L_enc D"]]: + """Forward pass of the full model. + + Args: + input_ids: Encoder input token IDs [B, L_enc]. + decoder_input_ids: Decoder input token IDs [B, L_dec]. + attention_mask: Encoder attention mask [B, L_enc]. + decoder_attention_mask: Decoder attention mask [B, L_dec]. + position_ids: Encoder position indices [B, L_enc]. + decoder_position_ids: Decoder position indices [B, L_dec]. + pixel_values: Optional input images [B, N, H, W, C] where N is images per batch. + encoder_outputs: Pre-computed encoder outputs (for generation). + decode: Whether to use KV cache. + deterministic: Whether to run in deterministic mode. + + Returns: + Tuple of (decoder hidden states, encoder hidden states). + """ + # Encode if not provided + if encoder_outputs is None: + encoder_outputs = self.encoder( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + images=pixel_values, # Note: encoder uses 'images' parameter + deterministic=deterministic, + ) + + # Decode + decoder_outputs = self.decoder( + decoder_input_ids, + encoder_hidden_states=encoder_outputs, + attention_mask=decoder_attention_mask, + encoder_attention_mask=attention_mask, + position_ids=decoder_position_ids, + decode=decode, + ) + + return decoder_outputs, encoder_outputs From 02331b3934f7ca441621723ab0e06c42b073caa7 Mon Sep 17 00:00:00 2001 From: Jaonary Rabarisoa Date: Tue, 27 Jan 2026 15:15:05 +0100 Subject: [PATCH 3/5] Refactor T5Gemma2 model components and add parameter handling and test code --- bonsai/models/t5gemma2/modeling.py | 160 ++----- bonsai/models/t5gemma2/params.py | 551 ++++++++++++++++++++++ bonsai/models/t5gemma2/tests/run_model.py | 517 ++++++++++++++++++++ 3 files changed, 1106 insertions(+), 122 deletions(-) create mode 100644 bonsai/models/t5gemma2/params.py create mode 100644 bonsai/models/t5gemma2/tests/run_model.py diff --git a/bonsai/models/t5gemma2/modeling.py b/bonsai/models/t5gemma2/modeling.py index 0bc278d4..b504951e 100644 --- a/bonsai/models/t5gemma2/modeling.py +++ b/bonsai/models/t5gemma2/modeling.py @@ -365,14 +365,14 @@ def __init__( else: gating_shape = (2, features, hidden_dim) - self.gating_einsum = Gemma3Einsum( + self.gating_einsum = T5Gemma2Einsum( shape=gating_shape, kernel_init=kernel_init, dtype=dtype, rngs=rngs, ) - self.linear = Gemma3Einsum( + self.linear = T5Gemma2Einsum( shape=(hidden_dim, features), kernel_init=kernel_init, dtype=dtype, @@ -530,29 +530,32 @@ def make_merged_attention_mask( def make_decode_mode_self_mask( batch_size: int, query_len: int, - cache_len: int, + max_cache_len: int, + current_cache_len: int, ) -> Bool[Array, "B 1 Q K"]: """Creates self-attention mask for decode mode (with KV cache). During decode mode, the query attends to all previously cached keys. - Since we're generating tokens sequentially and causality is already - enforced by the order of generation, all cached positions are valid. + We mask out positions that haven't been written to yet (>= current_cache_len). Args: batch_size: Batch size. query_len: Number of query positions (typically 1 for decode mode). - cache_len: Total length of cached keys (cache_index + query_len). + max_cache_len: Total allocated size of the cache. + current_cache_len: Current valid length of the cache. Returns: - Mask of shape [B, 1, query_len, cache_len] with all True. + Mask of shape [B, 1, query_len, max_cache_len]. """ - return jnp.ones((batch_size, 1, query_len, cache_len), dtype=jnp.bool_) + k_pos = jnp.arange(max_cache_len) + mask = k_pos < current_cache_len + return jnp.broadcast_to(mask[None, None, None, :], (batch_size, 1, query_len, max_cache_len)) def make_decode_mode_sliding_mask( batch_size: int, query_pos: int, - cache_len: int, + max_cache_len: int, sliding_window: int, ) -> Bool[Array, "B 1 1 K"]: """Creates sliding window mask for decode mode (with KV cache). @@ -563,20 +566,20 @@ def make_decode_mode_sliding_mask( Args: batch_size: Batch size. query_pos: Position of the current query (typically cache_index). - cache_len: Total length of cached keys. + max_cache_len: Total allocated size of the cache. sliding_window: Size of the sliding window. Returns: - Mask of shape [B, 1, 1, cache_len] with True for valid positions. + Mask of shape [B, 1, 1, max_cache_len]. """ - # Key positions: 0, 1, 2, ..., cache_len-1 - k_positions = jnp.arange(cache_len) + # Key positions: 0, 1, 2, ..., max_cache_len-1 + k_positions = jnp.arange(max_cache_len) # Within sliding window: |query_pos - k_pos| < sliding_window in_window = jnp.abs(query_pos - k_positions) < sliding_window # Also must be causal (k_pos <= query_pos) is_causal = k_positions <= query_pos mask = in_window & is_causal - return jnp.broadcast_to(mask[None, None, None, :], (batch_size, 1, 1, cache_len)) + return jnp.broadcast_to(mask[None, None, None, :], (batch_size, 1, 1, max_cache_len)) # ============================================================================= @@ -695,7 +698,7 @@ def __call__( deterministic: bool | None = None, ) -> Float[Array, "B L D"]: x = self.dense1(x) - x = nnx.gelu(x) + x = jax.nn.gelu(x) x = self.dropout(x, deterministic=deterministic) return self.dense2(x) @@ -741,7 +744,7 @@ def __init__( rngs=rngs, ) - self.mlp = Gemma3VisionMLP( + self.mlp = T5Gemma2VisionMLP( width=width, mlp_dim=mlp_dim, dropout=dropout, @@ -836,7 +839,7 @@ def __call__(self, x: jax.Array) -> jax.Array: assert not cur_width % output_width, f"{cur_width=} {output_width=}" window = cur_width // output_width window_shape = (window, window) - x = nnx.avg_pool(x, window_shape=window_shape, strides=window_shape) + x = nnx.avg_pool(x, window_shape, strides=window_shape, padding="VALID") batch_size, height, width, embed_dim = x.shape return jnp.reshape(x, (batch_size, height * width, embed_dim)) @@ -850,7 +853,7 @@ class T5Gemma2VisionSoftTokenizer(nnx.Module): def __init__( self, - config: VisionEncoderConfig, + config: T5Gemma2VisionConfig, *, dtype: jnp.dtype = jnp.float32, rngs: nnx.Rngs, @@ -882,7 +885,7 @@ def __init__( self.dropout = nnx.Dropout(rate=config.dropout, rngs=rngs) - self.transformer = Gemma3VisionEncoder( + self.transformer = T5Gemma2VisionEncoder( width=config.width, depth=config.depth, mlp_dim=config.mlp_dim, @@ -892,7 +895,7 @@ def __init__( rngs=rngs, ) - self.vision_exit = Gemma3VisionExit(output_length=256, rngs=rngs) + self.vision_exit = T5Gemma2VisionExit(output_length=256, rngs=rngs) def __call__( self, @@ -972,7 +975,7 @@ class T5Gemma2VisionEmbedder(nnx.Module): def __init__( self, *, - vision_config: VisionEncoderConfig, + vision_config: T5Gemma2VisionConfig, embed_dim: int, freeze_params: bool = True, dtype: jnp.dtype = jnp.float32, @@ -1010,81 +1013,6 @@ def __call__( # ============================================================================= # Multimodal Utilities # ============================================================================= -def _get_new_text_positions(*, offset_on: np.ndarray, offset_by: int) -> np.ndarray: - """Create the positions of the new tokens.""" - offset = np.cumsum(offset_on, axis=-1) * offset_by - new_positions = np.arange(offset_on.shape[-1]) + offset - new_positions -= offset_by * offset_on - return new_positions - - -def _insert_sequence( - tokens: np.ndarray, - *, - at: int, - sequence: list[int], - max_num_images: int, -) -> np.ndarray: - """Inserts a sequence of tokens at all occurrences of a specific token.""" - original_dim = tokens.ndim - if original_dim == 1: - tokens = tokens[None, :] - - batch_size, length = tokens.shape - mm_tokens_to_insert = np.array(sequence) - offset_by = len(mm_tokens_to_insert) - 1 - length_with_mm = length + max_num_images * offset_by - mm_start = tokens == at - - new_tokens = np.zeros((batch_size, length_with_mm), dtype=np.int64) - new_text_pos = _get_new_text_positions(offset_on=mm_start, offset_by=offset_by) - np.put_along_axis(new_tokens, new_text_pos, tokens, axis=1) - - batch_indices_to_zero, _ = np.where(mm_start) - new_pos_to_zero = new_text_pos[mm_start] - if batch_indices_to_zero.size > 0: - new_tokens[batch_indices_to_zero, new_pos_to_zero] = 0 - - batch_indices, seq_indices = np.nonzero(mm_start) - - if batch_indices.size > 0: - intra_batch_img_idx = np.cumsum(mm_start, axis=1)[mm_start] - 1 - final_img_start_pos = seq_indices + intra_batch_img_idx * offset_by - indices_to_insert = final_img_start_pos[:, None] + np.arange( - len(mm_tokens_to_insert), - ) - new_tokens[batch_indices[:, None], indices_to_insert] = mm_tokens_to_insert - - if original_dim == 1: - new_tokens = np.squeeze(new_tokens) - return new_tokens - - -def add_extra_tokens_for_images( - tokens: np.ndarray | list, - *, - max_num_images: int = 1, -) -> np.ndarray: - """Add extra image tokens to text tokens.""" - - mm_tokens = [ - NEW_LINE_TOKEN, - START_OF_IMAGE_TOKEN, - *[IMAGE_PLACEHOLDER_TOKEN] * NUM_PLACEHOLDER_TOKENS_PER_IMAGE, - END_OF_IMAGE_TOKEN, - NEW_LINE_TOKEN, - ] - if not isinstance(tokens, np.ndarray): - tokens = np.asarray(tokens) - - return _insert_sequence( - at=START_OF_IMAGE_TOKEN, - sequence=mm_tokens, - tokens=tokens, - max_num_images=max_num_images, - ) - - def merge_mm_embeddings( text_embeddings: jnp.ndarray, vision_embeddings: jnp.ndarray, @@ -1802,29 +1730,9 @@ def __call__( self.cached_self_key[...] = self_key_updated self.cached_self_value[...] = self_value_updated - # Slice out the valid portion of self-attention cache - # Use dynamic_slice to get [0:cur_index+seq_len] - new_cache_len = cur_index + seq_len - self_key = jax.lax.dynamic_slice( - self_key_updated, - (0, 0, 0, 0), - ( - self_key_updated.shape[0], - new_cache_len, - self.num_kv_heads, - self.head_dim, - ), - ) - self_value = jax.lax.dynamic_slice( - self_value_updated, - (0, 0, 0, 0), - ( - self_value_updated.shape[0], - new_cache_len, - self.num_kv_heads, - self.head_dim, - ), - ) + # Use the full updated cache (static shape) for JIT compatibility + self_key = self_key_updated + self_value = self_value_updated # Cross-attention: compute and cache on first call (when index is 0) # After that, reuse cached values @@ -2191,15 +2099,23 @@ def __call__( if decode: # Decode mode: create masks accounting for cached sequence length cache_index = self.get_cache_index() - cache_len = cache_index + dec_seq_len # Total self-attention key length + # Get max_decode_length (static) from the first block's cache + max_decode_len = self.blocks[0].attn.cached_self_key.shape[1] + # Current valid length (dynamic) + current_len = cache_index + dec_seq_len # Full attention: new tokens can attend to all cached + encoder - full_decoder_mask = make_decode_mode_self_mask(batch_size, dec_seq_len, cache_len) + full_decoder_mask = make_decode_mode_self_mask( + batch_size, + dec_seq_len, + max_decode_len, + current_len, + ) # Sliding attention: limited to sliding window sliding_decoder_mask = make_decode_mode_sliding_mask( batch_size, cache_index, # Current query position - cache_len, + max_decode_len, self.config.sliding_window, ) diff --git a/bonsai/models/t5gemma2/params.py b/bonsai/models/t5gemma2/params.py new file mode 100644 index 00000000..b41295dc --- /dev/null +++ b/bonsai/models/t5gemma2/params.py @@ -0,0 +1,551 @@ +# Copyright 2025 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import re +from enum import Enum + +import jax +import safetensors +from etils import epath +from flax import nnx + +from bonsai.models.t5gemma2 import modeling as model_lib +from bonsai.models.t5gemma2.modeling import T5Gemma2Config + + +def _get_key_and_transform_mapping(cfg: T5Gemma2Config): + """Returns mapping from checkpoint keys to model keys with transformations.""" + enc = cfg.encoder.text_config + dec = cfg.decoder + vis = cfg.encoder.vision_config + + class Transform(Enum): + """Transformations for model parameters. + + Format: (permute_axes, reshape_shape, reshape_first) + - If reshape_first=True: reshape then permute + - If reshape_first=False: permute then reshape + """ + + NONE = None + SCALE = None + EMBED = None + + # Encoder attention transforms + # q_proj: (num_q_heads * head_dim, embed_dim) -> (num_q_heads, embed_dim, head_dim) + # reshape to (num_q_heads, head_dim, embed_dim), permute (0, 2, 1) + ENC_ATTN_Q = ( + (0, 2, 1), + (enc.num_attention_heads, enc.head_dim, enc.embed_dim), + True, + ) + # k/v_proj: (num_kv_heads * head_dim, embed_dim) -> (num_kv_heads, embed_dim, head_dim) + ENC_ATTN_KV = ( + (0, 2, 1), + (enc.num_key_value_heads, enc.head_dim, enc.embed_dim), + True, + ) + # o_proj: (embed_dim, num_q_heads * head_dim) -> (num_q_heads, head_dim, embed_dim) + # reshape to (embed_dim, num_q_heads, head_dim), permute (1, 2, 0) + ENC_ATTN_OUT = ( + (1, 2, 0), + (enc.embed_dim, enc.num_attention_heads, enc.head_dim), + True, + ) + + # Decoder attention transforms (same structure as encoder, separate k/v) + DEC_ATTN_Q = ( + (0, 2, 1), + (dec.num_attention_heads, dec.head_dim, dec.embed_dim), + True, + ) + DEC_ATTN_KV = ( + (0, 2, 1), + (dec.num_key_value_heads, dec.head_dim, dec.embed_dim), + True, + ) + DEC_ATTN_OUT = ( + (1, 2, 0), + (dec.embed_dim, dec.num_attention_heads, dec.head_dim), + True, + ) + + # MLP transforms - gating uses no transform (transpose_gating_einsum=True) + MLP_GATE = None # checkpoint is (hidden, embed), model is (2, hidden, embed) after stack + MLP_DOWN = ((1, 0), None, False) # transpose + + # Vision transforms + VIT_LINEAR = ((1, 0), None, False) + VIT_BIAS = None + VIT_CONV = ((2, 3, 1, 0), None, False) # OIHW -> HWIO + VIT_POS_EMBED = None # Will handle separately + + # Maping of checkpoint_keys -> (model_keys, Transform) + return { + r"model\.encoder\.embed_tokens\.weight": ( + "decoder.embedder.embedding", + Transform.EMBED, + ), + r"model\.encoder\.embed_tokens\.eoi_embedding": ( + "decoder.embedder.eoi_embedding", + Transform.EMBED, + ), + # === Encoder Layers === + r"model\.encoder\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": ( + r"encoder.blocks.\1.attn.q_einsum.w", + Transform.ENC_ATTN_Q, + ), + r"model\.encoder\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": ( + r"encoder.blocks.\1.attn.kv_einsum.w.__stacked_k", + Transform.ENC_ATTN_KV, + ), + r"model\.encoder\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": ( + r"encoder.blocks.\1.attn.kv_einsum.w.__stacked_v", + Transform.ENC_ATTN_KV, + ), + r"model\.encoder\.layers\.([0-9]+)\.self_attn\.o_proj\.weight": ( + r"encoder.blocks.\1.attn.attn_vec_einsum.w", + Transform.ENC_ATTN_OUT, + ), + r"model\.encoder\.layers\.([0-9]+)\.self_attn\.q_norm\.weight": ( + r"encoder.blocks.\1.attn._query_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.layers\.([0-9]+)\.self_attn\.k_norm\.weight": ( + r"encoder.blocks.\1.attn._key_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.layers\.([0-9]+)\.mlp\.gate_proj\.weight": ( + r"encoder.blocks.\1.mlp.gating_einsum.w.__stacked_gate", + Transform.MLP_GATE, + ), + r"model\.encoder\.layers\.([0-9]+)\.mlp\.up_proj\.weight": ( + r"encoder.blocks.\1.mlp.gating_einsum.w.__stacked_up", + Transform.MLP_GATE, + ), + r"model\.encoder\.layers\.([0-9]+)\.mlp\.down_proj\.weight": ( + r"encoder.blocks.\1.mlp.linear.w", + Transform.MLP_DOWN, + ), + r"model\.encoder\.layers\.([0-9]+)\.pre_self_attn_layernorm\.weight": ( + r"encoder.blocks.\1.pre_attention_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.layers\.([0-9]+)\.post_self_attn_layernorm\.weight": ( + r"encoder.blocks.\1.post_attention_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.layers\.([0-9]+)\.pre_feedforward_layernorm\.weight": ( + r"encoder.blocks.\1.pre_ffw_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.layers\.([0-9]+)\.post_feedforward_layernorm\.weight": ( + r"encoder.blocks.\1.post_ffw_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.norm\.weight": ("encoder.norm.scale", Transform.SCALE), + # === Decoder Layers (uses separate einsum modules for k, v) === + # Note: Decoder embeddings are tied with encoder (shared embedder) + r"model\.decoder\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": ( + r"decoder.blocks.\1.attn.q_proj.w", + Transform.DEC_ATTN_Q, + ), + r"model\.decoder\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": ( + r"decoder.blocks.\1.attn.k_proj.w", + Transform.DEC_ATTN_KV, + ), + r"model\.decoder\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": ( + r"decoder.blocks.\1.attn.v_proj.w", + Transform.DEC_ATTN_KV, + ), + r"model\.decoder\.layers\.([0-9]+)\.self_attn\.o_proj\.weight": ( + r"decoder.blocks.\1.attn.o_proj.w", + Transform.DEC_ATTN_OUT, + ), + r"model\.decoder\.layers\.([0-9]+)\.self_attn\.q_norm\.weight": ( + r"decoder.blocks.\1.attn.q_norm.scale", + Transform.SCALE, + ), + r"model\.decoder\.layers\.([0-9]+)\.self_attn\.k_norm\.weight": ( + r"decoder.blocks.\1.attn.k_norm.scale", + Transform.SCALE, + ), + r"model\.decoder\.layers\.([0-9]+)\.mlp\.gate_proj\.weight": ( + r"decoder.blocks.\1.mlp.gating_einsum.w.__stacked_gate", + Transform.MLP_GATE, + ), + r"model\.decoder\.layers\.([0-9]+)\.mlp\.up_proj\.weight": ( + r"decoder.blocks.\1.mlp.gating_einsum.w.__stacked_up", + Transform.MLP_GATE, + ), + r"model\.decoder\.layers\.([0-9]+)\.mlp\.down_proj\.weight": ( + r"decoder.blocks.\1.mlp.linear.w", + Transform.MLP_DOWN, + ), + r"model\.decoder\.layers\.([0-9]+)\.pre_self_attn_layernorm\.weight": ( + r"decoder.blocks.\1.pre_attention_norm.scale", + Transform.SCALE, + ), + r"model\.decoder\.layers\.([0-9]+)\.post_self_attn_layernorm\.weight": ( + r"decoder.blocks.\1.post_attention_norm.scale", + Transform.SCALE, + ), + r"model\.decoder\.layers\.([0-9]+)\.pre_feedforward_layernorm\.weight": ( + r"decoder.blocks.\1.pre_ffw_norm.scale", + Transform.SCALE, + ), + r"model\.decoder\.layers\.([0-9]+)\.post_feedforward_layernorm\.weight": ( + r"decoder.blocks.\1.post_ffw_norm.scale", + Transform.SCALE, + ), + r"model\.decoder\.norm\.weight": ("decoder.norm.scale", Transform.SCALE), + # === Vision Tower === + r"model\.encoder\.vision_tower\.vision_model\.embeddings\.patch_embedding\.weight": ( + "encoder.vision_embedder.soft_tokenizer.embedding.kernel", + Transform.VIT_CONV, + ), + r"model\.encoder\.vision_tower\.vision_model\.embeddings\.patch_embedding\.bias": ( + "encoder.vision_embedder.soft_tokenizer.embedding.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.embeddings\.position_embedding\.weight": ( + "encoder.vision_embedder.soft_tokenizer.pos_embedding", + Transform.VIT_POS_EMBED, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm1\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.layer_norm1.scale", + Transform.SCALE, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm1\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.layer_norm1.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm2\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.layer_norm2.scale", + Transform.SCALE, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm2\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.layer_norm2.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.query.kernel", + Transform.VIT_LINEAR, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.q_proj\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.query.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.key.kernel", + Transform.VIT_LINEAR, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.k_proj\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.key.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.value.kernel", + Transform.VIT_LINEAR, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.v_proj\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.value.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.out_proj\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.out.kernel", + Transform.VIT_LINEAR, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.out_proj\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mha.out.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc1\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mlp.dense1.kernel", + Transform.VIT_LINEAR, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc1\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mlp.dense1.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc2\.weight": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mlp.dense2.kernel", + Transform.VIT_LINEAR, + ), + r"model\.encoder\.vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc2\.bias": ( + r"encoder.vision_embedder.soft_tokenizer.transformer.blocks.\1.mlp.dense2.bias", + Transform.VIT_BIAS, + ), + r"model\.encoder\.vision_tower\.vision_model\.post_layernorm\.weight": ( + "encoder.vision_embedder.soft_tokenizer.transformer.encoder_norm.scale", + Transform.SCALE, + ), + r"model\.encoder\.vision_tower\.vision_model\.post_layernorm\.bias": ( + "encoder.vision_embedder.soft_tokenizer.transformer.encoder_norm.bias", + Transform.VIT_BIAS, + ), + # === Multi-modal Projector === + r"model\.encoder\.multi_modal_projector\.mm_input_projection_weight": ( + "encoder.vision_embedder.soft_tokens_embedder.mm_input_projection.w", + Transform.NONE, + ), + r"model\.encoder\.multi_modal_projector\.mm_soft_emb_norm\.weight": ( + "encoder.vision_embedder.soft_tokens_embedder.mm_soft_embedding_norm.scale", + Transform.SCALE, + ), + }, vis + + +def _torch_key_to_jax_key(mapping, source_key): + """Convert checkpoint key to model key using regex substitution.""" + subs = [ + (re.sub(pat, repl, source_key), transform) + for pat, (repl, transform) in mapping.items() + if re.match(pat, source_key) + ] + if len(subs) != 1: + if len(subs) == 0: + return None, None + raise ValueError(f"Multiple matches for key: {source_key}") + return subs[0] + + +def _assign_weights(keys, tensor, state_dict, st_key, transform, sharding_dict): + """Recursively descend into state_dict and assign the (possibly transformed) tensor.""" + key, *rest = keys + if not rest: + if transform is not None: + permute, reshape, reshape_first = transform + if reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + if permute: + tensor = tensor.transpose(permute) + if not reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + if tensor.shape != state_dict[key].shape: + raise ValueError( + f"Shape mismatch for {st_key}: {tensor.shape} vs {state_dict[key].shape}" + ) + if sharding_dict is not None: + state_dict[key] = jax.device_put(tensor, sharding_dict[key]) + else: + state_dict[key] = jax.device_put(tensor) + else: + next_sharding = sharding_dict[key] if sharding_dict is not None else None + _assign_weights(rest, tensor, state_dict[key], st_key, transform, next_sharding) + + +def _stoi(s): + """String to int if possible.""" + try: + return int(s) + except ValueError: + return s + + +def create_model_from_safe_tensors( + file_dir: str, cfg: T5Gemma2Config, mesh: jax.sharding.Mesh | None = None +) -> model_lib.T5Gemma2: + """Load tensors from the safetensors file and create a T5Gemma2 model.""" + files = list(epath.Path(file_dir).expanduser().glob("*.safetensors")) + if not files: + raise ValueError(f"No safetensors found in {file_dir}") + + # Create actual model (not eval_shape) to preserve non-parameter arrays + model = model_lib.T5Gemma2(cfg, rngs=nnx.Rngs(params=0, dropout=0)) + graph_def, state = nnx.split(model) + state_dict = state.to_pure_dict() + sharding = ( + nnx.get_named_sharding(state, mesh).to_pure_dict() if mesh is not None else None + ) + + key_mapping, vis_cfg = _get_key_and_transform_mapping(cfg) + conversion_errors = [] + + # Buffer for stacked parameters (k+v for encoder, gate+up for mlp) + stacked_buffers = {} + + # Vision config for reshaping + vit_num_heads = vis_cfg.num_heads if vis_cfg else 16 + vit_head_dim = vis_cfg.width // vis_cfg.num_heads if vis_cfg else 72 + + for f in files: + with safetensors.safe_open(f, framework="numpy") as sf: + for torch_key in sf.keys(): + tensor = sf.get_tensor(torch_key) + + jax_key, transform = _torch_key_to_jax_key(key_mapping, torch_key) + if jax_key is None: + conversion_errors.append(f"No mapping found for: {torch_key}") + continue + + # Skip vision weights if vision is not enabled + if "vision_embedder" in jax_key and cfg.encoder.vision_config is None: + continue + + # Handle stacked parameters (use __stacked_ prefix to avoid conflicts) + if ".__stacked_k" in jax_key or ".__stacked_v" in jax_key: + if ".__stacked_k" in jax_key: + base_key = jax_key.replace(".__stacked_k", "") + idx = 0 + else: + base_key = jax_key.replace(".__stacked_v", "") + idx = 1 + + # Apply transform + if transform.value is not None: + permute, reshape, reshape_first = transform.value + if reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + if permute: + tensor = tensor.transpose(permute) + if not reshape_first and reshape is not None: + tensor = tensor.reshape(reshape) + + if base_key not in stacked_buffers: + stacked_buffers[base_key] = {} + stacked_buffers[base_key][idx] = tensor + continue + + if ".__stacked_gate" in jax_key or ".__stacked_up" in jax_key: + if ".__stacked_gate" in jax_key: + base_key = jax_key.replace(".__stacked_gate", "") + idx = 0 + else: + base_key = jax_key.replace(".__stacked_up", "") + idx = 1 + + # No transform for MLP gating (already correct shape) + if base_key not in stacked_buffers: + stacked_buffers[base_key] = {} + stacked_buffers[base_key][idx] = tensor + continue + + # Special handling for vision attention weights + if ( + "mha.query.kernel" in jax_key + or "mha.key.kernel" in jax_key + or "mha.value.kernel" in jax_key + ): + # q/k/v: (num_heads * head_dim, in_features) -> (in_features, num_heads, head_dim) + tensor = tensor.T # transpose + in_features = tensor.shape[0] + tensor = tensor.reshape(in_features, vit_num_heads, vit_head_dim) + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights( + keys, tensor, state_dict, torch_key, None, sharding + ) + except Exception as e: + full_jax_key = ".".join([str(k) for k in keys]) + conversion_errors.append( + f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" + ) + continue + + if "mha.out.kernel" in jax_key: + # out: (out_features, num_heads * head_dim) -> (num_heads, head_dim, out_features) + tensor = ( + tensor.T + ) # transpose to (num_heads * head_dim, out_features) + out_features = tensor.shape[1] + tensor = tensor.reshape(vit_num_heads, vit_head_dim, out_features) + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights( + keys, tensor, state_dict, torch_key, None, sharding + ) + except Exception as e: + full_jax_key = ".".join([str(k) for k in keys]) + conversion_errors.append( + f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" + ) + continue + + if ( + "mha.query.bias" in jax_key + or "mha.key.bias" in jax_key + or "mha.value.bias" in jax_key + ): + # bias: (num_heads * head_dim,) -> (num_heads, head_dim) + tensor = tensor.reshape(vit_num_heads, vit_head_dim) + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights( + keys, tensor, state_dict, torch_key, None, sharding + ) + except Exception as e: + full_jax_key = ".".join([str(k) for k in keys]) + conversion_errors.append( + f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" + ) + continue + + # Special handling for position embedding (add batch dim) + if "pos_embedding" in jax_key and "position_embedding" in torch_key: + tensor = tensor[None, :, :] # (4096, 1152) -> (1, 4096, 1152) + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights( + keys, tensor, state_dict, torch_key, None, sharding + ) + except Exception as e: + full_jax_key = ".".join([str(k) for k in keys]) + conversion_errors.append( + f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" + ) + continue + + keys = [_stoi(k) for k in jax_key.split(".")] + try: + _assign_weights( + keys, tensor, state_dict, torch_key, transform.value, sharding + ) + except Exception as e: + full_jax_key = ".".join([str(k) for k in keys]) + conversion_errors.append( + f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" + ) + gc.collect() + + # Now assign stacked parameters + for base_key, indices in stacked_buffers.items(): + if 0 not in indices or 1 not in indices: + conversion_errors.append(f"Incomplete stacked parameter: {base_key}") + continue + + # Stack the tensors + stacked = jax.numpy.stack([indices[0], indices[1]], axis=0) + + keys = [_stoi(k) for k in base_key.split(".")] + try: + _assign_weights( + keys, stacked, state_dict, f"stacked:{base_key}", None, sharding + ) + except Exception as e: + full_jax_key = ".".join([str(k) for k in keys]) + conversion_errors.append( + f"Failed stacked '{base_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" + ) + + if conversion_errors: + full_error_log = "\n".join(conversion_errors) + raise RuntimeError( + f"Encountered {len(conversion_errors)} weight conversion errors:\n{full_error_log}" + ) + + + gc.collect() + return nnx.merge(graph_def, state_dict) diff --git a/bonsai/models/t5gemma2/tests/run_model.py b/bonsai/models/t5gemma2/tests/run_model.py new file mode 100644 index 00000000..8e26c36a --- /dev/null +++ b/bonsai/models/t5gemma2/tests/run_model.py @@ -0,0 +1,517 @@ +# Copyright 2025 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Test script for T5Gemma2 model (multimodal). + +Usage: + python -m bonsai.models.t5gemma2.tests.run_model + python -m bonsai.models.t5gemma2.tests.run_model --demo image + python -m bonsai.models.t5gemma2.tests.run_model --demo translate +""" + +import argparse + +import jax +import jax.numpy as jnp +import numpy as np +import requests +from huggingface_hub import snapshot_download +from PIL import Image +from transformers import AutoTokenizer + +from bonsai.models.t5gemma2 import modeling, params + + +# ============================================================================= +# Special tokens from modeling.py +# ============================================================================= + +BOS_TOKEN = modeling.BOS_TOKEN # 2 +EOS_TOKEN = modeling.EOS_TOKEN # 1 +NEW_LINE_TOKEN = modeling.NEW_LINE_TOKEN # 108 +START_OF_IMAGE_TOKEN = modeling.START_OF_IMAGE_TOKEN # 255999 +END_OF_IMAGE_TOKEN = modeling.END_OF_IMAGE_TOKEN # 256000 +IMAGE_PLACEHOLDER_TOKEN = modeling.IMAGE_PLACEHOLDER_TOKEN # 256001 + + +# ============================================================================= +# Image token utilities (from text.py, inlined for standalone usage) +# ============================================================================= + + +def _get_new_text_positions(*, offset_on: np.ndarray, offset_by: int) -> np.ndarray: + """Create the positions of the new tokens.""" + offset = np.cumsum(offset_on, axis=-1) * offset_by + new_positions = np.arange(offset_on.shape[-1]) + offset + new_positions -= offset_by * offset_on + return new_positions + + +def _insert_sequence( + tokens: np.ndarray, + *, + at: int, + sequence: list[int], + max_num_images: int, +) -> np.ndarray: + """Inserts a sequence of tokens at all occurrences of a specific token.""" + original_dim = tokens.ndim + if original_dim == 1: + tokens = tokens[None, :] + + batch_size, length = tokens.shape + mm_tokens_to_insert = np.array(sequence) + offset_by = len(mm_tokens_to_insert) - 1 + length_with_mm = length + max_num_images * offset_by + mm_start = tokens == at + + new_tokens = np.zeros((batch_size, length_with_mm), dtype=np.int64) + new_text_pos = _get_new_text_positions(offset_on=mm_start, offset_by=offset_by) + np.put_along_axis(new_tokens, new_text_pos, tokens, axis=1) + + batch_indices_to_zero, _ = np.where(mm_start) + new_pos_to_zero = new_text_pos[mm_start] + if batch_indices_to_zero.size > 0: + new_tokens[batch_indices_to_zero, new_pos_to_zero] = 0 + + batch_indices, seq_indices = np.nonzero(mm_start) + + if batch_indices.size > 0: + intra_batch_img_idx = np.cumsum(mm_start, axis=1)[mm_start] - 1 + final_img_start_pos = seq_indices + intra_batch_img_idx * offset_by + indices_to_insert = final_img_start_pos[:, None] + np.arange( + len(mm_tokens_to_insert), + ) + new_tokens[batch_indices[:, None], indices_to_insert] = mm_tokens_to_insert + + if original_dim == 1: + new_tokens = np.squeeze(new_tokens) + return new_tokens + + +def add_extra_tokens_for_images( + tokens: np.ndarray | list, + *, + new_line_token: int, + start_of_image_token: int, + end_of_image_token: int, + image_placeholder_token: int, + num_placeholder_tokens_per_image: int = 256, + max_num_images: int = 1, +) -> np.ndarray: + """Add extra image tokens to text tokens. + + Expands token into the full image token sequence: + \\n [256 x placeholder] \\n + """ + mm_tokens = [ + new_line_token, + start_of_image_token, + *[image_placeholder_token] * num_placeholder_tokens_per_image, + end_of_image_token, + new_line_token, + ] + if not isinstance(tokens, np.ndarray): + tokens = np.asarray(tokens) + + return _insert_sequence( + at=start_of_image_token, + sequence=mm_tokens, + tokens=tokens, + max_num_images=max_num_images, + ) + + +# ============================================================================= +# Tokenization utilities +# ============================================================================= + + +def tokenize(tokenizer, text: str) -> jnp.ndarray: + """Tokenize text with BOS token prepended. + + Args: + tokenizer: HuggingFace tokenizer. + text: Input text to tokenize. + + Returns: + Token IDs as jnp.ndarray of shape [1, seq_len]. + """ + token_ids = tokenizer.encode(text, add_special_tokens=False) + # Prepend BOS token + token_ids = [BOS_TOKEN] + token_ids + return jnp.array([token_ids], dtype=jnp.int32) + + +def create_multimodal_input( + tokenizer, + text: str, + image_path: str | None = None, + *, + image_size: int = 896, +) -> tuple[jnp.ndarray, jnp.ndarray | None]: + """Create batched input_ids and pixel_values for the model. + + Handles both text tokenization and image preprocessing. When an image is + provided and text contains , the function: + 1. Loads and preprocesses the image + 2. Tokenizes the text + 3. Expands into the full image token sequence + + Args: + tokenizer: HuggingFace tokenizer. + text: Text prompt (may contain marker). + image_path: Optional path to image file or URL. + image_size: Target image size (default 896 for T5Gemma2). + + Returns: + Tuple of (input_ids, pixel_values): + - input_ids: Token IDs of shape [1, seq_len] + - pixel_values: Image array of shape [1, 1, H, W, C] or None if no image + """ + # Tokenize text with BOS + token_ids = tokenizer.encode(text, add_special_tokens=False) + token_ids = [BOS_TOKEN] + token_ids + + # Process image if provided and text contains image marker + pixel_values = None + if image_path is not None and "" in text: + # Load image + if image_path.startswith("http"): + image = Image.open(requests.get(image_path, stream=True).raw) + else: + image = Image.open(image_path) + + # Convert to RGB and resize + image = image.convert("RGB") + image = image.resize((image_size, image_size), Image.BILINEAR) + + # Normalize to [-1, 1] and format as [1, 1, H, W, C] + image_array = np.array(image, dtype=np.float32) / 127.5 - 1.0 + pixel_values = jnp.array(image_array[None, None, ...], dtype=jnp.float32) + + # Expand into full image token sequence + token_ids = add_extra_tokens_for_images( + token_ids, + new_line_token=NEW_LINE_TOKEN, + start_of_image_token=START_OF_IMAGE_TOKEN, + end_of_image_token=END_OF_IMAGE_TOKEN, + image_placeholder_token=IMAGE_PLACEHOLDER_TOKEN, + max_num_images=1, + ) + + # Format as batched array [1, seq_len] + input_ids = jnp.array([token_ids], dtype=jnp.int32) + + return input_ids, pixel_values + + +# ============================================================================= +# Generation +# ============================================================================= + + +def greedy_generate( + model: modeling.T5Gemma2, + encoder_input_ids: jnp.ndarray, + max_new_tokens: int = 50, + eos_token_ids: int | list[int] = EOS_TOKEN, + pixel_values: jnp.ndarray | None = None, + use_cache: bool = True, +) -> jnp.ndarray: + """Greedy decoding for T5Gemma2. + + Args: + model: T5Gemma2 model instance. + encoder_input_ids: Encoder input token IDs [B, L]. + max_new_tokens: Maximum number of tokens to generate. + eos_token_ids: End of sequence token ID(s). + pixel_values: Optional image pixel values [B, N, H, W, C]. + use_cache: Whether to use KV cache for faster decoding. + + Returns: + Generated token IDs [B, num_generated_tokens]. + """ + if isinstance(eos_token_ids, int): + eos_token_ids = [eos_token_ids] + + batch_size = encoder_input_ids.shape[0] + encoder_seq_len = encoder_input_ids.shape[1] + + # Create encoder attention mask + encoder_mask = jnp.ones_like(encoder_input_ids, dtype=jnp.bool_) + + # Encode input (with optional images) + encoder_outputs = model.encoder( + encoder_input_ids, + attention_mask=encoder_mask, + images=pixel_values, + ) + + # Initialize cache if using cached decoding + if use_cache: + model.init_cache( + batch_size=batch_size, + max_decode_length=max_new_tokens + 1, + encoder_seq_length=encoder_seq_len, + ) + + # Start with BOS token + decoder_input_ids = jnp.full((batch_size, 1), BOS_TOKEN, dtype=jnp.int32) + generated_ids = [] + + for step in range(max_new_tokens): + if use_cache: + if step == 0: + # First step: prefill with BOS + decoder_outputs = model.decoder( + decoder_input_ids, + encoder_hidden_states=encoder_outputs, + encoder_attention_mask=encoder_mask, + decode=True, + ) + else: + # Subsequent steps: feed only the new token + decoder_outputs = model.decoder( + next_token[:, None], + encoder_hidden_states=encoder_outputs, + encoder_attention_mask=encoder_mask, + decode=True, + ) + else: + # No cache: recompute full sequence each step + decoder_outputs = model.decoder( + decoder_input_ids, + encoder_hidden_states=encoder_outputs, + encoder_attention_mask=encoder_mask, + decode=False, + ) + + # Get logits from the last position + last_hidden = decoder_outputs[:, -1:, :] + # Use embedding weights for output projection (tied embeddings) + embed_table = model.decoder.embedder.embedding[...] + logits = jnp.einsum("btd,vd->btv", last_hidden, embed_table) + + # Greedy selection + next_token = jnp.argmax(logits[:, -1, :], axis=-1) + generated_ids.append(next_token) + + if not use_cache: + decoder_input_ids = jnp.concatenate( + [decoder_input_ids, next_token[:, None]], + axis=1, + ) + + # Check for EOS + is_eos = jnp.any(jnp.array([next_token == eos for eos in eos_token_ids])) + if is_eos: + break + + return jnp.stack(generated_ids, axis=1) + + +# ============================================================================= +# Demo functions +# ============================================================================= + + +def run_model(): + """Run T5Gemma2 model with example prompts.""" + # Download model checkpoint + model_ckpt_path = snapshot_download("google/t5gemma-2-270m-270m") + config = modeling.T5Gemma2Config.t5gemma2_270m_270m(with_vision=True) + + # Example queries for translation (text-only) + queries = [ + "Translate English to French: Hello, how are you?", + "Translate English to German: The weather is nice today.", + ] + + print("Loading tokenizer...") + tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) + + print("Loading model from checkpoint...") + model = params.create_model_from_safe_tensors(model_ckpt_path, config) + print("Model loaded!") + + generate_steps = 32 + + for query in queries: + print(f"\nInput: {query}") + + # Tokenize + tokens = tokenize(tokenizer, query) + print(f"Token shape: {tokens.shape}") + + # Generate + generated = greedy_generate( + model, + tokens, + max_new_tokens=generate_steps, + use_cache=True, + ) + + # Decode output + output_tokens = jax.device_get(generated[0]) + # Find EOS and truncate + eos_indices = np.where(output_tokens == EOS_TOKEN)[0] + if eos_indices.size > 0: + output_tokens = output_tokens[: eos_indices[0]] + + decoded = tokenizer.decode(output_tokens, skip_special_tokens=True) + print(f"Output: {decoded}\n") + + +def run_multimodal_demo(): + """Run a multimodal demo with images.""" + model_ckpt_path = snapshot_download("google/t5gemma-2-270m-270m") + config = modeling.T5Gemma2Config.t5gemma2_270m_270m(with_vision=True) + + print("Loading tokenizer...") + tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) + + print("Loading model from checkpoint...") + model = params.create_model_from_safe_tensors(model_ckpt_path, config) + print("Model loaded!") + + # Test images with prompts + image_prompts = [ + ( + "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg", + " in this image, there is", + ), + ] + + print("=" * 80) + print("Multimodal Demo - Image Captioning") + print("=" * 80) + + for image_url, text in image_prompts: + print(f"\nImage: {image_url[:60]}...") + print(f"Prompt: {text}") + + # Create multimodal input + input_ids, pixel_values = create_multimodal_input( + tokenizer, text, image_path=image_url + ) + + print(f"Input IDs shape: {input_ids.shape}") + if pixel_values is not None: + print(f"Pixel values shape: {pixel_values.shape}") + + # Generate + generated = greedy_generate( + model, + input_ids, + max_new_tokens=30, + pixel_values=pixel_values, + use_cache=True, + ) + + output_tokens = jax.device_get(generated[0]) + decoded = tokenizer.decode(output_tokens, skip_special_tokens=True) + print(f"Output: {decoded}") + print("-" * 80) + + print("\nMultimodal demo complete!") + + +def run_translation_demo(): + """Run a translation demo with few-shot examples.""" + model_ckpt_path = snapshot_download("google/t5gemma-2-270m-270m") + config = modeling.T5Gemma2Config.t5gemma2_270m_270m(with_vision=True) + + print("Loading tokenizer...") + tokenizer = AutoTokenizer.from_pretrained(model_ckpt_path) + + print("Loading model from checkpoint...") + model = params.create_model_from_safe_tensors(model_ckpt_path, config) + print("Model loaded!") + + # Few-shot translation prompt (English to Chinese) + prompt = """Translate the text from English to Chinese (zh_CN). + + English: Hook the vacuum up to it, it helps keep the dust down, and you can prep your concrete that way also. We prefer to do it with the hand grinders ourself, so either way will work pretty good. Usually if it's got an epoxy coating already, the hand grinders work a little better, but this thing works really good too. So after we do that, we fix all the cracks, fix all the divots with our patch repair material, and then we grind them smooth. And then we clean the concrete and get it ready for the first coating. This is going to be a 100% solids epoxy, so it goes on in its different procedures, different stages. + Chinese (zh_CN): 把吸尘器连在上面，这样可以减少灰尘，你还可以顺便处理一下混凝土。我们更喜欢用手持式研磨机，两种方法的效果都不错。通常，如果表面已经有环氧树脂涂层的话，手持式研磨机的效果会更好一些，但这台机器的效果也很好。处理完之后，我们用修补材料把所有的裂缝和坑洼都补好，再把它们磨平。然后清洁混凝土表面，为涂刷第一层涂料做好准备。我们这里用的是 100% 固体环氧树脂，所以涂刷的时候会有不同的步骤和阶段。 + + English: The Zoroastrian text, Vendidad, states that Yima built an underground city on the orders of the god Ahura Mazda, to protect his people from a catastrophic winter. Much like the account of Noah in the Bible, Yima was instructed to collect pairs of the best animals and people as well as the best seeds in order to reseed the Earth after the winter cataclysm. This was before the last Ice Age, 110,000 years ago. + Chinese (zh_CN): 琐罗亚斯德教经典《万迪达德》中说，Yima 奉阿胡拉·马兹达神之命建造了一座地下城市，以保护他的人民躲避一场灾难性的寒冬。就像《圣经》中诺亚的故事一样，Yima 被指示收集最好的动物、人类以及最好的种子，以便在冬季灾难过后重新在地球上播种。这发生在最后一个冰河时代之前，也就是 11 万年前。 + + English: Okay, so let me explain. All right, so the problem is, if you look inside there... You see the wood siding? There's the old siding, and it butts up to the shingles there. And then I put this over it. And what happens is the dirt collects there, to that flashing. + Chinese (zh_CN): 好的，我来解释一下。问题是，你们看里面……看到木头墙板了吗？那是原来的墙板，紧挨着那边的瓦。然后我把这个盖在上面。结果灰尘就堆积在那儿，堆积到泛水板上。 + + English: Hey guys, Thunder E here, and welcome to the video you've been waiting for. I am talking about gaming on the ASUS ROG Phone 5. Now, the ROG Phone series is well known for its gaming powers, but in this video, we're going to find out if the ROG Phone 5 is truly taking back the crown as the king of gaming phones. + Chinese (zh_CN): 大家好，我是雷霆 E，欢迎收看大家期待已久的视频。今天要评测的是华硕 ROG Phone 5 的游戏性能。ROG Phone 系列手机一直以其强大的游戏性能而闻名，那么，ROG Phone 5 能否真正加冕"游戏手机之王"？我们拭目以待。 + + English: It is December 1997, and the Imperial Sugar Company is acquiring a new production site at Port Wentworth from Savannah Foods and Industries Incorporated. There is nothing really of note here. It was doing what businesses do, and that is acquiring to expand. The site has been home to food production and processing since the early 1900s. Savannah Industries Incorporated began construction of granulated sugar production facilities at Port Wentworth during the 1910s, completing it in 1917. + Chinese (zh_CN): 那是 1997 年 12 月，帝国糖业公司正从萨凡纳食品和工业有限公司手中收购位于温特沃斯港的一个新生产基地。这的确没什么值得注意的，它做的只是一家公司都会做的事情，那就是通过收购来扩张。该基地自 20 世纪初以来一直是食品生产和加工的场所。萨凡纳工业有限公司在 20 世纪 10 年代开始在温特沃斯港建造砂糖生产设施，并于 1917 年竣工。 + + English: Time for the Scotty Kilmer channel. Does your car have faded paint on it? Then stay tuned, because today I'm going to show you how to polish off faded paint. And all it takes is a bucket of water, a polisher, and a bottle of this Meguiar's Ultimate Compound. + Chinese (zh_CN):""" + + print("=" * 60) + print("Translation Demo: English to Chinese (zh_CN)") + print("=" * 60) + + tokens = tokenize(tokenizer, prompt) + print(f"\nInput length: {tokens.shape[1]} tokens") + + print("\nGenerating (max 100 tokens, stops at EOS or newline)...") + generated = greedy_generate( + model, + tokens, + max_new_tokens=100, + eos_token_ids=[EOS_TOKEN, NEW_LINE_TOKEN], + use_cache=True, + ) + + output_tokens = jax.device_get(generated[0]) + decoded = tokenizer.decode(output_tokens, skip_special_tokens=True) + + print(f"\n{'=' * 60}") + print("Generated Translation:") + print("=" * 60) + print(decoded) + print("=" * 60) + + +# ============================================================================= +# Main +# ============================================================================= + + +def main(): + parser = argparse.ArgumentParser(description="Run T5Gemma2 model tests") + parser.add_argument( + "--demo", + type=str, + choices=["image", "translate"], + default=None, + help="Run a demo: 'image' for multimodal, 'translate' for translation", + ) + args = parser.parse_args() + + if args.demo == "image": + run_multimodal_demo() + elif args.demo == "translate": + run_translation_demo() + else: + print("Running default text-only demo...") + print("Use --demo image for multimodal or --demo translate for translation\n") + run_model() + + +if __name__ == "__main__": + main() From 4b04264449f3bbfac5ebdc99c2cd4533694c13bd Mon Sep 17 00:00:00 2001 From: Jaonary Rabarisoa Date: Tue, 27 Jan 2026 15:58:42 +0100 Subject: [PATCH 4/5] Fix formating --- bonsai/models/t5gemma2/modeling.py | 197 +++++++++++++--------- bonsai/models/t5gemma2/params.py | 77 ++------- bonsai/models/t5gemma2/tests/run_model.py | 5 +- 3 files changed, 137 insertions(+), 142 deletions(-) diff --git a/bonsai/models/t5gemma2/modeling.py b/bonsai/models/t5gemma2/modeling.py index b504951e..2a21fec8 100644 --- a/bonsai/models/t5gemma2/modeling.py +++ b/bonsai/models/t5gemma2/modeling.py @@ -11,14 +11,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +from __future__ import annotations import functools from dataclasses import dataclass, field from enum import Enum import jax import jax.numpy as jnp -import numpy as np from flax import nnx from flax.nnx.module import first_from from jaxtyping import Array, Bool, Float, Int @@ -36,15 +35,19 @@ class AttentionType(Enum): # Attention pattern: 5 local sliding + 1 global _ATTN_PATTERN = ( - AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, - AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, AttentionType.GLOBAL, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.GLOBAL, ) def _make_layer_types(num_layers: int) -> tuple[AttentionType, ...]: """Generate attention types for all layers using the standard pattern.""" n = len(_ATTN_PATTERN) - return (_ATTN_PATTERN * (num_layers // n) + _ATTN_PATTERN[: num_layers % n]) + return _ATTN_PATTERN * (num_layers // n) + _ATTN_PATTERN[: num_layers % n] @dataclass(frozen=True) @@ -248,13 +251,13 @@ def _get_rope_scale_factor( def apply_rope( - inputs: Float[Array, "B L N H"], - positions: Float[Array, "B L"], + inputs: Float[Array, "B L N H"], # noqa: F722 + positions: Float[Array, "B L"], # noqa: F722 *, base_frequency: int, scale_factor: float = 1.0, rope_proportion: float = 1.0, -) -> Float[Array, "B L N H"]: +) -> Float[Array, "B L N H"]: # noqa: F722 """Applies Rotary Position Embeddings (RoPE). Args: @@ -323,18 +326,20 @@ def __init__( self, features: int, *, + epsilon: float = 1e-6, scale_init: nnx.Initializer = ZEROS_INIT, dtype: jnp.dtype = jnp.float32, rngs: nnx.Rngs, ): + self.epsilon = epsilon self.scale = nnx.Param(scale_init(rngs.params(), (features,), dtype)) - def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: + def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: # noqa: F722 var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) # Jax.lax.rsqrt is used because it returns different floats than # jnp.reciprocal(jnp.sqrt(var + 1e-06)) - normed_inputs = x * jax.lax.rsqrt(var + 1e-06) + normed_inputs = x * jax.lax.rsqrt(var + self.epsilon) # normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is # a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to @@ -379,7 +384,7 @@ def __init__( rngs=rngs, ) - def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: + def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: # noqa: F722 """Applies feed-forward transformation. Args: @@ -400,8 +405,8 @@ def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: def make_bidirectional_mask( - input_mask: Bool[Array, "B L"], -) -> Bool[Array, "B 1 L L"]: + input_mask: Bool[Array, "B L"], # noqa: F722 +) -> Bool[Array, "B 1 L L"]: # noqa: F722 """Creates a bidirectional attention mask for encoder. Args: @@ -416,8 +421,8 @@ def make_bidirectional_mask( def make_causal_mask( - input_mask: Bool[Array, "B L"], -) -> Bool[Array, "B 1 L L"]: + input_mask: Bool[Array, "B L"], # noqa: F722 +) -> Bool[Array, "B 1 L L"]: # noqa: F722 """Creates a causal attention mask for decoder. Args: @@ -435,11 +440,11 @@ def make_causal_mask( def make_sliding_window_mask( - positions: Int[Array, "B L"], + positions: Int[Array, "B L"], # noqa: F722 sliding_window: int, *, bidirectional: bool = False, -) -> Bool[Array, "B 1 L L"]: +) -> Bool[Array, "B 1 L L"]: # noqa: F722 """Creates a sliding window mask. Args: @@ -475,10 +480,10 @@ def make_sliding_window_mask( def make_sliding_window_causal_mask( - input_mask: Bool[Array, "B L"], - positions: Int[Array, "B L"], + input_mask: Bool[Array, "B L"], # noqa: F722 + positions: Int[Array, "B L"], # noqa: F722 sliding_window: int, -) -> Bool[Array, "B 1 L L"]: +) -> Bool[Array, "B 1 L L"]: # noqa: F722 """Creates a causal mask with sliding window for decoder. Combines causal mask (lower triangular) with sliding window constraint. @@ -502,9 +507,9 @@ def make_sliding_window_causal_mask( def make_merged_attention_mask( - decoder_mask: Bool[Array, "B 1 L_dec L_dec"], - encoder_mask: Bool[Array, "B 1 1 L_enc"], -) -> Bool[Array, "B 1 L_dec L_combined"]: + decoder_mask: Bool[Array, "B 1 L_dec L_dec"], # noqa: F722 + encoder_mask: Bool[Array, "B 1 1 L_enc"], # noqa: F722 +) -> Bool[Array, "B 1 L_dec L_combined"]: # noqa: F722 """Creates merged attention mask for decoder's merged attention. Concatenates the decoder self-attention mask with the cross-attention @@ -532,7 +537,7 @@ def make_decode_mode_self_mask( query_len: int, max_cache_len: int, current_cache_len: int, -) -> Bool[Array, "B 1 Q K"]: +) -> Bool[Array, "B 1 Q K"]: # noqa: F722 """Creates self-attention mask for decode mode (with KV cache). During decode mode, the query attends to all previously cached keys. @@ -557,7 +562,7 @@ def make_decode_mode_sliding_mask( query_pos: int, max_cache_len: int, sliding_window: int, -) -> Bool[Array, "B 1 1 K"]: +) -> Bool[Array, "B 1 1 K"]: # noqa: F722 """Creates sliding window mask for decode mode (with KV cache). During decode mode with sliding window, the query can only attend to @@ -619,7 +624,7 @@ def __init__( ZEROS_INIT(rngs.params(), (embed_dim,), dtype), ) - def __call__(self, input_ids: Int[Array, "B L"]) -> Float[Array, "B L D"]: + def __call__(self, input_ids: Int[Array, "B L"]) -> Float[Array, "B L D"]: # noqa: F722 """Embed input tokens with scaling. Args: @@ -648,7 +653,7 @@ def _posemb_sincos_2d( width: int, temperature: float = 10_000.0, dtype: jnp.dtype = jnp.float32, -) -> Float[Array, "1 M D"]: +) -> Float[Array, "1 M D"]: # noqa: F722 """Sinusoidal 2D position embeddings (MoCo v3 style).""" y, x = jnp.mgrid[:h, :w] @@ -693,10 +698,10 @@ def __init__( def __call__( self, - x: Float[Array, "B L D"], + x: Float[Array, "B L D"], # noqa: F722 *, deterministic: bool | None = None, - ) -> Float[Array, "B L D"]: + ) -> Float[Array, "B L D"]: # noqa: F722 x = self.dense1(x) x = jax.nn.gelu(x) x = self.dropout(x, deterministic=deterministic) @@ -754,10 +759,10 @@ def __init__( def __call__( self, - x: Float[Array, "B L D"], + x: Float[Array, "B L D"], # noqa: F722 *, deterministic: bool | None = None, - ) -> Float[Array, "B L D"]: + ) -> Float[Array, "B L D"]: # noqa: F722 y = self.layer_norm1(x) y = self.mha(y, deterministic=deterministic) y = self.dropout1(y, deterministic=deterministic) @@ -899,10 +904,10 @@ def __init__( def __call__( self, - images: Float[Array, "B N H W C"], + images: Float[Array, "B N H W C"], # noqa: F722 *, deterministic: bool | None = None, - ) -> Float[Array, "B N P D"]: + ) -> Float[Array, "B N P D"]: # noqa: F722 if len(images.shape) == 4: images = images[:, None, :] b, n, h, w, c = images.shape @@ -964,7 +969,7 @@ def __init__( rngs=rngs, ) - def __call__(self, x: Float[Array, "B N P Dv"]) -> Float[Array, "B N P De"]: + def __call__(self, x: Float[Array, "B N P Dv"]) -> Float[Array, "B N P De"]: # noqa: F722 x = self.mm_soft_embedding_norm(x) return self.mm_input_projection("...tm,md->...td", x) @@ -998,10 +1003,10 @@ def __init__( def __call__( self, - images: Float[Array, "B N H W C"], + images: Float[Array, "B N H W C"], # noqa: F722 *, deterministic: bool | None = None, - ) -> Float[Array, "B N P De"]: + ) -> Float[Array, "B N P De"]: # noqa: F722 soft_tokens = self.soft_tokenizer(images, deterministic=deterministic) if self.freeze_params: @@ -1067,6 +1072,7 @@ def __init__( *, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, + rms_norm_eps: float = 1e-6, attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_qk_norm: bool = False, @@ -1083,6 +1089,7 @@ def __init__( self.query_pre_attn_scalar = query_pre_attn_scalar self.rope_base_frequency = rope_base_frequency self.rope_scale_factor = rope_scale_factor + self.rms_norm_eps = rms_norm_eps self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.use_qk_norm = use_qk_norm @@ -1122,12 +1129,14 @@ def __init__( if use_qk_norm: self._query_norm = T5Gemma2RMSNorm( head_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, ) self._key_norm = T5Gemma2RMSNorm( head_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1143,10 +1152,10 @@ def use_gqa(self) -> bool: def __call__( self, - x: Float[Array, "B L D"], - segment_pos: Float[Array, "B L"], - attn_mask: Float[Array, "B L L"], - ) -> Float[Array, "B L D"]: + x: Float[Array, "B L D"], # noqa: F722 + segment_pos: Float[Array, "B L"], # noqa: F722 + attn_mask: Float[Array, "B L L"], # noqa: F722 + ) -> Float[Array, "B L D"]: # noqa: F722 """Applies bidirectional multi-head attention. Args: @@ -1222,9 +1231,9 @@ def __call__( # Compute attention output if self.use_gqa: - b, t, kg, h = probs.shape + b, t, kg, s = probs.shape probs = probs.reshape( - (b, t, self.num_kv_heads, int(kg / self.num_kv_heads), h), + (b, t, self.num_kv_heads, int(kg / self.num_kv_heads), s), ) encoded = jnp.einsum("BTKGS,BSKH->BTKGH", probs, value_proj) b, t, k, g, h = encoded.shape @@ -1257,6 +1266,7 @@ def __init__( transpose_gating_einsum: bool, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, + rms_norm_eps: float = 1e-6, attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_qk_norm: bool = False, @@ -1270,6 +1280,7 @@ def __init__( self.pre_attention_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1284,6 +1295,7 @@ def __init__( query_pre_attn_scalar=query_pre_attn_scalar, rope_base_frequency=rope_base_frequency, rope_scale_factor=rope_scale_factor, + rms_norm_eps=rms_norm_eps, attn_logits_soft_cap=attn_logits_soft_cap, sliding_window_size=sliding_window_size, use_qk_norm=use_qk_norm, @@ -1295,6 +1307,7 @@ def __init__( if use_post_attn_norm: self.post_attention_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1302,6 +1315,7 @@ def __init__( self.pre_ffw_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1319,6 +1333,7 @@ def __init__( if use_post_ffw_norm: self.post_ffw_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1326,10 +1341,10 @@ def __init__( def __call__( self, - x: Float[Array, "B L D"], - segment_pos: Int[Array, "B L"], - attn_mask: Bool[Array, "B L L"], - ) -> Float[Array, "B L D"]: + x: Float[Array, "B L D"], # noqa: F722 + segment_pos: Int[Array, "B L"], # noqa: F722 + attn_mask: Bool[Array, "B L L"], # noqa: F722 + ) -> Float[Array, "B L D"]: # noqa: F722 """Apply encoder layer. Args: @@ -1424,6 +1439,7 @@ def __init__( text_config, attention_types[i], ), + rms_norm_eps=text_config.rms_norm_eps, attn_logits_soft_cap=text_config.attn_logit_softcapping, sliding_window_size=( text_config.sliding_window if attention_types[i] == AttentionType.LOCAL_SLIDING else None @@ -1437,24 +1453,29 @@ def __init__( ) # Final normalization - self.norm = T5Gemma2RMSNorm(text_config.embed_dim, dtype=dtype, rngs=rngs) + self.norm = T5Gemma2RMSNorm( + text_config.embed_dim, + epsilon=text_config.rms_norm_eps, + dtype=dtype, + rngs=rngs, + ) def __call__( self, - tokens: Int[Array, "B L"], - attention_mask: Bool[Array, "B L"] | None = None, - position_ids: Int[Array, "B L"] | None = None, - images: Float[Array, "B N H W C"] | None = None, + tokens: Int[Array, "B L"], # noqa: F722 + attention_mask: Bool[Array, "B L"] | None = None, # noqa: F722 + position_ids: Int[Array, "B L"] | None = None, # noqa: F722 + images: Float[Array, "B N H W C"] | None = None, # noqa: F722 *, deterministic: bool = True, - ) -> Float[Array, "B L D"]: + ) -> Float[Array, "B L D"]: # noqa: F722 """Forward pass of the encoder. Args: - input_ids: Input token IDs [B, L]. + tokens: Input token IDs [B, L]. attention_mask: Attention mask [B, L]. position_ids: Position indices [B, L]. - pixel_values: Optional input images [B, N, H, W, C] where N is images per batch. + images: Optional input images [B, N, H, W, C] where N is images per batch. deterministic: Whether to run in deterministic mode. Returns: @@ -1532,6 +1553,7 @@ def __init__( *, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, + rms_norm_eps: float = 1e-6, attn_logits_soft_cap: float | None = None, use_qk_norm: bool = True, kernel_init: nnx.Initializer = NORMAL_INIT, @@ -1547,6 +1569,7 @@ def __init__( self.query_pre_attn_scalar = query_pre_attn_scalar self.rope_base_frequency = rope_base_frequency self.rope_scale_factor = rope_scale_factor + self.rms_norm_eps = rms_norm_eps self.attn_logits_soft_cap = attn_logits_soft_cap self.use_qk_norm = use_qk_norm self.decode = decode @@ -1583,12 +1606,14 @@ def __init__( if use_qk_norm: self.q_norm = T5Gemma2RMSNorm( head_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, ) self.k_norm = T5Gemma2RMSNorm( head_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1647,13 +1672,13 @@ def init_cache( def __call__( self, - hidden_states: Float[Array, "B L_dec D"], - encoder_hidden_states: Float[Array, "B L_enc D"], - position_ids: Int[Array, "B L_dec"], - merged_attention_mask: Bool[Array, "B 1 L_dec L_combined"] | None = None, + hidden_states: Float[Array, "B L_dec D"], # noqa: F722 + encoder_hidden_states: Float[Array, "B L_enc D"], # noqa: F722 + position_ids: Int[Array, "B L_dec"], # noqa: F722 + merged_attention_mask: Bool[Array, "B 1 L_dec L_combined"] | None = None, # noqa: F722 *, decode: bool | None = None, - ) -> Float[Array, "B L_dec D"]: + ) -> Float[Array, "B L_dec D"]: # noqa: F722 """Apply merged self-attention and cross-attention. Args: @@ -1824,6 +1849,7 @@ def __init__( query_pre_attn_scalar: float, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, + rms_norm_eps: float = 1e-6, attn_logits_soft_cap: float | None = None, use_qk_norm: bool = True, kernel_init: nnx.Initializer = NORMAL_INIT, @@ -1839,6 +1865,7 @@ def __init__( # Pre-attention norm self.pre_attention_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1853,6 +1880,7 @@ def __init__( query_pre_attn_scalar=query_pre_attn_scalar, rope_base_frequency=rope_base_frequency, rope_scale_factor=rope_scale_factor, + rms_norm_eps=rms_norm_eps, attn_logits_soft_cap=attn_logits_soft_cap, use_qk_norm=use_qk_norm, kernel_init=kernel_init, @@ -1866,6 +1894,7 @@ def __init__( if use_post_attn_norm: self.post_attention_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1874,6 +1903,7 @@ def __init__( # Pre-FFW norm self.pre_ffw_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1893,6 +1923,7 @@ def __init__( if use_post_ffw_norm: self.post_ffw_norm = T5Gemma2RMSNorm( embed_dim, + epsilon=rms_norm_eps, dtype=dtype, scale_init=scale_init, rngs=rngs, @@ -1900,13 +1931,13 @@ def __init__( def __call__( self, - x: Float[Array, "B L_dec D"], - encoder_hidden_states: Float[Array, "B L_enc D"], - segment_pos: Int[Array, "B L_dec"], - attn_mask: Bool[Array, "B 1 L_dec L_combined"] | None = None, + x: Float[Array, "B L_dec D"], # noqa: F722 + encoder_hidden_states: Float[Array, "B L_enc D"], # noqa: F722 + segment_pos: Int[Array, "B L_dec"], # noqa: F722 + attn_mask: Bool[Array, "B 1 L_dec L_combined"] | None = None, # noqa: F722 *, decode: bool | None = None, - ) -> Float[Array, "B L_dec D"]: + ) -> Float[Array, "B L_dec D"]: # noqa: F722 """Apply decoder layer. Args: @@ -1994,6 +2025,7 @@ def __init__( query_pre_attn_scalar=config.query_pre_attn_scalar, rope_base_frequency=_get_rope_base_frequency(config, attention_types[i]), rope_scale_factor=_get_rope_scale_factor(config, attention_types[i]), + rms_norm_eps=config.rms_norm_eps, attn_logits_soft_cap=config.attn_logit_softcapping, use_qk_norm=True, decode=decode, @@ -2005,7 +2037,12 @@ def __init__( ) # Final normalization - self.norm = T5Gemma2RMSNorm(config.embed_dim, dtype=dtype, rngs=rngs) + self.norm = T5Gemma2RMSNorm( + config.embed_dim, + epsilon=config.rms_norm_eps, + dtype=dtype, + rngs=rngs, + ) def init_cache( self, @@ -2044,14 +2081,14 @@ def get_cache_index(self) -> int: def __call__( self, - input_ids: Int[Array, "B L_dec"], - encoder_hidden_states: Float[Array, "B L_enc D"], - attention_mask: Bool[Array, "B L_dec"] | None = None, - encoder_attention_mask: Bool[Array, "B L_enc"] | None = None, - position_ids: Int[Array, "B L_dec"] | None = None, + input_ids: Int[Array, "B L_dec"], # noqa: F722 + encoder_hidden_states: Float[Array, "B L_enc D"], # noqa: F722 + attention_mask: Bool[Array, "B L_dec"] | None = None, # noqa: F722 + encoder_attention_mask: Bool[Array, "B L_enc"] | None = None, # noqa: F722 + position_ids: Int[Array, "B L_dec"] | None = None, # noqa: F722 *, decode: bool | None = None, - ) -> Float[Array, "B L_dec D"]: + ) -> Float[Array, "B L_dec D"]: # noqa: F722 """Forward pass of the decoder. Args: @@ -2239,18 +2276,18 @@ def init_cache( def __call__( self, - input_ids: Int[Array, "B L_enc"], - decoder_input_ids: Int[Array, "B L_dec"], - attention_mask: Bool[Array, "B L_enc"] | None = None, - decoder_attention_mask: Bool[Array, "B L_dec"] | None = None, - position_ids: Int[Array, "B L_enc"] | None = None, - decoder_position_ids: Int[Array, "B L_dec"] | None = None, - pixel_values: Float[Array, "B N H W C"] | None = None, - encoder_outputs: Float[Array, "B L_enc D"] | None = None, + input_ids: Int[Array, "B L_enc"], # noqa: F722 + decoder_input_ids: Int[Array, "B L_dec"], # noqa: F722 + attention_mask: Bool[Array, "B L_enc"] | None = None, # noqa: F722 + decoder_attention_mask: Bool[Array, "B L_dec"] | None = None, # noqa: F722 + position_ids: Int[Array, "B L_enc"] | None = None, # noqa: F722 + decoder_position_ids: Int[Array, "B L_dec"] | None = None, # noqa: F722 + pixel_values: Float[Array, "B N H W C"] | None = None, # noqa: F722 + encoder_outputs: Float[Array, "B L_enc D"] | None = None, # noqa: F722 *, decode: bool | None = None, deterministic: bool = True, - ) -> tuple[Float[Array, "B L_dec D"], Float[Array, "B L_enc D"]]: + ) -> tuple[Float[Array, "B L_dec D"], Float[Array, "B L_enc D"]]: # noqa: F722 """Forward pass of the full model. Args: diff --git a/bonsai/models/t5gemma2/params.py b/bonsai/models/t5gemma2/params.py index b41295dc..9be32ceb 100644 --- a/bonsai/models/t5gemma2/params.py +++ b/bonsai/models/t5gemma2/params.py @@ -335,9 +335,7 @@ def _assign_weights(keys, tensor, state_dict, st_key, transform, sharding_dict): if not reshape_first and reshape is not None: tensor = tensor.reshape(reshape) if tensor.shape != state_dict[key].shape: - raise ValueError( - f"Shape mismatch for {st_key}: {tensor.shape} vs {state_dict[key].shape}" - ) + raise ValueError(f"Shape mismatch for {st_key}: {tensor.shape} vs {state_dict[key].shape}") if sharding_dict is not None: state_dict[key] = jax.device_put(tensor, sharding_dict[key]) else: @@ -367,9 +365,7 @@ def create_model_from_safe_tensors( model = model_lib.T5Gemma2(cfg, rngs=nnx.Rngs(params=0, dropout=0)) graph_def, state = nnx.split(model) state_dict = state.to_pure_dict() - sharding = ( - nnx.get_named_sharding(state, mesh).to_pure_dict() if mesh is not None else None - ) + sharding = nnx.get_named_sharding(state, mesh).to_pure_dict() if mesh is not None else None key_mapping, vis_cfg = _get_key_and_transform_mapping(cfg) conversion_errors = [] @@ -434,63 +430,41 @@ def create_model_from_safe_tensors( continue # Special handling for vision attention weights - if ( - "mha.query.kernel" in jax_key - or "mha.key.kernel" in jax_key - or "mha.value.kernel" in jax_key - ): + if "mha.query.kernel" in jax_key or "mha.key.kernel" in jax_key or "mha.value.kernel" in jax_key: # q/k/v: (num_heads * head_dim, in_features) -> (in_features, num_heads, head_dim) tensor = tensor.T # transpose in_features = tensor.shape[0] tensor = tensor.reshape(in_features, vit_num_heads, vit_head_dim) keys = [_stoi(k) for k in jax_key.split(".")] try: - _assign_weights( - keys, tensor, state_dict, torch_key, None, sharding - ) + _assign_weights(keys, tensor, state_dict, torch_key, None, sharding) except Exception as e: full_jax_key = ".".join([str(k) for k in keys]) - conversion_errors.append( - f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" - ) + conversion_errors.append(f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}") continue if "mha.out.kernel" in jax_key: # out: (out_features, num_heads * head_dim) -> (num_heads, head_dim, out_features) - tensor = ( - tensor.T - ) # transpose to (num_heads * head_dim, out_features) + tensor = tensor.T # transpose to (num_heads * head_dim, out_features) out_features = tensor.shape[1] tensor = tensor.reshape(vit_num_heads, vit_head_dim, out_features) keys = [_stoi(k) for k in jax_key.split(".")] try: - _assign_weights( - keys, tensor, state_dict, torch_key, None, sharding - ) + _assign_weights(keys, tensor, state_dict, torch_key, None, sharding) except Exception as e: full_jax_key = ".".join([str(k) for k in keys]) - conversion_errors.append( - f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" - ) + conversion_errors.append(f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}") continue - if ( - "mha.query.bias" in jax_key - or "mha.key.bias" in jax_key - or "mha.value.bias" in jax_key - ): + if "mha.query.bias" in jax_key or "mha.key.bias" in jax_key or "mha.value.bias" in jax_key: # bias: (num_heads * head_dim,) -> (num_heads, head_dim) tensor = tensor.reshape(vit_num_heads, vit_head_dim) keys = [_stoi(k) for k in jax_key.split(".")] try: - _assign_weights( - keys, tensor, state_dict, torch_key, None, sharding - ) + _assign_weights(keys, tensor, state_dict, torch_key, None, sharding) except Exception as e: full_jax_key = ".".join([str(k) for k in keys]) - conversion_errors.append( - f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" - ) + conversion_errors.append(f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}") continue # Special handling for position embedding (add batch dim) @@ -498,26 +472,18 @@ def create_model_from_safe_tensors( tensor = tensor[None, :, :] # (4096, 1152) -> (1, 4096, 1152) keys = [_stoi(k) for k in jax_key.split(".")] try: - _assign_weights( - keys, tensor, state_dict, torch_key, None, sharding - ) + _assign_weights(keys, tensor, state_dict, torch_key, None, sharding) except Exception as e: full_jax_key = ".".join([str(k) for k in keys]) - conversion_errors.append( - f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" - ) + conversion_errors.append(f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}") continue keys = [_stoi(k) for k in jax_key.split(".")] try: - _assign_weights( - keys, tensor, state_dict, torch_key, transform.value, sharding - ) + _assign_weights(keys, tensor, state_dict, torch_key, transform.value, sharding) except Exception as e: full_jax_key = ".".join([str(k) for k in keys]) - conversion_errors.append( - f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" - ) + conversion_errors.append(f"Failed '{torch_key}' -> '{full_jax_key}': {type(e).__name__}: {e}") gc.collect() # Now assign stacked parameters @@ -531,21 +497,14 @@ def create_model_from_safe_tensors( keys = [_stoi(k) for k in base_key.split(".")] try: - _assign_weights( - keys, stacked, state_dict, f"stacked:{base_key}", None, sharding - ) + _assign_weights(keys, stacked, state_dict, f"stacked:{base_key}", None, sharding) except Exception as e: full_jax_key = ".".join([str(k) for k in keys]) - conversion_errors.append( - f"Failed stacked '{base_key}' -> '{full_jax_key}': {type(e).__name__}: {e}" - ) + conversion_errors.append(f"Failed stacked '{base_key}' -> '{full_jax_key}': {type(e).__name__}: {e}") if conversion_errors: full_error_log = "\n".join(conversion_errors) - raise RuntimeError( - f"Encountered {len(conversion_errors)} weight conversion errors:\n{full_error_log}" - ) - + raise RuntimeError(f"Encountered {len(conversion_errors)} weight conversion errors:\n{full_error_log}") gc.collect() return nnx.merge(graph_def, state_dict) diff --git a/bonsai/models/t5gemma2/tests/run_model.py b/bonsai/models/t5gemma2/tests/run_model.py index 8e26c36a..076a35d4 100644 --- a/bonsai/models/t5gemma2/tests/run_model.py +++ b/bonsai/models/t5gemma2/tests/run_model.py @@ -270,6 +270,7 @@ def greedy_generate( # Start with BOS token decoder_input_ids = jnp.full((batch_size, 1), BOS_TOKEN, dtype=jnp.int32) generated_ids = [] + next_token = None # Track the last generated token for cached decoding for step in range(max_new_tokens): if use_cache: @@ -403,9 +404,7 @@ def run_multimodal_demo(): print(f"Prompt: {text}") # Create multimodal input - input_ids, pixel_values = create_multimodal_input( - tokenizer, text, image_path=image_url - ) + input_ids, pixel_values = create_multimodal_input(tokenizer, text, image_path=image_url) print(f"Input IDs shape: {input_ids.shape}") if pixel_values is not None: From e534b2667400d57c0c2da2081b1a39e812fe9e95 Mon Sep 17 00:00:00 2001 From: Jaonary Rabarisoa Date: Mon, 2 Feb 2026 10:03:46 +0100 Subject: [PATCH 5/5] Add sharding configuration classes, update model parameters and add test --- bonsai/models/t5gemma2/README.md | 12 +- bonsai/models/t5gemma2/modeling.py | 739 +++++------------- bonsai/models/t5gemma2/tests/run_model.py | 24 +- .../t5gemma2/tests/test_outputs_t5gemma2.py | 145 ++++ 4 files changed, 370 insertions(+), 550 deletions(-) create mode 100644 bonsai/models/t5gemma2/tests/test_outputs_t5gemma2.py diff --git a/bonsai/models/t5gemma2/README.md b/bonsai/models/t5gemma2/README.md index 2c5d8e54..872390e1 100644 --- a/bonsai/models/t5gemma2/README.md +++ b/bonsai/models/t5gemma2/README.md @@ -1 +1,11 @@ -# T5Gemma2 in JAX (WIP) \ No newline at end of file +# T5Gemma2 in JAX + +This directory contains a pure JAX implementation of the [T5Gemma2 model](https://huggingface.co/collections/google/t5gemma-2-release-6839e38ad1e09ed3703c47e7), using the [Flax NNX](https://flax.readthedocs.io/en/stable/index.html) API. + +### Running this model + +```sh +python -m bonsai.models.t5gemma2.tests.run_model +python -m bonsai.models.t5gemma2.tests.run_model --demo image +python -m bonsai.models.t5gemma2.tests.run_model --demo translate +``` diff --git a/bonsai/models/t5gemma2/modeling.py b/bonsai/models/t5gemma2/modeling.py index 2a21fec8..23ae620b 100644 --- a/bonsai/models/t5gemma2/modeling.py +++ b/bonsai/models/t5gemma2/modeling.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from __future__ import annotations + import functools from dataclasses import dataclass, field from enum import Enum @@ -20,19 +20,21 @@ import jax.numpy as jnp from flax import nnx from flax.nnx.module import first_from +from jax import P +from jax.sharding import PartitionSpec from jaxtyping import Array, Bool, Float, Int -# ============================================================================= -# Configuration Classes -# ============================================================================= - - class AttentionType(Enum): GLOBAL = "global" LOCAL_SLIDING = "local_sliding" +class ShardMode(Enum): + FSDP = "fsdp" + TP = "tp" + + # Attention pattern: 5 local sliding + 1 global _ATTN_PATTERN = ( AttentionType.LOCAL_SLIDING, @@ -45,11 +47,101 @@ class AttentionType(Enum): def _make_layer_types(num_layers: int) -> tuple[AttentionType, ...]: - """Generate attention types for all layers using the standard pattern.""" n = len(_ATTN_PATTERN) return _ATTN_PATTERN * (num_layers // n) + _ATTN_PATTERN[: num_layers % n] +@dataclass(slots=True, frozen=True) +class VisionShardingCfg: + attn_kernel: PartitionSpec | None = None + attn_bias: PartitionSpec | None = None + attn_qk_activation: PartitionSpec | None = None + fc1_kernel: PartitionSpec | None = None + fc1_bias: PartitionSpec | None = None + fc2_kernel: PartitionSpec | None = None + fc2_bias: PartitionSpec | None = None + activation: PartitionSpec | None = None + layer_norm: PartitionSpec | None = None + emb_patch_kernel: PartitionSpec | None = None + emb_patch_bias: PartitionSpec | None = None + emb_pos_kernel: PartitionSpec | None = None + + @staticmethod + def no_sharding(): + return VisionShardingCfg() + + @staticmethod + def default(use_fsdp: bool, use_tp: bool): + fsdp = ShardMode.FSDP.value if use_fsdp else None + tp = ShardMode.TP.value if use_tp else None + return VisionShardingCfg( + attn_kernel=P(tp, fsdp), + attn_bias=P(tp), + attn_qk_activation=P(fsdp, tp), + fc1_kernel=P(fsdp, tp), + fc1_bias=P(tp), + fc2_kernel=P(tp, fsdp), + fc2_bias=P(tp), + activation=P(fsdp, None, tp), + layer_norm=P(tp), + emb_patch_kernel=P(None, None, None, tp), + emb_patch_bias=P(tp), + emb_pos_kernel=P(None, tp), + ) + + +@dataclass(slots=True, frozen=True) +class TextShardingCfg: + attn_kernel: PartitionSpec | None = None + attn_bias: PartitionSpec | None = None + attn_qk_activation: PartitionSpec | None = None + down_kernel: PartitionSpec | None = None + down_bias: PartitionSpec | None = None + up_gate_kernel: PartitionSpec | None = None + up_gate_bias: PartitionSpec | None = None + activation: PartitionSpec | None = None + norm: PartitionSpec | None = None + cache: PartitionSpec | None = None + emb_kernel: PartitionSpec | None = None + + @staticmethod + def no_sharding(): + return TextShardingCfg() + + @staticmethod + def default(use_fsdp: bool, use_tp: bool): + fsdp = ShardMode.FSDP.value if use_fsdp else None + tp = ShardMode.TP.value if use_tp else None + return TextShardingCfg( + attn_kernel=P(tp, fsdp), + attn_bias=P(tp), + attn_qk_activation=P(fsdp, None, tp), + down_kernel=P(tp, fsdp), + down_bias=P(tp), + up_gate_kernel=P(fsdp, tp), + up_gate_bias=P(tp), + activation=P(fsdp, None, tp), + norm=P(tp), + cache=P(fsdp, None, tp, None), + emb_kernel=P(None, tp), + ) + + +@dataclass(slots=True, frozen=True) +class MMShardingCfg: + mmp_norm: PartitionSpec | None = None + mmp_weight: PartitionSpec | None = None + + @staticmethod + def no_sharding(): + return MMShardingCfg() + + @staticmethod + def default(use_tp: bool): + tp = ShardMode.TP.value if use_tp else None + return MMShardingCfg(mmp_norm=P(tp), mmp_weight=P(tp)) + + @dataclass(frozen=True) class RoPEParameters: rope_type: str = "default" @@ -74,6 +166,7 @@ class T5Gemma2VisionConfig: num_heads: int = 16 posemb: str = "learn" dropout: float = 0.0 + shd_cfg: VisionShardingCfg = field(default_factory=VisionShardingCfg) @dataclass(frozen=True) @@ -93,6 +186,7 @@ class T5Gemma2TextConfig: bos_token_id: int = 2 eos_token_id: int = 1 attn_logit_softcapping: float | None = None + shd_cfg: TextShardingCfg = field(default_factory=TextShardingCfg) @functools.cached_property def query_pre_attn_scalar(self) -> float: @@ -110,7 +204,7 @@ class T5Gemma2EncoderConfig: @dataclass(frozen=True) class T5Gemma2DecoderConfig(T5Gemma2TextConfig): - """Decoder config - inherits all fields from T5Gemma2TextConfig.""" + pass @dataclass(frozen=True) @@ -119,6 +213,7 @@ class T5Gemma2Config: decoder: T5Gemma2DecoderConfig eoi_token_index: int = 256000 pad_token_id: int = 0 + shd_cfg: MMShardingCfg = field(default_factory=MMShardingCfg) @classmethod def _from_params( @@ -130,8 +225,20 @@ def _from_params( num_key_value_heads: int, sliding_window: int, with_vision: bool = True, + use_fsdp: bool = False, + use_tp: bool = False, ) -> "T5Gemma2Config": layer_types = _make_layer_types(num_layers) + + if use_fsdp or use_tp: + text_shd = TextShardingCfg.default(use_fsdp, use_tp) + vision_shd = VisionShardingCfg.default(use_fsdp, use_tp) + mm_shd = MMShardingCfg.default(use_tp) + else: + text_shd = TextShardingCfg.no_sharding() + vision_shd = VisionShardingCfg.no_sharding() + mm_shd = MMShardingCfg.no_sharding() + text_cfg = T5Gemma2TextConfig( num_hidden_layers=num_layers, embed_dim=embed_dim, @@ -140,8 +247,9 @@ def _from_params( num_key_value_heads=num_key_value_heads, sliding_window=sliding_window, layer_types=layer_types, + shd_cfg=text_shd, ) - vision_cfg = T5Gemma2VisionConfig() if with_vision else None + vision_cfg = T5Gemma2VisionConfig(shd_cfg=vision_shd) if with_vision else None return cls( encoder=T5Gemma2EncoderConfig(text_config=text_cfg, vision_config=vision_cfg), decoder=T5Gemma2DecoderConfig( @@ -152,30 +260,28 @@ def _from_params( num_key_value_heads=num_key_value_heads, sliding_window=sliding_window, layer_types=layer_types, + shd_cfg=text_shd, ), + shd_cfg=mm_shd, ) @classmethod - def t5gemma2_270m_270m(cls, with_vision: bool = True) -> "T5Gemma2Config": - return cls._from_params(18, 640, 2048, 4, 1, 512, with_vision) + def t5gemma2_270m_270m( + cls, with_vision: bool = True, use_fsdp: bool = False, use_tp: bool = False + ) -> "T5Gemma2Config": + return cls._from_params(18, 640, 2048, 4, 1, 512, with_vision, use_fsdp, use_tp) @classmethod - def t5gemma2_1b_1b(cls, with_vision: bool = True) -> "T5Gemma2Config": - return cls._from_params(26, 1152, 6912, 4, 1, 512, with_vision) + def t5gemma2_1b_1b(cls, with_vision: bool = True, use_fsdp: bool = False, use_tp: bool = False) -> "T5Gemma2Config": + return cls._from_params(26, 1152, 6912, 4, 1, 512, with_vision, use_fsdp, use_tp) @classmethod - def t5gemma2_4b_4b(cls, with_vision: bool = True) -> "T5Gemma2Config": - return cls._from_params(34, 2560, 10240, 8, 4, 1024, with_vision) - + def t5gemma2_4b_4b(cls, with_vision: bool = True, use_fsdp: bool = False, use_tp: bool = False) -> "T5Gemma2Config": + return cls._from_params(34, 2560, 10240, 8, 4, 1024, with_vision, use_fsdp, use_tp) -# ============================================================================= -# Constants -# ============================================================================= -# Large negative number for masking in attention -K_MASK = -2.3819763e38 +_K_MASK = jnp.finfo(jnp.bfloat16).min -# Special tokens BOS_TOKEN = 2 EOS_TOKEN = 1 NEW_LINE_TOKEN = 108 @@ -185,30 +291,14 @@ def t5gemma2_4b_4b(cls, with_vision: bool = True) -> "T5Gemma2Config": IMAGE_PLACEHOLDER_TOKEN = 256001 # Placeholder for image. Different from Gemma3. NUM_PLACEHOLDER_TOKENS_PER_IMAGE = 256 -# Default initializers NORMAL_INIT = nnx.initializers.normal() ZEROS_INIT = nnx.initializers.zeros_init() -# ============================================================================= -# Helper Functions -# ============================================================================= - - def _get_rope_base_frequency( text_config: T5Gemma2TextConfig, attn_type: AttentionType, ) -> int: - """Get RoPE base frequency for a given attention type. - - Args: - text_config: Text configuration containing rope_parameters. - attn_type: The attention type (GLOBAL or LOCAL_SLIDING). - - Returns: - The base frequency for RoPE. - """ - # Map attention type to rope parameter key if attn_type == AttentionType.GLOBAL: key = "full_attention" else: @@ -217,26 +307,13 @@ def _get_rope_base_frequency( rope_params = text_config.rope_parameters.get(key) if rope_params is not None: return int(rope_params.rope_theta) - return 10_000 # Default + return 10_000 def _get_rope_scale_factor( text_config: T5Gemma2TextConfig, attn_type: AttentionType, ) -> float: - """Get RoPE scale factor for a given attention type. - - For "linear" rope_type, this returns the factor from config. - For "default" rope_type, returns 1.0 (no scaling). - - Args: - text_config: Text configuration containing rope_parameters. - attn_type: The attention type (GLOBAL or LOCAL_SLIDING). - - Returns: - The scale factor for RoPE. - """ - # Map attention type to rope parameter key if attn_type == AttentionType.GLOBAL: key = "full_attention" else: @@ -244,10 +321,9 @@ def _get_rope_scale_factor( rope_params = text_config.rope_parameters.get(key) if rope_params is not None: - # "linear" rope_type uses factor, "default" uses 1.0 if rope_params.rope_type == "linear": return rope_params.factor - return 1.0 # Default (no scaling) + return 1.0 def apply_rope( @@ -258,18 +334,6 @@ def apply_rope( scale_factor: float = 1.0, rope_proportion: float = 1.0, ) -> Float[Array, "B L N H"]: # noqa: F722 - """Applies Rotary Position Embeddings (RoPE). - - Args: - inputs: Array of shape [B, L, N, H]. - positions: Array of shape [B, L]. - base_frequency: Base frequency used to compute rotations. - scale_factor: Scale factor for positional interpolation. - rope_proportion: Proportion of head dimension to apply RoPE to. - - Returns: - Array of shape [B, L, N, H]. - """ head_dim = inputs.shape[-1] rope_angles = int(rope_proportion * head_dim // 2) nope_angles = head_dim // 2 - rope_angles @@ -297,14 +361,7 @@ def apply_rope( return out.astype(inputs.dtype) -# ============================================================================= -# Base Layers -# ============================================================================= - - class T5Gemma2Einsum(nnx.Module): - """Parameterized einsum layer.""" - def __init__( self, shape: tuple[int, ...], @@ -315,13 +372,12 @@ def __init__( ): self.w = nnx.Param(kernel_init(rngs.params(), shape, dtype)) - def __call__(self, eqn: str, x: jax.Array) -> jax.Array: - return jnp.einsum(eqn, x, self.w.value) + @jax.named_scope("einsum") + def __call__(self, eqn: str, x: jax.Array, *, out_sharding=None) -> jax.Array: + return jnp.einsum(eqn, x, self.w[...], out_sharding=out_sharding) class T5Gemma2RMSNorm(nnx.Module): - """RMS Normalization layer.""" - def __init__( self, features: int, @@ -334,29 +390,22 @@ def __init__( self.epsilon = epsilon self.scale = nnx.Param(scale_init(rngs.params(), (features,), dtype)) + @jax.named_scope("rms_norm") def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: # noqa: F722 var = jnp.mean(jnp.square(x), axis=-1, keepdims=True) - - # Jax.lax.rsqrt is used because it returns different floats than - # jnp.reciprocal(jnp.sqrt(var + 1e-06)) normed_inputs = x * jax.lax.rsqrt(var + self.epsilon) - - # normed_inputs is a rank-K tensor, K > 1 (K is typically 2 or 3). scale is - # a rank-1 tensor. To avoid implicit rank-promotion, reshape scale to - # a (1, ..., 1, D) tensor, so the rank of scale matches normed_inputs. - scale = jnp.expand_dims(self.scale.value, axis=range(len(x.shape) - 1)) + scale = jnp.expand_dims(self.scale[...], axis=range(len(x.shape) - 1)) return normed_inputs * (1 + scale) class T5Gemma2FeedForward(nnx.Module): - """Feed-forward module with gated activation.""" - def __init__( self, features: int, hidden_dim: int, *, transpose_gating_einsum: bool, + shd_cfg: TextShardingCfg | None = None, kernel_init: nnx.Initializer = NORMAL_INIT, dtype: jnp.dtype = jnp.float32, rngs: nnx.Rngs, @@ -364,6 +413,7 @@ def __init__( self.features = features self.hidden_dim = hidden_dim self.transpose_gating_einsum = transpose_gating_einsum + self.shd_cfg = shd_cfg if transpose_gating_einsum: gating_shape = (2, hidden_dim, features) @@ -384,38 +434,18 @@ def __init__( rngs=rngs, ) + @jax.named_scope("feed_forward") def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: # noqa: F722 - """Applies feed-forward transformation. - - Args: - x: Input of shape [batch_size, seq_len, features]. - - Returns: - Output of shape [batch_size, seq_len, features]. - """ eq = "...F,NHF->...NH" if self.transpose_gating_einsum else "...F,NFH->...NH" gate = self.gating_einsum(eq, x) activations = jax.nn.gelu(gate[..., 0, :]) * gate[..., 1, :] - return self.linear("...H,HF->...F", activations) - - -# ============================================================================= -# Attention Masks -# ============================================================================= + shd = self.shd_cfg.activation if self.shd_cfg is not None else None + return self.linear("...H,HF->...F", activations, out_sharding=shd) def make_bidirectional_mask( input_mask: Bool[Array, "B L"], # noqa: F722 ) -> Bool[Array, "B 1 L L"]: # noqa: F722 - """Creates a bidirectional attention mask for encoder. - - Args: - input_mask: Boolean mask where True indicates valid tokens. - - Returns: - Attention mask of shape [B, 1, L, L]. - """ - # [B, L] -> [B, 1, 1, L] mask = input_mask[:, None, None, :] return mask @@ -423,18 +453,8 @@ def make_bidirectional_mask( def make_causal_mask( input_mask: Bool[Array, "B L"], # noqa: F722 ) -> Bool[Array, "B 1 L L"]: # noqa: F722 - """Creates a causal attention mask for decoder. - - Args: - input_mask: Boolean mask where True indicates valid tokens. - - Returns: - Causal attention mask of shape [B, 1, L, L]. - """ seq_len = input_mask.shape[-1] - # Create causal mask: [L, L] causal = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) - # Combine with input mask: [B, 1, L, L] mask = input_mask[:, None, None, :] & causal[None, None, :, :] return mask @@ -445,26 +465,10 @@ def make_sliding_window_mask( *, bidirectional: bool = False, ) -> Bool[Array, "B 1 L L"]: # noqa: F722 - """Creates a sliding window mask. - - Args: - positions: Position indices of shape [B, L]. - sliding_window: Size of the sliding window. - bidirectional: If True, creates a symmetric bidirectional window where - the total window size equals `sliding_window` (split evenly on both - sides). If False, creates a causal window of size `sliding_window`. - - Returns: - Sliding window mask of shape [B, 1, L, L]. - """ - # [B, L, 1] and [B, 1, L] q_pos = positions[:, :, None] k_pos = positions[:, None, :] if bidirectional: - # Bidirectional: symmetric window centered on query position - # Total window size = sliding_window, split evenly on both sides - # Matches PyTorch: left_window = (sw + 1) // 2, right_window = sw // 2 + 1 left_window = (sliding_window + 1) // 2 right_window = sliding_window // 2 + 1 dist = q_pos - k_pos @@ -472,7 +476,6 @@ def make_sliding_window_mask( right_mask = (dist < 0) & (-dist < right_window) mask = left_mask | right_mask else: - # Causal: can only attend to past tokens within window dist = jnp.abs(q_pos - k_pos) mask = dist < sliding_window @@ -484,25 +487,8 @@ def make_sliding_window_causal_mask( positions: Int[Array, "B L"], # noqa: F722 sliding_window: int, ) -> Bool[Array, "B 1 L L"]: # noqa: F722 - """Creates a causal mask with sliding window for decoder. - - Combines causal mask (lower triangular) with sliding window constraint. - - Args: - input_mask: Boolean mask where True indicates valid tokens. - positions: Position indices of shape [B, L]. - sliding_window: Size of the sliding window. - - Returns: - Sliding window causal mask of shape [B, 1, L, L]. - """ - # Start with causal mask causal_mask = make_causal_mask(input_mask) - - # Create sliding window mask sliding_mask = make_sliding_window_mask(positions, sliding_window) - - # Combine: must satisfy both causal AND sliding window return causal_mask & sliding_mask @@ -510,25 +496,9 @@ def make_merged_attention_mask( decoder_mask: Bool[Array, "B 1 L_dec L_dec"], # noqa: F722 encoder_mask: Bool[Array, "B 1 1 L_enc"], # noqa: F722 ) -> Bool[Array, "B 1 L_dec L_combined"]: # noqa: F722 - """Creates merged attention mask for decoder's merged attention. - - Concatenates the decoder self-attention mask with the cross-attention - mask to form a single mask for the merged attention operation. - - Args: - decoder_mask: Causal mask for decoder self-attention [B, 1, L_dec, L_dec]. - encoder_mask: Mask for encoder hidden states [B, 1, 1, L_enc]. - - Returns: - Merged mask of shape [B, 1, L_dec, L_dec + L_enc]. - """ batch_size, _, seq_len, _ = decoder_mask.shape enc_len = encoder_mask.shape[-1] - - # Broadcast encoder mask to [B, 1, L_dec, L_enc] cross_mask = jnp.broadcast_to(encoder_mask, (batch_size, 1, seq_len, enc_len)) - - # Concatenate along key dimension return jnp.concatenate([decoder_mask, cross_mask], axis=-1) @@ -538,20 +508,6 @@ def make_decode_mode_self_mask( max_cache_len: int, current_cache_len: int, ) -> Bool[Array, "B 1 Q K"]: # noqa: F722 - """Creates self-attention mask for decode mode (with KV cache). - - During decode mode, the query attends to all previously cached keys. - We mask out positions that haven't been written to yet (>= current_cache_len). - - Args: - batch_size: Batch size. - query_len: Number of query positions (typically 1 for decode mode). - max_cache_len: Total allocated size of the cache. - current_cache_len: Current valid length of the cache. - - Returns: - Mask of shape [B, 1, query_len, max_cache_len]. - """ k_pos = jnp.arange(max_cache_len) mask = k_pos < current_cache_len return jnp.broadcast_to(mask[None, None, None, :], (batch_size, 1, query_len, max_cache_len)) @@ -563,42 +519,14 @@ def make_decode_mode_sliding_mask( max_cache_len: int, sliding_window: int, ) -> Bool[Array, "B 1 1 K"]: # noqa: F722 - """Creates sliding window mask for decode mode (with KV cache). - - During decode mode with sliding window, the query can only attend to - positions within the sliding window range. - - Args: - batch_size: Batch size. - query_pos: Position of the current query (typically cache_index). - max_cache_len: Total allocated size of the cache. - sliding_window: Size of the sliding window. - - Returns: - Mask of shape [B, 1, 1, max_cache_len]. - """ - # Key positions: 0, 1, 2, ..., max_cache_len-1 k_positions = jnp.arange(max_cache_len) - # Within sliding window: |query_pos - k_pos| < sliding_window in_window = jnp.abs(query_pos - k_positions) < sliding_window - # Also must be causal (k_pos <= query_pos) is_causal = k_positions <= query_pos mask = in_window & is_causal return jnp.broadcast_to(mask[None, None, None, :], (batch_size, 1, 1, max_cache_len)) -# ============================================================================= -# Text Embeddings -# ============================================================================= - - class T5Gemma2ScaledWordEmbedding(nnx.Module): - """Scaled word embedding with special EOI (end-of-image) token. - - The embeddings are scaled by the provided embed_scale, and the EOI token - has a separate learnable embedding. - """ - def __init__( self, vocab_size: int, @@ -612,7 +540,6 @@ def __init__( self.vocab_size = vocab_size self.embed_dim = embed_dim self.padding_idx = padding_idx - # If embed_scale is not provided, default to sqrt(embed_dim) if embed_scale is None: embed_scale = embed_dim**0.5 self.embed_scale = jnp.array(embed_scale, dtype=dtype) @@ -625,27 +552,11 @@ def __init__( ) def __call__(self, input_ids: Int[Array, "B L"]) -> Float[Array, "B L D"]: # noqa: F722 - """Embed input tokens with scaling. - - Args: - input_ids: Input token IDs of shape [B, L]. - - Returns: - Embeddings of shape [B, L, hidden_size]. - """ embeddings = self.embedding[input_ids] * self.embed_scale - - # Replace EOI token embeddings eoi_mask = input_ids == END_OF_IMAGE_TOKEN - return jnp.where(eoi_mask[..., None], self.eoi_embedding[...], embeddings) -# ============================================================================= -# Vision Components -# ============================================================================= - - def _posemb_sincos_2d( h: int, w: int, @@ -654,7 +565,6 @@ def _posemb_sincos_2d( temperature: float = 10_000.0, dtype: jnp.dtype = jnp.float32, ) -> Float[Array, "1 M D"]: # noqa: F722 - """Sinusoidal 2D position embeddings (MoCo v3 style).""" y, x = jnp.mgrid[:h, :w] assert width % 4 == 0, "Width must be mult of 4 for sincos posemb" @@ -667,8 +577,6 @@ def _posemb_sincos_2d( class T5Gemma2VisionMLP(nnx.Module): - """MLP for Vision Transformer.""" - def __init__( self, *, @@ -709,8 +617,6 @@ def __call__( class T5Gemma2VisionEncoderBlock(nnx.Module): - """Transformer encoder block for Vision Transformer.""" - def __init__( self, *, @@ -775,8 +681,6 @@ def __call__( class T5Gemma2VisionEncoder(nnx.Module): - """Vision Transformer Encoder.""" - def __init__( self, *, @@ -823,7 +727,7 @@ def __call__( class T5Gemma2VisionExit(nnx.Module): - """Vision exit layer - spatially pools soft tokens to output length.""" + """Spatially pools soft tokens to output length.""" def __init__(self, output_length: int = 256, *, rngs: nnx.Rngs): self.output_length = output_length @@ -850,11 +754,7 @@ def __call__(self, x: jax.Array) -> jax.Array: class T5Gemma2VisionSoftTokenizer(nnx.Module): - """Vision soft tokenizer (ViT trained with SigLiP). - - Transforms images into soft tokens that can be embedded into the - text embedding space. - """ + """Vision soft tokenizer (ViT trained with SigLiP).""" def __init__( self, @@ -917,7 +817,7 @@ def __call__( bn, h, w, c = x.shape x = jnp.reshape(x, [bn, h * w, c]) - x = x + self.pos_embedding.value + x = x + self.pos_embedding[...] x = self.dropout(x, deterministic=deterministic) x = self.transformer(x, deterministic=deterministic) x = self.vision_exit(x) @@ -933,7 +833,6 @@ def _get_posemb( width: int, dtype: jnp.dtype = jnp.float32, ) -> nnx.Param: - """Returns the position embedding.""" if typ == "learn": shape = (1, seqshape[0] * seqshape[1], width) initializer = nnx.initializers.normal(stddev=1 / (width**0.5)) @@ -953,9 +852,11 @@ def __init__( embed_dim: int, *, soft_tokens_dim: int, + mm_shd_cfg: MMShardingCfg | None = None, dtype: jnp.dtype = jnp.float32, rngs: nnx.Rngs, ): + self.mm_shd_cfg = mm_shd_cfg self.mm_soft_embedding_norm = T5Gemma2RMSNorm( soft_tokens_dim, scale_init=nnx.initializers.zeros_init(), @@ -971,7 +872,8 @@ def __init__( def __call__(self, x: Float[Array, "B N P Dv"]) -> Float[Array, "B N P De"]: # noqa: F722 x = self.mm_soft_embedding_norm(x) - return self.mm_input_projection("...tm,md->...td", x) + shd = self.mm_shd_cfg.mmp_weight if self.mm_shd_cfg is not None else None + return self.mm_input_projection("...tm,md->...td", x, out_sharding=shd) class T5Gemma2VisionEmbedder(nnx.Module): @@ -982,6 +884,7 @@ def __init__( *, vision_config: T5Gemma2VisionConfig, embed_dim: int, + mm_shd_cfg: MMShardingCfg | None = None, freeze_params: bool = True, dtype: jnp.dtype = jnp.float32, rngs: nnx.Rngs, @@ -997,6 +900,7 @@ def __init__( self.soft_tokens_embedder = T5Gemma2VisionSoftTokensEmbedder( embed_dim, soft_tokens_dim=vision_config.width, + mm_shd_cfg=mm_shd_cfg, dtype=dtype, rngs=rngs, ) @@ -1015,15 +919,11 @@ def __call__( return self.soft_tokens_embedder(soft_tokens) -# ============================================================================= -# Multimodal Utilities -# ============================================================================= def merge_mm_embeddings( text_embeddings: jnp.ndarray, vision_embeddings: jnp.ndarray, mask: jnp.ndarray, ) -> jnp.ndarray: - """Merge text and multimodal (vision) embeddings.""" return jax.vmap(_merge_mm_embeddings_inner, in_axes=(0, 0, 0))( text_embeddings, vision_embeddings, @@ -1036,7 +936,6 @@ def _merge_mm_embeddings_inner( vision_embeddings: jnp.ndarray, mask: jnp.ndarray, ) -> jnp.ndarray: - """Merge embeddings without batch dimension.""" num_images, num_toks_per_image, d = vision_embeddings.shape vision_embeddings = jnp.reshape( vision_embeddings, @@ -1049,18 +948,7 @@ def _merge_mm_embeddings_inner( return merged.at[0].set(first_pos) -# ============================================================================= -# Encoder Components -# ============================================================================= - - class T5Gemma2EncoderAttention(nnx.Module): - """Bidirectional self-attention for T5Gemma2 encoder. - - This is separate from Gemma3Attention because the encoder needs - bidirectional sliding window attention, not causal sliding window. - """ - def __init__( self, num_q_heads: int, @@ -1070,6 +958,7 @@ def __init__( attn_type: AttentionType, query_pre_attn_scalar: float, *, + shd_cfg: TextShardingCfg | None = None, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, rms_norm_eps: float = 1e-6, @@ -1093,6 +982,7 @@ def __init__( self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.use_qk_norm = use_qk_norm + self.shd_cfg = shd_cfg self.attn_vec_einsum = T5Gemma2Einsum( shape=(num_q_heads, head_dim, features), @@ -1101,7 +991,6 @@ def __init__( rngs=rngs, ) - # Check if we can use combined QKV projection if num_kv_heads == num_q_heads: self.qkv_einsum = T5Gemma2Einsum( shape=(3, num_q_heads, features, head_dim), @@ -1150,35 +1039,25 @@ def use_qkv_einsum(self) -> bool: def use_gqa(self) -> bool: return self.num_kv_heads != self.num_q_heads and self.num_kv_heads > 1 + @jax.named_scope("encoder_attention") def __call__( self, x: Float[Array, "B L D"], # noqa: F722 segment_pos: Float[Array, "B L"], # noqa: F722 attn_mask: Float[Array, "B L L"], # noqa: F722 ) -> Float[Array, "B L D"]: # noqa: F722 - """Applies bidirectional multi-head attention. + shd = self.shd_cfg.activation if self.shd_cfg is not None else None - Args: - x: Input sequence of shape [batch_size, seq_len, embed_dim]. - segment_pos: Absolute positions of shape [batch_size, seq_len]. - attn_mask: Attention mask of shape [batch_size, seq_len, seq_len]. - - Returns: - Output sequence of shape [batch_size, seq_len, embed_dim]. - """ - # Project Q, K, V if self.use_qkv_einsum: query_proj, key_proj, value_proj = self.qkv_einsum("BTD,SNDH->SBTNH", x) else: - query_proj = self.q_einsum("BTD,NDH->BTNH", x) + query_proj = self.q_einsum("BTD,NDH->BTNH", x, out_sharding=shd) key_proj, value_proj = self.kv_einsum("BSD,CKDH->CBSKH", x) - # Apply QK normalization if self.use_qk_norm: query_proj = self._query_norm(query_proj) key_proj = self._key_norm(key_proj) - # Apply RoPE query_proj = apply_rope( query_proj, segment_pos, @@ -1194,7 +1073,6 @@ def __call__( scale_factor=self.rope_scale_factor, ) - # Compute attention logits if self.use_gqa: b, t, kg, h = query_scaled.shape query_scaled = query_scaled.reshape( @@ -1206,18 +1084,15 @@ def __call__( else: logits = jnp.einsum("BTNH,BSNH->BTNS", query_scaled, key_proj) - # Apply soft capping if self.attn_logits_soft_cap is not None: logits = jnp.tanh(logits / self.attn_logits_soft_cap) logits = logits * self.attn_logits_soft_cap - # Apply bidirectional sliding window mask for LOCAL_SLIDING layers if self.attn_type == AttentionType.LOCAL_SLIDING: if self.sliding_window_size is None: raise ValueError( "sliding_window_size must be set for LOCAL_SLIDING attention", ) - # Use make_sliding_window_mask with bidirectional=True, squeeze to [B, L, L] sliding_mask = make_sliding_window_mask( segment_pos, self.sliding_window_size, @@ -1225,11 +1100,9 @@ def __call__( )[:, 0, :, :] attn_mask = attn_mask * sliding_mask - # Apply mask and softmax - padded_logits = jnp.where(jnp.expand_dims(attn_mask, -2), logits, K_MASK) + padded_logits = jnp.where(jnp.expand_dims(attn_mask, -2), logits, _K_MASK) probs = jax.nn.softmax(padded_logits, axis=-1).astype(key_proj.dtype) - # Compute attention output if self.use_gqa: b, t, kg, s = probs.shape probs = probs.reshape( @@ -1241,16 +1114,10 @@ def __call__( else: encoded = jnp.einsum("BTNS,BSNH->BTNH", probs, value_proj) - return self.attn_vec_einsum("BTNH,NHD->BTD", encoded) + return self.attn_vec_einsum("BTNH,NHD->BTD", encoded, out_sharding=shd) class T5Gemma2EncoderBlock(nnx.Module): - """Encoder transformer block with bidirectional attention. - - Uses T5Gemma2EncoderAttention which handles bidirectional sliding window - correctly (unlike Gemma3Attention which uses causal sliding window). - """ - def __init__( self, num_q_heads: int, @@ -1264,6 +1131,7 @@ def __init__( attn_type: AttentionType, query_pre_attn_scalar: float, transpose_gating_einsum: bool, + shd_cfg: TextShardingCfg | None = None, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, rms_norm_eps: float = 1e-6, @@ -1293,6 +1161,7 @@ def __init__( num_kv_heads=num_kv_heads, attn_type=attn_type, query_pre_attn_scalar=query_pre_attn_scalar, + shd_cfg=shd_cfg, rope_base_frequency=rope_base_frequency, rope_scale_factor=rope_scale_factor, rms_norm_eps=rms_norm_eps, @@ -1325,6 +1194,7 @@ def __init__( features=embed_dim, hidden_dim=hidden_dim, transpose_gating_einsum=transpose_gating_einsum, + shd_cfg=shd_cfg, kernel_init=kernel_init, dtype=dtype, rngs=rngs, @@ -1339,23 +1209,13 @@ def __init__( rngs=rngs, ) + @jax.named_scope("encoder_block") def __call__( self, x: Float[Array, "B L D"], # noqa: F722 segment_pos: Int[Array, "B L"], # noqa: F722 attn_mask: Bool[Array, "B L L"], # noqa: F722 ) -> Float[Array, "B L D"]: # noqa: F722 - """Apply encoder layer. - - Args: - x: Input hidden states [B, L, D]. - segment_pos: Position indices [B, L]. - attn_mask: Bidirectional attention mask [B, L, L]. - - Returns: - Output hidden states [B, L, D]. - """ - # Attention block inputs_normalized = self.pre_attention_norm(x) attn_output = self.attn(inputs_normalized, segment_pos, attn_mask) @@ -1364,7 +1224,6 @@ def __call__( attn_output = attn_output + x - # Feed-forward block outputs = self.pre_ffw_norm(attn_output) outputs = self.mlp(outputs) @@ -1375,15 +1234,11 @@ def __call__( class T5Gemma2Encoder(nnx.Module): - """T5Gemma2 Encoder with optional vision support. - - Processes text (and optional images) using bidirectional self-attention. - """ - def __init__( self, config: T5Gemma2EncoderConfig, *, + mm_shd_cfg: MMShardingCfg | None = None, embedder: nnx.Module | None = None, dtype: jnp.dtype = jnp.float32, dtype_mm: jnp.dtype = jnp.float32, @@ -1391,8 +1246,8 @@ def __init__( ): self.config = config text_config = config.text_config + shd_cfg = text_config.shd_cfg if text_config.shd_cfg.activation is not None else None - # Text embeddings self.embedder = embedder or T5Gemma2ScaledWordEmbedding( vocab_size=text_config.vocab_size, embed_dim=text_config.embed_dim, @@ -1401,23 +1256,21 @@ def __init__( rngs=rngs, ) - # Vision encoder (optional) - uses Gemma3 vision component self.vision_embedder = nnx.data(None) if config.vision_config is not None: self.vision_embedder = T5Gemma2VisionEmbedder( vision_config=config.vision_config, embed_dim=text_config.embed_dim, + mm_shd_cfg=mm_shd_cfg, freeze_params=True, dtype=dtype_mm, rngs=rngs, ) - # Determine attention types per layer attention_types = text_config.layer_types if not attention_types: attention_types = tuple(AttentionType.GLOBAL for _ in range(text_config.num_hidden_layers)) - # Encoder layers (use Gemma3Block via alias) self.blocks = nnx.List( [ T5Gemma2EncoderBlock( @@ -1431,6 +1284,7 @@ def __init__( attn_type=attention_types[i], query_pre_attn_scalar=text_config.query_pre_attn_scalar, transpose_gating_einsum=True, + shd_cfg=shd_cfg, rope_base_frequency=_get_rope_base_frequency( text_config, attention_types[i], @@ -1452,7 +1306,6 @@ def __init__( ] ) - # Final normalization self.norm = T5Gemma2RMSNorm( text_config.embed_dim, epsilon=text_config.rms_norm_eps, @@ -1469,32 +1322,16 @@ def __call__( *, deterministic: bool = True, ) -> Float[Array, "B L D"]: # noqa: F722 - """Forward pass of the encoder. - - Args: - tokens: Input token IDs [B, L]. - attention_mask: Attention mask [B, L]. - position_ids: Position indices [B, L]. - images: Optional input images [B, N, H, W, C] where N is images per batch. - deterministic: Whether to run in deterministic mode. - - Returns: - Encoder hidden states [B, L, D]. - """ batch_size, seq_len = tokens.shape - # Create position IDs if not provided if position_ids is None: position_ids = jnp.arange(seq_len)[None, :].repeat(batch_size, axis=0) - # Create attention mask if not provided if attention_mask is None: attention_mask = jnp.ones((batch_size, seq_len), dtype=jnp.bool_) - # Embed tokens x = self.embedder(tokens) - # Process images and merge with text embeddings if images is not None and self.vision_embedder is not None: image_embeddings = self.vision_embedder(images, deterministic=deterministic) x = merge_mm_embeddings( @@ -1503,42 +1340,22 @@ def __call__( tokens == IMAGE_PLACEHOLDER_TOKEN, ) - # Create bidirectional attention mask for encoder [B, L, L] - # Each query can attend to all valid (non-padded) key positions attn_mask = make_bidirectional_mask(attention_mask) - # Squeeze to [B, L, L] if it's [B, 1, L, L] if attn_mask.ndim == 4: attn_mask = attn_mask[:, 0, :, :] - # Ensure mask is [B, L, L] by broadcasting if needed if attn_mask.shape[-2] == 1: - # [B, 1, L] -> [B, L, L] attn_mask = jnp.broadcast_to(attn_mask, (batch_size, seq_len, seq_len)) - # Apply encoder layers - # T5Gemma2EncoderBlock handles bidirectional sliding window internally for block in self.blocks: x = block(x, position_ids, attn_mask) - # Final normalization return self.norm(x) -# ============================================================================= -# Decoder Components -# ============================================================================= - - class T5Gemma2MergedAttention(nnx.Module): """Merged self-attention and cross-attention for decoder. - Combines self-attention and cross-attention in a single operation: - - Query comes from decoder hidden states - - Keys/Values are concatenation of [self_kv, cross_kv] - - Mask is concatenation of [causal_mask, encoder_mask] - - This is more efficient than separate attention operations. - Uses nnx.Cache for autoregressive decoding. Call init_cache() before using decode=True. """ @@ -1551,6 +1368,7 @@ def __init__( head_dim: int, query_pre_attn_scalar: float, *, + shd_cfg: TextShardingCfg | None = None, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, rms_norm_eps: float = 1e-6, @@ -1573,10 +1391,10 @@ def __init__( self.attn_logits_soft_cap = attn_logits_soft_cap self.use_qk_norm = use_qk_norm self.decode = decode + self.shd_cfg = shd_cfg self.num_kv_groups = num_q_heads // num_kv_heads - # Projections self.q_proj = T5Gemma2Einsum( shape=(num_q_heads, features, head_dim), kernel_init=kernel_init, @@ -1602,7 +1420,6 @@ def __init__( rngs=rngs, ) - # QK normalization if use_qk_norm: self.q_norm = T5Gemma2RMSNorm( head_dim, @@ -1619,14 +1436,10 @@ def __init__( rngs=rngs, ) - # Cache for autoregressive decoding (nnx.Cache pattern like Gemma3) - # Self-attention cache self.cached_self_key: nnx.Cache[Array] | None = nnx.data(None) self.cached_self_value: nnx.Cache[Array] | None = nnx.data(None) - # Cross-attention cache (computed once from encoder outputs) self.cached_cross_key: nnx.Cache[Array] | None = nnx.data(None) self.cached_cross_value: nnx.Cache[Array] | None = nnx.data(None) - # Cache index for self-attention self.cache_index: nnx.Cache[Array] | None = nnx.data(None) def init_cache( @@ -1637,39 +1450,29 @@ def init_cache( encoder_seq_length: int, dtype: jnp.dtype = jnp.float32, ) -> None: - """Initialize KV caches for autoregressive decoding. - - Must be called before using decode=True. + cache_shd = self.shd_cfg.cache if self.shd_cfg is not None else None - Args: - batch_size: Number of sequences in the batch. - max_decode_length: Maximum decoder sequence length. - encoder_seq_length: Encoder sequence length (for cross-attention). - dtype: Data type for cache arrays. - """ - # Self-attention cache shape: [B, max_decode_len, num_kv_heads, head_dim] self_cache_shape = ( batch_size, max_decode_length, self.num_kv_heads, self.head_dim, ) - self.cached_self_key = nnx.Cache(jnp.zeros(self_cache_shape, dtype)) - self.cached_self_value = nnx.Cache(jnp.zeros(self_cache_shape, dtype)) + self.cached_self_key = nnx.Cache(jnp.zeros(self_cache_shape, dtype, out_sharding=cache_shd)) + self.cached_self_value = nnx.Cache(jnp.zeros(self_cache_shape, dtype, out_sharding=cache_shd)) - # Cross-attention cache shape: [B, enc_seq_len, num_kv_heads, head_dim] cross_cache_shape = ( batch_size, encoder_seq_length, self.num_kv_heads, self.head_dim, ) - self.cached_cross_key = nnx.Cache(jnp.zeros(cross_cache_shape, dtype)) - self.cached_cross_value = nnx.Cache(jnp.zeros(cross_cache_shape, dtype)) + self.cached_cross_key = nnx.Cache(jnp.zeros(cross_cache_shape, dtype, out_sharding=cache_shd)) + self.cached_cross_value = nnx.Cache(jnp.zeros(cross_cache_shape, dtype, out_sharding=cache_shd)) - # Cache index tracks current position in self-attention cache self.cache_index = nnx.Cache(jnp.array(0, dtype=jnp.int32)) + @jax.named_scope("merged_attention") def __call__( self, hidden_states: Float[Array, "B L_dec D"], # noqa: F722 @@ -1679,31 +1482,17 @@ def __call__( *, decode: bool | None = None, ) -> Float[Array, "B L_dec D"]: # noqa: F722 - """Apply merged self-attention and cross-attention. - - Args: - hidden_states: Decoder hidden states [B, L_dec, D]. - encoder_hidden_states: Encoder outputs [B, L_enc, D]. - position_ids: Decoder position indices [B, L_dec]. - merged_attention_mask: Merged mask [B, 1, L_dec, L_combined]. - decode: Whether to use KV cache for autoregressive decoding. - - Returns: - Output hidden states [B, L_dec, D]. - """ seq_len = hidden_states.shape[1] + shd = self.shd_cfg.activation if self.shd_cfg is not None else None - # Project decoder Q, K, V (self-attention) - query = self.q_proj("BTD,NDH->BTNH", hidden_states) - self_key = self.k_proj("BTD,NDH->BTNH", hidden_states) - self_value = self.v_proj("BTD,NDH->BTNH", hidden_states) + query = self.q_proj("BTD,NDH->BTNH", hidden_states, out_sharding=shd) + self_key = self.k_proj("BTD,NDH->BTNH", hidden_states, out_sharding=shd) + self_value = self.v_proj("BTD,NDH->BTNH", hidden_states, out_sharding=shd) - # QK normalization for self-attention if self.use_qk_norm: query = self.q_norm(query) self_key = self.k_norm(self_key) - # Apply RoPE to query and self-attention key query = apply_rope( query, position_ids, @@ -1717,14 +1506,12 @@ def __call__( scale_factor=self.rope_scale_factor, ) - # Determine decode mode decode = first_from( decode, self.decode, error_msg="No decode argument provided to T5Gemma2MergedAttention", ) - # Handle caching for autoregressive decoding if decode: if ( self.cached_self_key is None @@ -1740,7 +1527,6 @@ def __call__( cur_index = self.cache_index[...] slice_indices = (0, cur_index, 0, 0) - # Update self-attention cache using dynamic_update_slice self_key_updated = jax.lax.dynamic_update_slice( self.cached_self_key[...], self_key, @@ -1755,86 +1541,66 @@ def __call__( self.cached_self_key[...] = self_key_updated self.cached_self_value[...] = self_value_updated - # Use the full updated cache (static shape) for JIT compatibility self_key = self_key_updated self_value = self_value_updated - # Cross-attention: compute and cache on first call (when index is 0) - # After that, reuse cached values def compute_cross_kv(): - cross_k = self.k_proj("BTD,NDH->BTNH", encoder_hidden_states) - cross_v = self.v_proj("BTD,NDH->BTNH", encoder_hidden_states) + cross_k = self.k_proj("BTD,NDH->BTNH", encoder_hidden_states, out_sharding=shd) + cross_v = self.v_proj("BTD,NDH->BTNH", encoder_hidden_states, out_sharding=shd) if self.use_qk_norm: cross_k = self.k_norm(cross_k) return cross_k, cross_v - # Use lax.cond to conditionally compute cross-attention KV cross_key, cross_value = jax.lax.cond( cur_index == 0, compute_cross_kv, lambda: (self.cached_cross_key[...], self.cached_cross_value[...]), ) - # Update cross cache (will be no-op after first call due to cond) self.cached_cross_key[...] = cross_key self.cached_cross_value[...] = cross_value - # Update cache index self.cache_index[...] = cur_index + seq_len else: - # No caching - compute cross-attention KV directly - cross_key = self.k_proj("BTD,NDH->BTNH", encoder_hidden_states) - cross_value = self.v_proj("BTD,NDH->BTNH", encoder_hidden_states) + cross_key = self.k_proj("BTD,NDH->BTNH", encoder_hidden_states, out_sharding=shd) + cross_value = self.v_proj("BTD,NDH->BTNH", encoder_hidden_states, out_sharding=shd) if self.use_qk_norm: cross_key = self.k_norm(cross_key) - # Scale query query = query * self.query_pre_attn_scalar - # Concatenate self and cross KV key = jnp.concatenate([self_key, cross_key], axis=1) value = jnp.concatenate([self_value, cross_value], axis=1) - # Transpose to [B, N, T, H] for standard attention computation query = jnp.transpose(query, (0, 2, 1, 3)) # [B, N, T, H] key = jnp.transpose(key, (0, 2, 1, 3)) # [B, K, S, H] value = jnp.transpose(value, (0, 2, 1, 3)) # [B, K, S, H] - # Expand KV heads for GQA if self.num_kv_groups > 1: key = jnp.repeat(key, self.num_kv_groups, axis=1) value = jnp.repeat(value, self.num_kv_groups, axis=1) - # Compute attention scores: [B, N, T, S] attn_weights = jnp.einsum("BNTH,BNSH->BNTS", query, key) - # Apply softcapping if self.attn_logits_soft_cap is not None: attn_weights = jnp.tanh(attn_weights / self.attn_logits_soft_cap) attn_weights = attn_weights * self.attn_logits_soft_cap - # Apply attention mask if merged_attention_mask is not None: - attn_weights = jnp.where(merged_attention_mask, attn_weights, K_MASK) + attn_weights = jnp.where(merged_attention_mask, attn_weights, _K_MASK) - # Softmax attn_weights = jax.nn.softmax(attn_weights, axis=-1).astype(value.dtype) - # Apply attention to values: [B, N, T, H] attn_output = jnp.einsum("BNTS,BNSH->BNTH", attn_weights, value) - # Transpose back to [B, T, N, H] attn_output = jnp.transpose(attn_output, (0, 2, 1, 3)) - # Output projection - output = self.o_proj("BTNH,NHD->BTD", attn_output) + output = self.o_proj("BTNH,NHD->BTD", attn_output, out_sharding=shd) return output class T5Gemma2DecoderBlock(nnx.Module): - """Decoder transformer block with merged self/cross attention.""" - def __init__( self, num_q_heads: int, @@ -1847,6 +1613,7 @@ def __init__( use_post_attn_norm: bool = True, use_post_ffw_norm: bool = True, query_pre_attn_scalar: float, + shd_cfg: TextShardingCfg | None = None, rope_base_frequency: int = 10_000, rope_scale_factor: float = 1.0, rms_norm_eps: float = 1e-6, @@ -1862,7 +1629,6 @@ def __init__( self.use_post_attn_norm = use_post_attn_norm self.use_post_ffw_norm = use_post_ffw_norm - # Pre-attention norm self.pre_attention_norm = T5Gemma2RMSNorm( embed_dim, epsilon=rms_norm_eps, @@ -1871,13 +1637,13 @@ def __init__( rngs=rngs, ) - # Merged attention (self + cross) self.attn = T5Gemma2MergedAttention( num_q_heads=num_q_heads, num_kv_heads=num_kv_heads, features=embed_dim, head_dim=head_dim, query_pre_attn_scalar=query_pre_attn_scalar, + shd_cfg=shd_cfg, rope_base_frequency=rope_base_frequency, rope_scale_factor=rope_scale_factor, rms_norm_eps=rms_norm_eps, @@ -1890,7 +1656,6 @@ def __init__( rngs=rngs, ) - # Post-attention norm if use_post_attn_norm: self.post_attention_norm = T5Gemma2RMSNorm( embed_dim, @@ -1900,7 +1665,6 @@ def __init__( rngs=rngs, ) - # Pre-FFW norm self.pre_ffw_norm = T5Gemma2RMSNorm( embed_dim, epsilon=rms_norm_eps, @@ -1909,17 +1673,16 @@ def __init__( rngs=rngs, ) - # Feed-forward self.mlp = T5Gemma2FeedForward( features=embed_dim, hidden_dim=hidden_dim, transpose_gating_einsum=True, + shd_cfg=shd_cfg, kernel_init=kernel_init, dtype=dtype, rngs=rngs, ) - # Post-FFW norm if use_post_ffw_norm: self.post_ffw_norm = T5Gemma2RMSNorm( embed_dim, @@ -1929,6 +1692,7 @@ def __init__( rngs=rngs, ) + @jax.named_scope("decoder_block") def __call__( self, x: Float[Array, "B L_dec D"], # noqa: F722 @@ -1938,19 +1702,6 @@ def __call__( *, decode: bool | None = None, ) -> Float[Array, "B L_dec D"]: # noqa: F722 - """Apply decoder layer. - - Args: - x: Decoder hidden states [B, L_dec, D]. - encoder_hidden_states: Encoder outputs [B, L_enc, D]. - segment_pos: Decoder position indices [B, L_dec]. - attn_mask: Merged attention mask [B, 1, L_dec, L_combined]. - decode: Whether to use KV cache. - - Returns: - Output hidden states [B, L_dec, D]. - """ - # Merged attention block inputs_normalized = self.pre_attention_norm(x) attn_output = self.attn( inputs_normalized, @@ -1965,7 +1716,6 @@ def __call__( attn_output += x - # Feed-forward block outputs = self.pre_ffw_norm(attn_output) outputs = self.mlp(outputs) @@ -1978,12 +1728,6 @@ def __call__( class T5Gemma2Decoder(nnx.Module): - """T5Gemma2 Decoder with merged self/cross attention. - - Uses nnx.Cache for autoregressive decoding. Call init_cache() before - using decode=True. - """ - def __init__( self, config: T5Gemma2DecoderConfig, @@ -1995,8 +1739,8 @@ def __init__( ): self.config = config self.dtype = dtype + shd_cfg = config.shd_cfg if config.shd_cfg.activation is not None else None - # Text embeddings self.embedder = embedder or T5Gemma2ScaledWordEmbedding( vocab_size=config.vocab_size, embed_dim=config.embed_dim, @@ -2005,12 +1749,10 @@ def __init__( rngs=rngs, ) - # Determine attention types per layer attention_types = config.layer_types if not attention_types: attention_types = tuple(AttentionType.GLOBAL for _ in range(config.num_hidden_layers)) - # Decoder layers self.blocks = nnx.List( [ T5Gemma2DecoderBlock( @@ -2023,6 +1765,7 @@ def __init__( use_post_attn_norm=True, use_post_ffw_norm=True, query_pre_attn_scalar=config.query_pre_attn_scalar, + shd_cfg=shd_cfg, rope_base_frequency=_get_rope_base_frequency(config, attention_types[i]), rope_scale_factor=_get_rope_scale_factor(config, attention_types[i]), rms_norm_eps=config.rms_norm_eps, @@ -2036,7 +1779,6 @@ def __init__( ] ) - # Final normalization self.norm = T5Gemma2RMSNorm( config.embed_dim, epsilon=config.rms_norm_eps, @@ -2052,16 +1794,6 @@ def init_cache( encoder_seq_length: int, dtype: jnp.dtype | None = None, ) -> None: - """Initialize KV caches for all decoder layers. - - Must be called before using decode=True. - - Args: - batch_size: Number of sequences in the batch. - max_decode_length: Maximum decoder sequence length. - encoder_seq_length: Encoder sequence length (for cross-attention). - dtype: Data type for cache arrays. Defaults to model dtype. - """ if dtype is None: dtype = self.dtype @@ -2074,7 +1806,6 @@ def init_cache( ) def get_cache_index(self) -> int: - """Get current cache index from the first layer.""" if len(self.blocks) > 0 and self.blocks[0].attn.cache_index is not None: return int(self.blocks[0].attn.cache_index[...]) return 0 @@ -2089,27 +1820,12 @@ def __call__( *, decode: bool | None = None, ) -> Float[Array, "B L_dec D"]: # noqa: F722 - """Forward pass of the decoder. - - Args: - input_ids: Decoder input token IDs [B, L_dec]. - encoder_hidden_states: Encoder outputs [B, L_enc, D]. - attention_mask: Decoder attention mask [B, L_dec]. - encoder_attention_mask: Encoder attention mask [B, L_enc]. - position_ids: Decoder position indices [B, L_dec]. - decode: Whether to use KV cache. - - Returns: - Decoder hidden states [B, L_dec, D]. - """ batch_size, dec_seq_len = input_ids.shape enc_seq_len = encoder_hidden_states.shape[1] - # Create position IDs if not provided if position_ids is None: cache_index = self.get_cache_index() if decode and cache_index > 0: - # During autoregressive decoding, position is cache_index position_ids = jnp.full( (batch_size, dec_seq_len), cache_index, @@ -2121,7 +1837,6 @@ def __call__( axis=0, ) - # Create attention masks if not provided if attention_mask is None: attention_mask = jnp.ones((batch_size, dec_seq_len), dtype=jnp.bool_) if encoder_attention_mask is None: @@ -2130,57 +1845,43 @@ def __call__( dtype=jnp.bool_, ) - # Create encoder mask for cross-attention (bidirectional) encoder_mask = encoder_attention_mask[:, None, None, :] if decode: - # Decode mode: create masks accounting for cached sequence length cache_index = self.get_cache_index() - # Get max_decode_length (static) from the first block's cache max_decode_len = self.blocks[0].attn.cached_self_key.shape[1] - # Current valid length (dynamic) current_len = cache_index + dec_seq_len - # Full attention: new tokens can attend to all cached + encoder full_decoder_mask = make_decode_mode_self_mask( batch_size, dec_seq_len, max_decode_len, current_len, ) - # Sliding attention: limited to sliding window sliding_decoder_mask = make_decode_mode_sliding_mask( batch_size, - cache_index, # Current query position + cache_index, max_decode_len, self.config.sliding_window, ) - # Merged masks with encoder (broadcast encoder mask for query length) enc_mask_broadcast = jnp.broadcast_to(encoder_mask, (batch_size, 1, dec_seq_len, enc_seq_len)) merged_mask_full = jnp.concatenate([full_decoder_mask, enc_mask_broadcast], axis=-1) merged_mask_sliding = jnp.concatenate([sliding_decoder_mask, enc_mask_broadcast], axis=-1) else: - # Non-decode mode: standard causal masks - # Full attention: standard causal mask full_decoder_mask = make_causal_mask(attention_mask) - # Sliding attention: causal + sliding window sliding_decoder_mask = make_sliding_window_causal_mask( attention_mask, position_ids, self.config.sliding_window, ) - # Create merged masks for each attention type merged_mask_full = make_merged_attention_mask(full_decoder_mask, encoder_mask) merged_mask_sliding = make_merged_attention_mask(sliding_decoder_mask, encoder_mask) - # Embed tokens hidden_states = self.embedder(input_ids) - # Apply decoder layers with per-layer masks for block in self.blocks: - # Select mask based on layer's attention type if block.attn_type == AttentionType.LOCAL_SLIDING: block_mask = merged_mask_sliding else: @@ -2194,25 +1895,12 @@ def __call__( decode=decode, ) - # Final normalization hidden_states = self.norm(hidden_states) return hidden_states -# ============================================================================= -# Full Model -# ============================================================================= - - class T5Gemma2(nnx.Module): - """T5Gemma2 Encoder-Decoder Model. - - Combines the encoder and decoder into a full sequence-to-sequence model. - Uses nnx.Cache for autoregressive decoding. Call init_cache() before - using decode=True. - """ - def __init__( self, config: T5Gemma2Config, @@ -2224,7 +1912,6 @@ def __init__( ): self.config = config - # Use tied embedding for encoder and decoder self.embedder = T5Gemma2ScaledWordEmbedding( vocab_size=config.encoder.text_config.vocab_size, embed_dim=config.encoder.text_config.embed_dim, @@ -2233,16 +1920,15 @@ def __init__( rngs=rngs, ) - # Encoder self.encoder = T5Gemma2Encoder( config.encoder, + mm_shd_cfg=config.shd_cfg if config.shd_cfg.mmp_weight is not None else None, embedder=self.embedder, dtype=dtype, dtype_mm=dtype_mm, rngs=rngs, ) - # Decoder self.decoder = T5Gemma2Decoder( config.decoder, embedder=self.embedder, @@ -2259,14 +1945,6 @@ def init_cache( encoder_seq_length: int, dtype: jnp.dtype | None = None, ) -> None: - """Initialize decoder KV caches for autoregressive decoding. - - Args: - batch_size: Number of sequences in the batch. - max_decode_length: Maximum decoder sequence length. - encoder_seq_length: Encoder sequence length. - dtype: Data type for cache arrays. - """ self.decoder.init_cache( batch_size=batch_size, max_decode_length=max_decode_length, @@ -2287,35 +1965,16 @@ def __call__( *, decode: bool | None = None, deterministic: bool = True, - ) -> tuple[Float[Array, "B L_dec D"], Float[Array, "B L_enc D"]]: # noqa: F722 - """Forward pass of the full model. - - Args: - input_ids: Encoder input token IDs [B, L_enc]. - decoder_input_ids: Decoder input token IDs [B, L_dec]. - attention_mask: Encoder attention mask [B, L_enc]. - decoder_attention_mask: Decoder attention mask [B, L_dec]. - position_ids: Encoder position indices [B, L_enc]. - decoder_position_ids: Decoder position indices [B, L_dec]. - pixel_values: Optional input images [B, N, H, W, C] where N is images per batch. - encoder_outputs: Pre-computed encoder outputs (for generation). - decode: Whether to use KV cache. - deterministic: Whether to run in deterministic mode. - - Returns: - Tuple of (decoder hidden states, encoder hidden states). - """ - # Encode if not provided + ) -> tuple[Float[Array, "B L_dec V"], Float[Array, "B L_enc D"]]: # noqa: F722 if encoder_outputs is None: encoder_outputs = self.encoder( input_ids, attention_mask=attention_mask, position_ids=position_ids, - images=pixel_values, # Note: encoder uses 'images' parameter + images=pixel_values, deterministic=deterministic, ) - # Decode decoder_outputs = self.decoder( decoder_input_ids, encoder_hidden_states=encoder_outputs, @@ -2325,4 +1984,16 @@ def __call__( decode=decode, ) - return decoder_outputs, encoder_outputs + embed_table = self.embedder.embedding[...] + logits = jnp.einsum("btd,vd->btv", decoder_outputs, embed_table) + return logits, encoder_outputs + + +@jax.jit +def forward( + model: T5Gemma2, + encoder_input_ids: Array, + decoder_input_ids: Array, + pixel_values: Array | None = None, +) -> tuple[Array, Array]: + return model(encoder_input_ids, decoder_input_ids, pixel_values=pixel_values) diff --git a/bonsai/models/t5gemma2/tests/run_model.py b/bonsai/models/t5gemma2/tests/run_model.py index 076a35d4..d000e938 100644 --- a/bonsai/models/t5gemma2/tests/run_model.py +++ b/bonsai/models/t5gemma2/tests/run_model.py @@ -276,35 +276,29 @@ def greedy_generate( if use_cache: if step == 0: # First step: prefill with BOS - decoder_outputs = model.decoder( + logits, _ = model( + encoder_input_ids, decoder_input_ids, - encoder_hidden_states=encoder_outputs, - encoder_attention_mask=encoder_mask, + encoder_outputs=encoder_outputs, decode=True, ) else: # Subsequent steps: feed only the new token - decoder_outputs = model.decoder( + logits, _ = model( + encoder_input_ids, next_token[:, None], - encoder_hidden_states=encoder_outputs, - encoder_attention_mask=encoder_mask, + encoder_outputs=encoder_outputs, decode=True, ) else: # No cache: recompute full sequence each step - decoder_outputs = model.decoder( + logits, _ = model( + encoder_input_ids, decoder_input_ids, - encoder_hidden_states=encoder_outputs, - encoder_attention_mask=encoder_mask, + encoder_outputs=encoder_outputs, decode=False, ) - # Get logits from the last position - last_hidden = decoder_outputs[:, -1:, :] - # Use embedding weights for output projection (tied embeddings) - embed_table = model.decoder.embedder.embedding[...] - logits = jnp.einsum("btd,vd->btv", last_hidden, embed_table) - # Greedy selection next_token = jnp.argmax(logits[:, -1, :], axis=-1) generated_ids.append(next_token) diff --git a/bonsai/models/t5gemma2/tests/test_outputs_t5gemma2.py b/bonsai/models/t5gemma2/tests/test_outputs_t5gemma2.py new file mode 100644 index 00000000..717ae4fc --- /dev/null +++ b/bonsai/models/t5gemma2/tests/test_outputs_t5gemma2.py @@ -0,0 +1,145 @@ +# Copyright 2025 The JAX Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import jax +import jax.numpy as jnp +import numpy as np +import torch +from absl.testing import absltest +from flax import nnx +from huggingface_hub import snapshot_download +from transformers import AutoTokenizer, T5Gemma2Model + +from bonsai.models.t5gemma2 import modeling, params + + +class TestModuleForwardPasses(absltest.TestCase): + def setUp(self): + super().setUp() + jax.config.update("jax_default_matmul_precision", "float32") + model_name: str = "google/t5gemma-2-270m-270m" + self.tokenizer = AutoTokenizer.from_pretrained(model_name) + model_ckpt_path = snapshot_download(model_name) + + # HuggingFace reference model + self.torch_model = T5Gemma2Model.from_pretrained(model_name, torch_dtype=torch.float32).eval() + + # Bonsai model (no vision to match text-only comparison) + self.bonsai_config = modeling.T5Gemma2Config.t5gemma2_270m_270m(with_vision=False) + graph_def, state = nnx.split(params.create_model_from_safe_tensors(model_ckpt_path, self.bonsai_config)) + state = jax.tree.map(lambda x: x.astype(jnp.float32) if isinstance(x, jax.Array) else x, state) + self.nnx_model = nnx.merge(graph_def, state) + + self.batch_size = 2 + self.num_input_tokens = 8 + self.relaxed_tol = 1e-3 + + self.enc_cfg = self.bonsai_config.encoder.text_config + + def _to_torch(self, jax_arr): + return torch.tensor(np.array(jax_arr, dtype=np.float32)) + + def _assert_close(self, jax_out, torch_out): + torch.testing.assert_close( + self._to_torch(jax_out), + torch_out.float(), + rtol=self.relaxed_tol, + atol=self.relaxed_tol, + check_dtype=False, + ) + + def test_embedder(self): + nm = self.nnx_model.embedder + tm = self.torch_model.encoder.embed_tokens + + tx = torch.randint(0, self.enc_cfg.vocab_size, size=(self.batch_size, self.num_input_tokens)) + jx = jnp.array(tx.cpu().numpy()) + + self._assert_close(nm(jx), tm(tx)) + + def test_rms_norm(self): + nm = self.nnx_model.encoder.blocks[0].pre_attention_norm + tm = self.torch_model.encoder.layers[0].pre_self_attn_layernorm + + shape = (self.batch_size, self.num_input_tokens, self.enc_cfg.embed_dim) + jx = jax.random.normal(jax.random.key(0), shape=shape) + tx = self._to_torch(jx) + + self._assert_close(nm(jx), tm(tx)) + + def test_feed_forward(self): + nm = self.nnx_model.encoder.blocks[0].mlp + tm = self.torch_model.encoder.layers[0].mlp + + shape = (self.batch_size, self.num_input_tokens, self.enc_cfg.embed_dim) + jx = jax.random.normal(jax.random.key(0), shape=shape) + tx = self._to_torch(jx) + + self._assert_close(nm(jx), tm(tx)) + + def test_full_encoder(self): + text = "Translate English to French: Hello, how are you?" + tokens = self.tokenizer.encode(text, add_special_tokens=False) + tokens = [modeling.BOS_TOKEN] + tokens + + jx = jnp.array([tokens], dtype=jnp.int32) + tx = torch.tensor([tokens], dtype=torch.long) + + jy = self.nnx_model.encoder(jx) + ty = self.torch_model.encoder(tx).last_hidden_state + self._assert_close(jy, ty) + + def test_full_encoder_batched(self): + texts = [ + "Translate English to French: Hello, how are you?", + "Translate English to German: The weather is nice today.", + ] + tokenized = [[modeling.BOS_TOKEN] + self.tokenizer.encode(t, add_special_tokens=False) for t in texts] + # Pad to same length + max_len = max(len(t) for t in tokenized) + padded = [t + [self.enc_cfg.pad_token_id] * (max_len - len(t)) for t in tokenized] + + jx = jnp.array(padded, dtype=jnp.int32) + tx = torch.tensor(padded, dtype=torch.long) + attention_mask = torch.tensor([[1] * len(t) + [0] * (max_len - len(t)) for t in tokenized]) + + jy = self.nnx_model.encoder(jx) + ty = self.torch_model.encoder(tx, attention_mask=attention_mask).last_hidden_state + self._assert_close(jy, ty) + + def test_full_model(self): + text = "Translate English to French: Hello, how are you?" + tokens = self.tokenizer.encode(text, add_special_tokens=False) + tokens = [modeling.BOS_TOKEN] + tokens + + encoder_ids = jnp.array([tokens], dtype=jnp.int32) + decoder_ids = jnp.array([[modeling.BOS_TOKEN]], dtype=jnp.int32) + + torch_encoder_ids = torch.tensor([tokens], dtype=torch.long) + torch_decoder_ids = torch.tensor([[modeling.BOS_TOKEN]], dtype=torch.long) + + # Bonsai model returns (logits, encoder_outputs) + jax_logits, _ = self.nnx_model(encoder_ids, decoder_ids, decode=False) + + # HF T5Gemma2Model returns decoder hidden states; compute logits via tied embeddings + torch_out = self.torch_model(input_ids=torch_encoder_ids, decoder_input_ids=torch_decoder_ids) + torch_hidden = torch_out.last_hidden_state + embed_table = self.torch_model.encoder.embed_tokens.weight + torch_logits = torch.einsum("btd,vd->btv", torch_hidden, embed_table) + + self._assert_close(jax_logits, torch_logits) + + +if __name__ == "__main__": + absltest.main()