diff --git a/bonsai/models/t5gemma2/README.md b/bonsai/models/t5gemma2/README.md new file mode 100644 index 00000000..872390e1 --- /dev/null +++ b/bonsai/models/t5gemma2/README.md @@ -0,0 +1,11 @@ +# 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/__init__.py b/bonsai/models/t5gemma2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bonsai/models/t5gemma2/modeling.py b/bonsai/models/t5gemma2/modeling.py new file mode 100644 index 00000000..23ae620b --- /dev/null +++ b/bonsai/models/t5gemma2/modeling.py @@ -0,0 +1,1999 @@ +# 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 +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 + + +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, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.LOCAL_SLIDING, + AttentionType.GLOBAL, +) + + +def _make_layer_types(num_layers: int) -> tuple[AttentionType, ...]: + 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" + 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 + shd_cfg: VisionShardingCfg = field(default_factory=VisionShardingCfg) + + +@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 + shd_cfg: TextShardingCfg = field(default_factory=TextShardingCfg) + + @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): + pass + + +@dataclass(frozen=True) +class T5Gemma2Config: + encoder: T5Gemma2EncoderConfig + decoder: T5Gemma2DecoderConfig + eoi_token_index: int = 256000 + pad_token_id: int = 0 + shd_cfg: MMShardingCfg = field(default_factory=MMShardingCfg) + + @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, + 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, + 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, + shd_cfg=text_shd, + ) + 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( + 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, + shd_cfg=text_shd, + ), + shd_cfg=mm_shd, + ) + + @classmethod + 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, 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, 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) + + +_K_MASK = jnp.finfo(jnp.bfloat16).min + +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 + +NORMAL_INIT = nnx.initializers.normal() +ZEROS_INIT = nnx.initializers.zeros_init() + + +def _get_rope_base_frequency( + text_config: T5Gemma2TextConfig, + attn_type: AttentionType, +) -> int: + 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 + + +def _get_rope_scale_factor( + text_config: T5Gemma2TextConfig, + attn_type: AttentionType, +) -> float: + 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: + if rope_params.rope_type == "linear": + return rope_params.factor + return 1.0 + + +def apply_rope( + 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"]: # noqa: F722 + 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) + + +class T5Gemma2Einsum(nnx.Module): + 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)) + + @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): + 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)) + + @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) + normed_inputs = x * jax.lax.rsqrt(var + self.epsilon) + scale = jnp.expand_dims(self.scale[...], axis=range(len(x.shape) - 1)) + return normed_inputs * (1 + scale) + + +class T5Gemma2FeedForward(nnx.Module): + 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, + ): + 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) + else: + gating_shape = (2, features, hidden_dim) + + self.gating_einsum = T5Gemma2Einsum( + shape=gating_shape, + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + self.linear = T5Gemma2Einsum( + shape=(hidden_dim, features), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + @jax.named_scope("feed_forward") + def __call__(self, x: Float[Array, "B L D"]) -> Float[Array, "B L D"]: # noqa: F722 + 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, :] + 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 + mask = input_mask[:, None, None, :] + return mask + + +def make_causal_mask( + input_mask: Bool[Array, "B L"], # noqa: F722 +) -> Bool[Array, "B 1 L L"]: # noqa: F722 + seq_len = input_mask.shape[-1] + causal = jnp.tril(jnp.ones((seq_len, seq_len), dtype=jnp.bool_)) + mask = input_mask[:, None, None, :] & causal[None, None, :, :] + return mask + + +def make_sliding_window_mask( + positions: Int[Array, "B L"], # noqa: F722 + sliding_window: int, + *, + bidirectional: bool = False, +) -> Bool[Array, "B 1 L L"]: # noqa: F722 + q_pos = positions[:, :, None] + k_pos = positions[:, None, :] + + if bidirectional: + 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: + 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"], # noqa: F722 + positions: Int[Array, "B L"], # noqa: F722 + sliding_window: int, +) -> Bool[Array, "B 1 L L"]: # noqa: F722 + causal_mask = make_causal_mask(input_mask) + sliding_mask = make_sliding_window_mask(positions, sliding_window) + return causal_mask & sliding_mask + + +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 + batch_size, _, seq_len, _ = decoder_mask.shape + enc_len = encoder_mask.shape[-1] + cross_mask = jnp.broadcast_to(encoder_mask, (batch_size, 1, seq_len, enc_len)) + return jnp.concatenate([decoder_mask, cross_mask], axis=-1) + + +def make_decode_mode_self_mask( + batch_size: int, + query_len: int, + max_cache_len: int, + current_cache_len: int, +) -> Bool[Array, "B 1 Q K"]: # noqa: F722 + 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, + max_cache_len: int, + sliding_window: int, +) -> Bool[Array, "B 1 1 K"]: # noqa: F722 + k_positions = jnp.arange(max_cache_len) + in_window = jnp.abs(query_pos - k_positions) < sliding_window + 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)) + + +class T5Gemma2ScaledWordEmbedding(nnx.Module): + 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 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"]: # noqa: F722 + embeddings = self.embedding[input_ids] * self.embed_scale + eoi_mask = input_ids == END_OF_IMAGE_TOKEN + return jnp.where(eoi_mask[..., None], self.eoi_embedding[...], embeddings) + + +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"]: # noqa: F722 + 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): + 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"], # noqa: F722 + *, + deterministic: bool | None = None, + ) -> Float[Array, "B L D"]: # noqa: F722 + x = self.dense1(x) + x = jax.nn.gelu(x) + x = self.dropout(x, deterministic=deterministic) + return self.dense2(x) + + +class T5Gemma2VisionEncoderBlock(nnx.Module): + 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 = T5Gemma2VisionMLP( + width=width, + mlp_dim=mlp_dim, + dropout=dropout, + dtype=dtype, + rngs=rngs, + ) + + def __call__( + self, + x: Float[Array, "B L D"], # noqa: F722 + *, + deterministic: bool | None = None, + ) -> Float[Array, "B L D"]: # noqa: F722 + 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): + 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): + """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, strides=window_shape, padding="VALID") + 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).""" + + def __init__( + self, + config: T5Gemma2VisionConfig, + *, + 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 = T5Gemma2VisionEncoder( + 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 = T5Gemma2VisionExit(output_length=256, rngs=rngs) + + def __call__( + self, + images: Float[Array, "B N H W C"], # noqa: F722 + *, + deterministic: bool | None = None, + ) -> Float[Array, "B N P D"]: # noqa: F722 + 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[...] + 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: + 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, + 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(), + 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"]: # noqa: F722 + x = self.mm_soft_embedding_norm(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): + """Vision embedder - tokenizes images and embeds into text space.""" + + def __init__( + self, + *, + vision_config: T5Gemma2VisionConfig, + embed_dim: int, + mm_shd_cfg: MMShardingCfg | None = None, + 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, + mm_shd_cfg=mm_shd_cfg, + dtype=dtype, + rngs=rngs, + ) + + def __call__( + self, + images: Float[Array, "B N H W C"], # noqa: F722 + *, + deterministic: bool | None = None, + ) -> Float[Array, "B N P De"]: # noqa: F722 + 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) + + +def merge_mm_embeddings( + text_embeddings: jnp.ndarray, + vision_embeddings: jnp.ndarray, + mask: jnp.ndarray, +) -> jnp.ndarray: + 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: + 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) + + +class T5Gemma2EncoderAttention(nnx.Module): + def __init__( + self, + num_q_heads: int, + num_kv_heads: int, + features: int, + head_dim: int, + 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, + 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.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 + self.shd_cfg = shd_cfg + + self.attn_vec_einsum = T5Gemma2Einsum( + shape=(num_q_heads, head_dim, features), + kernel_init=kernel_init, + dtype=dtype, + rngs=rngs, + ) + + 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, + 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, + ) + + @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 + + @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 + shd = self.shd_cfg.activation if self.shd_cfg is not None else None + + 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, out_sharding=shd) + key_proj, value_proj = self.kv_einsum("BSD,CKDH->CBSKH", x) + + if self.use_qk_norm: + query_proj = self._query_norm(query_proj) + key_proj = self._key_norm(key_proj) + + 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, + ) + + 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) + + 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 + + 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", + ) + sliding_mask = make_sliding_window_mask( + segment_pos, + self.sliding_window_size, + bidirectional=True, + )[:, 0, :, :] + attn_mask = attn_mask * sliding_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) + + if self.use_gqa: + b, t, kg, s = probs.shape + probs = probs.reshape( + (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 + 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, out_sharding=shd) + + +class T5Gemma2EncoderBlock(nnx.Module): + 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, + shd_cfg: TextShardingCfg | None = None, + 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, + 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, + epsilon=rms_norm_eps, + 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, + shd_cfg=shd_cfg, + 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, + scale_init=scale_init, + dtype=dtype, + rngs=rngs, + ) + + if use_post_attn_norm: + self.post_attention_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + self.pre_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + self.mlp = T5Gemma2FeedForward( + 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, + ) + + if use_post_ffw_norm: + self.post_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_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 + 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 + + 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): + 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, + rngs: nnx.Rngs, + ): + 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 + + 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, + ) + + 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, + ) + + attention_types = text_config.layer_types + if not attention_types: + attention_types = tuple(AttentionType.GLOBAL for _ in range(text_config.num_hidden_layers)) + + 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, + shd_cfg=shd_cfg, + rope_base_frequency=_get_rope_base_frequency( + text_config, + attention_types[i], + ), + rope_scale_factor=_get_rope_scale_factor( + 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 + ), + use_qk_norm=True, + dtype=dtype, + rngs=rngs, + ) + for i in range(text_config.num_hidden_layers) + ] + ) + + 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"], # 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"]: # noqa: F722 + batch_size, seq_len = tokens.shape + + if position_ids is None: + position_ids = jnp.arange(seq_len)[None, :].repeat(batch_size, axis=0) + + if attention_mask is None: + attention_mask = jnp.ones((batch_size, seq_len), dtype=jnp.bool_) + + x = self.embedder(tokens) + + 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, + ) + + attn_mask = make_bidirectional_mask(attention_mask) + if attn_mask.ndim == 4: + attn_mask = attn_mask[:, 0, :, :] + + if attn_mask.shape[-2] == 1: + attn_mask = jnp.broadcast_to(attn_mask, (batch_size, seq_len, seq_len)) + + for block in self.blocks: + x = block(x, position_ids, attn_mask) + + return self.norm(x) + + +class T5Gemma2MergedAttention(nnx.Module): + """Merged self-attention and cross-attention for decoder. + + 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, + *, + shd_cfg: TextShardingCfg | None = None, + 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, + 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.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 + self.shd_cfg = shd_cfg + + self.num_kv_groups = num_q_heads // num_kv_heads + + 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, + ) + + 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, + ) + + self.cached_self_key: nnx.Cache[Array] | None = nnx.data(None) + self.cached_self_value: nnx.Cache[Array] | None = nnx.data(None) + self.cached_cross_key: nnx.Cache[Array] | None = nnx.data(None) + self.cached_cross_value: nnx.Cache[Array] | None = nnx.data(None) + 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: + cache_shd = self.shd_cfg.cache if self.shd_cfg is not None else None + + 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, out_sharding=cache_shd)) + self.cached_self_value = nnx.Cache(jnp.zeros(self_cache_shape, dtype, out_sharding=cache_shd)) + + 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, out_sharding=cache_shd)) + self.cached_cross_value = nnx.Cache(jnp.zeros(cross_cache_shape, dtype, out_sharding=cache_shd)) + + 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 + 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"]: # noqa: F722 + seq_len = hidden_states.shape[1] + shd = self.shd_cfg.activation if self.shd_cfg is not None else None + + 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) + + if self.use_qk_norm: + query = self.q_norm(query) + self_key = self.k_norm(self_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, + ) + + decode = first_from( + decode, + self.decode, + error_msg="No decode argument provided to T5Gemma2MergedAttention", + ) + + 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) + + 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 + + self_key = self_key_updated + self_value = self_value_updated + + def compute_cross_kv(): + 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 + + cross_key, cross_value = jax.lax.cond( + cur_index == 0, + compute_cross_kv, + lambda: (self.cached_cross_key[...], self.cached_cross_value[...]), + ) + + self.cached_cross_key[...] = cross_key + self.cached_cross_value[...] = cross_value + + self.cache_index[...] = cur_index + seq_len + else: + 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) + + query = query * self.query_pre_attn_scalar + + key = jnp.concatenate([self_key, cross_key], axis=1) + value = jnp.concatenate([self_value, cross_value], axis=1) + + 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] + + 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) + + attn_weights = jnp.einsum("BNTH,BNSH->BNTS", query, key) + + 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 + + if merged_attention_mask is not None: + attn_weights = jnp.where(merged_attention_mask, attn_weights, _K_MASK) + + attn_weights = jax.nn.softmax(attn_weights, axis=-1).astype(value.dtype) + + attn_output = jnp.einsum("BNTS,BNSH->BNTH", attn_weights, value) + + attn_output = jnp.transpose(attn_output, (0, 2, 1, 3)) + + output = self.o_proj("BTNH,NHD->BTD", attn_output, out_sharding=shd) + + return output + + +class T5Gemma2DecoderBlock(nnx.Module): + 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, + shd_cfg: TextShardingCfg | None = None, + 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, + 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 + + self.pre_attention_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + 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, + 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, + ) + + if use_post_attn_norm: + self.post_attention_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + self.pre_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + 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, + ) + + if use_post_ffw_norm: + self.post_ffw_norm = T5Gemma2RMSNorm( + embed_dim, + epsilon=rms_norm_eps, + dtype=dtype, + scale_init=scale_init, + rngs=rngs, + ) + + @jax.named_scope("decoder_block") + def __call__( + self, + 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"]: # noqa: F722 + 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 + + 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): + 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 + shd_cfg = config.shd_cfg if config.shd_cfg.activation is not None else None + + 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, + ) + + attention_types = config.layer_types + if not attention_types: + attention_types = tuple(AttentionType.GLOBAL for _ in range(config.num_hidden_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, + 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, + 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) + ] + ) + + self.norm = T5Gemma2RMSNorm( + config.embed_dim, + epsilon=config.rms_norm_eps, + 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: + 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: + 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"], # 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"]: # noqa: F722 + batch_size, dec_seq_len = input_ids.shape + enc_seq_len = encoder_hidden_states.shape[1] + + if position_ids is None: + cache_index = self.get_cache_index() + if decode and cache_index > 0: + 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, + ) + + 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_, + ) + + encoder_mask = encoder_attention_mask[:, None, None, :] + + if decode: + cache_index = self.get_cache_index() + max_decode_len = self.blocks[0].attn.cached_self_key.shape[1] + current_len = cache_index + dec_seq_len + + full_decoder_mask = make_decode_mode_self_mask( + batch_size, + dec_seq_len, + max_decode_len, + current_len, + ) + sliding_decoder_mask = make_decode_mode_sliding_mask( + batch_size, + cache_index, + max_decode_len, + self.config.sliding_window, + ) + + 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: + full_decoder_mask = make_causal_mask(attention_mask) + sliding_decoder_mask = make_sliding_window_causal_mask( + attention_mask, + position_ids, + self.config.sliding_window, + ) + + merged_mask_full = make_merged_attention_mask(full_decoder_mask, encoder_mask) + merged_mask_sliding = make_merged_attention_mask(sliding_decoder_mask, encoder_mask) + + hidden_states = self.embedder(input_ids) + + for block in self.blocks: + 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, + ) + + hidden_states = self.norm(hidden_states) + + return hidden_states + + +class T5Gemma2(nnx.Module): + 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 + + 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, + ) + + 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, + ) + + 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: + 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"], # 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 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, + deterministic=deterministic, + ) + + 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, + ) + + 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/params.py b/bonsai/models/t5gemma2/params.py new file mode 100644 index 00000000..9be32ceb --- /dev/null +++ b/bonsai/models/t5gemma2/params.py @@ -0,0 +1,510 @@ +# 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/__init__.py b/bonsai/models/t5gemma2/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/bonsai/models/t5gemma2/tests/run_model.py b/bonsai/models/t5gemma2/tests/run_model.py new file mode 100644 index 00000000..d000e938 --- /dev/null +++ b/bonsai/models/t5gemma2/tests/run_model.py @@ -0,0 +1,510 @@ +# 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 = [] + next_token = None # Track the last generated token for cached decoding + + for step in range(max_new_tokens): + if use_cache: + if step == 0: + # First step: prefill with BOS + logits, _ = model( + encoder_input_ids, + decoder_input_ids, + encoder_outputs=encoder_outputs, + decode=True, + ) + else: + # Subsequent steps: feed only the new token + logits, _ = model( + encoder_input_ids, + next_token[:, None], + encoder_outputs=encoder_outputs, + decode=True, + ) + else: + # No cache: recompute full sequence each step + logits, _ = model( + encoder_input_ids, + decoder_input_ids, + encoder_outputs=encoder_outputs, + decode=False, + ) + + # 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() 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()