diff --git a/docs/gemma4_vision_port.md b/docs/gemma4_vision_port.md new file mode 100644 index 000000000..0dcfce8e0 --- /dev/null +++ b/docs/gemma4_vision_port.md @@ -0,0 +1,203 @@ +# Porting the real Gemma 4 vision tower (`gemma4_vision`) to Tunix + +## Why this exists + +The first attempt at "multi-modal (vision) support for Gemma 4" reused Gemma 3's +**SigLIP** encoder (`tunix/models/gemma3/vision.py`, `SigLIPConfig`). That is the +wrong architecture. The released `google/gemma-4-e2b-it` checkpoint declares its +own vision tower: + +``` +architectures : ['Gemma4ForConditionalGeneration'] +vision_config -> model_type: 'gemma4_vision', 16 layers, hidden 768, + intermediate 3072, patch 16 +audio_config -> model_type: 'gemma4_audio', 12 layers, hidden 1024 +``` + +`gemma4_vision` is **not** SigLIP: it uses gated MLPs (`gate/up/down`), 2D RoPE, +per-head q/k/v RMSNorm, a 4-norm sandwich block, a factored 2D learned position +table, and a spatial pooler. The checkpoint also has a `gemma4_audio` tower and +single-tensor `embed_vision` / `embed_audio` projectors (there is **no** +`multi_modal_projector` that the SigLIP loader expected). + +This document is the reverse-engineered spec (from HF +`transformers.models.gemma4.modeling_gemma4`, v5.x) and the staged plan to do it +correctly. + +## Scope + +* **In scope:** the vision tower + the `embed_vision` projector + the safetensors + loader mappings + wiring into `Gemma4.__call__` + parity validation. +* **Out of scope (for now):** the `gemma4_audio` Conformer tower. The loader + should gracefully **skip** all `model.audio_tower.*` and `model.embed_audio.*` + keys (the existing loader already logs skipped keys; that is acceptable for a + vision-only PR). + +## Architecture (as ported in `tunix/models/gemma4/vision_real.py`) + +`Gemma4VisionConfig` (e2b defaults): hidden 768, intermediate 3072, 16 layers, +12 heads, 12 kv-heads, head_dim 64, eps 1e-6, patch 16, position_embedding_size +10240, pooling_kernel_size 3, rope_theta 100.0, act `gelu_pytorch_tanh`, +`use_clipped_linears=False`, `standardize=False`. + +Forward pipeline (`Gemma4VisionModel`): + +1. **PatchEmbedder** — `pixel_values` are pre-flattened patches `[B, P, 3*16²]`. + Compute `2*(px-0.5)`, project with a bias-free `Linear(768→768)`, then add a + factored 2D position embedding: `one_hot(pos).permute @ table[2,10240,768]` + summed over the two spatial axes. Padding patches (pos == -1) get zeroed. +2. **Encoder** — 2D RoPE (`rotary_emb`) computes cos/sin from `pixel_position_ids` + (`inv_freq` over `spatial_dim = head_dim//2 = 32`, reused per axis). 16 + bidirectional sandwich layers: + `x = x + post_attn_norm(attn(input_norm(x)))`, then + `x = x + post_ffw_norm(mlp(pre_ffw_norm(x)))`. + * **Attention** — bias-free q/k/v/o projections; `q_norm`,`k_norm` are scaled + RMSNorm, `v_norm` is **unscaled** RMSNorm; multidim RoPE on q,k; `scaling = + 1.0` (not `head_dim**-0.5`); non-causal. + * **MLP** — `down(gelu_tanh(gate(x)) * up(x))`. +3. **Pooler** — zero padding, 2D average-pool patches into + `output_length = P / pooling_kernel²` soft tokens, scale by `sqrt(768)`. + +Projector (`Gemma4MultimodalEmbedder`, a.k.a. `embed_vision`): +`embedding_projection(embedding_pre_projection_norm(soft_tokens))`, where the +pre-projection norm is **unscaled** RMSNorm(768) and the projection is bias-free +`Linear(768 → text_hidden=1536)`. + +### Gemma-4-specific gotchas (verified against HF source) + +* `Gemma4RMSNorm` scales by `weight` **directly** (not `1 + weight` like Gemma + 1/2/3). Loader must not apply a +1. +* `with_scale=False` norms (`v_norm`, `embedding_pre_projection_norm`) have **no** + weight in the checkpoint — don't expect one. +* `Gemma4ClippableLinear` nests an inner `linear`, so checkpoint keys are + `...{proj}.linear.weight` (note the extra `.linear.`). For e2b + (`use_clipped_linears=False`) the `input_min/max`,`output_min/max` buffers are + ±inf no-ops and can be ignored. +* Attention `scaling = 1.0`. + +## Checkpoint → nnx key mapping (for the loader) + +All real keys are prefixed `model.` (the existing vision regexes in +`params_safetensors.py` lack this prefix and use `re.match`, so they would skip +every vision key — that must be fixed, e.g. with an optional `(?:model\.)?`). + +| Real checkpoint key | nnx param path | transform | +|---|---|---| +| `model.vision_tower.patch_embedder.input_proj.weight` | `vision_tower.patch_embedder.input_proj.kernel` | transpose (1,0) | +| `model.vision_tower.patch_embedder.position_embedding_table` | `vision_tower.patch_embedder.position_embedding_table` | none | +| `model.vision_tower.encoder.layers.N.self_attn.{q,k,v,o}_proj.linear.weight` | `...self_attn.{q,k,v,o}_proj.linear.kernel` | transpose (1,0) | +| `model.vision_tower.encoder.layers.N.self_attn.{q,k}_norm.weight` | `...self_attn.{q,k}_norm.scale` | none | +| `model.vision_tower.encoder.layers.N.mlp.{gate,up,down}_proj.linear.weight` | `...mlp.{gate,up,down}_proj.linear.kernel` | transpose (1,0) | +| `model.vision_tower.encoder.layers.N.{input_layernorm,post_attention_layernorm,pre_feedforward_layernorm,post_feedforward_layernorm}.weight` | `...{same}.scale` | none | +| `model.embed_vision.embedding_projection.weight` | `embed_vision.embedding_projection.kernel` | transpose (1,0) | +| `model.vision_tower.encoder.layers.N.self_attn.{q,k,v,o}_proj.{input,output}_{min,max}` | — | **skip** (±inf, e2b) | +| `model.audio_tower.*`, `model.embed_audio.*` | — | **skip** (out of scope) | + +`tests/models/gemma4/vision_real_test.py::test_module_tree_matches_checkpoint_keys` +pins the nnx side of this table. + +## Staged plan & status + +* **Stage 1 — vision tower port. DONE (wiring/shapes only).** + `tunix/models/gemma4/vision_real.py` builds and forward-passes; module tree + maps 1:1 to checkpoint keys. Tests: + `tests/models/gemma4/vision_real_test.py` (4 passing). **Numerics unvalidated.** +* **Stage 2 — image processor + loader. DONE (coverage-verified; resize + numerics pending).** + * `tunix/models/gemma4/image_processing.py` — NumPy port of the HF processor + (target-size, patchify, `(x,y)` position ids, padding = -1). Resize uses PIL + (not bit-exact with torchvision — see caveat). + * `tunix/models/gemma4/vision_params_safetensors.py` — `vision_key_mapping` + (`model.`-prefix-safe, `$`-anchored) + `create_vision_stack_from_safe_tensors` + loading a `Gemma4VisionStack`. Audio + clip buffers are skipped. + * Tests (`vision_params_safetensors_test.py`, `image_processing_test.py`): + every real vision key maps to a param, linears transpose / norms don't, + audio+clip buffers skip, no double-matches, **no uninitialised params**, and + processor output feeds the stack end-to-end. **Numeric parity still pending.** +* **Stage 3 — numeric parity (GATING). PORT MATH VALIDATED; real-checkpoint run pending.** + Two harnesses: + * `examples/gemma4/vision_parity_random_weights.py` — **checkpoint-free**. + Builds a small HF `Gemma4VisionModel`+`Gemma4MultimodalEmbedder` with random + weights, serializes them under real checkpoint key names, loads them into the + JAX `Gemma4VisionStack` via the production loader, and compares per-layer in + fp32. **Run in-sandbox (transformers 5.9.0 + torch 2.12): PARITY PASSED** — + `after_patch_embed` exact, per-layer max-abs ~1e-6, tower 1.5e-5, projector + 3.6e-7. This validates the port math + key mapping + loader against HF. + * `examples/gemma4/vision_parity_check.py` — same comparison against a **real** + `google/gemma-4-*-it` checkpoint (adds real-weight file loading + an actual + forward). Run: + ``` + pip install torch 'transformers==5.9.0' safetensors # 5.10.1 trips a torch._dynamo bug + python examples/gemma4/vision_parity_check.py --ckpt ~/gemma4-e2b + ``` + Status: the architecture/math is proven correct against HF. The only thing the + real-checkpoint run adds is confirming the real `.safetensors` files load and a + real image captions — do that before Stage 4. +* **Stage 4 — end-to-end. WIRED + merge-validated; full-model parity + caption pending.** + * `tunix/models/gemma4/model.py` — `Gemma4.__call__` gains a non-breaking + `input_embeddings` override (legacy SigLIP path untouched). + * `tunix/models/gemma4/multimodal.py` — `Gemma4Multimodal` composes the text + `Gemma4` with `Gemma4VisionStack`: embed tokens, run the vision stack, scatter + soft tokens at `tokens == image_token_id` (HF `masked_scatter` equivalent via + `merge_embeddings`), run the transformer on merged embeddings. Plus + `create_multimodal_from_safe_tensors` (loads text + vision from one checkpoint). + * `examples/gemma4/multimodal_generate.py` — single-image, no-padding, eager + greedy caption demo. + * Tests (`multimodal_test.py`, sandbox dry-run): soft tokens land exactly at + image positions, text positions are untouched, the image changes downstream + logits, the mask is bidirectional over the image span, and the greedy loop + runs end-to-end. **26/26 gemma4 tests pass.** + + Open items / honest limits: + * **Single, non-padded image only** — the merge assumes #valid-soft-tokens == + #placeholders; multi-image / padded batches need a valid-token gather first. + * **Real caption** needs the checkpoint (run `multimodal_generate.py`). + +* **Stage 4 follow-up — full-model parity vs HF `Gemma4Model.forward`.** + `examples/gemma4/multimodal_parity_random_weights.py` builds a tiny random HF + `Gemma4ForConditionalGeneration`, saves it as safetensors, loads into + `Gemma4Multimodal`, runs both, and diffs per-position logits + PLE. + + Findings against HF: + * **PLE token-identity branch: bit-exact (max=0).** HF + `Gemma4Model.forward` does NOT pass the merged image embeddings to PLE; + it substitutes `image_token_id → pad_token_id` in the ids handed to + `embed_tokens_per_layer`, while the context-projection branch sees the + merged embeddings (with vision at image positions). `Gemma4Multimodal. + _compute_per_layer_inputs` mirrors this exactly; locked in by + `tests/models/gemma4/multimodal_test.py:: + test_per_layer_inputs_substitutes_pad_at_image_positions`. + * **Bidirectional vs causal mask.** HF `Gemma4TextConfig. + use_bidirectional_attention == "vision"` controls this; smaller models + default to plain causal. Tunix matches via a `bidirectional_image_span` + flag on `Gemma4Multimodal` (default `False` = causal, like HF small). + * **Residual divergence at text positions after the image** (max ~9e-2 in + the random-weights probe) reproduces in a **pure-text** parity (same HF + weights, no image, no multimodal wrapper) — so this is a **pre-existing + Tunix-vs-HF divergence in `tunix.models.gemma4` text-model arithmetic** + (probably global-layer RoPE / proportional partial-rotary settings), not + something introduced by the vision port. Out of scope for this PR; worth + flagging upstream separately. + + Net: the multimodal wrapper itself now matches HF semantically on every + point we can prove without the real checkpoint. BOS (pure text, no image + influence) is exact (1.6e-7); image-position logit diff is `~4e-2` and + bounded by the pre-existing text-model divergence. + +## Validation prerequisites (Stage 3) + +Neither a clean Tunix venv nor this sandbox ships torch. To produce HF reference +activations you must, on a machine with the checkpoint: + +```bash +pip install 'torch' 'transformers==5.10.1' # or whatever ships Gemma4 +``` + +Then load the same safetensors into both stacks and diff. (A harness skeleton +will live in `examples/gemma4/vision_parity_check.py` once Stage 2 lands.) + +## Caveats + +* This is a from-source port; bit-exact parity is unproven until Stage 3. +* The official multimodal Gemma 4 implementation is likely in progress upstream + in Tunix (the text model is Google-authored). Coordinate before investing in + Stages 2-4 to avoid duplicating in-flight work. diff --git a/examples/gemma4/multimodal_generate.py b/examples/gemma4/multimodal_generate.py new file mode 100644 index 000000000..5de2e89d5 --- /dev/null +++ b/examples/gemma4/multimodal_generate.py @@ -0,0 +1,324 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""End-to-end caption demo for Gemma 4 multimodal (JIT + KV cache). + +Loads a real `google/gemma-4-*-it` checkpoint into `Gemma4Multimodal`, +processes one image, and greedy-decodes a caption using a prefill/decode split +with a fixed-size KV cache. JIT-compiled: first call compiles (~1-2 min for +2B), subsequent tokens decode in milliseconds. + +Usage: + source ~/tunix-venv/bin/activate + python examples/gemma4/multimodal_generate.py \\ + --ckpt ~/gemma4-e2b --image cat.jpg --prompt "Describe this image." +""" + +from __future__ import annotations + +import argparse +import functools +import sys +import time + +import numpy as np +import jax +import jax.numpy as jnp +from flax import nnx + + +# --------------------------------------------------------------------------- +# JIT-compiled prefill and decode (module-level so compilation is cached) +# --------------------------------------------------------------------------- + +@functools.partial(nnx.jit, static_argnames=('max_new_tokens',)) +def _prefill(model, tokens, pixel_values, pixel_position_ids, max_new_tokens): + """Full prefill: encodes image+text, populates KV cache. + + Returns (next_token_ids (B,), kv_cache). + """ + B, L = tokens.shape + cache_dtype = model.text_model.config.dtype # match computation dtype + cache = model.text_model.init_cache(B, L + max_new_tokens, cache_dtype) + attn_mask = model.get_attention_mask(tokens) # (B, L, L) causal + logits, new_cache = model( + tokens, pixel_values, pixel_position_ids, + cache=cache, attention_mask=attn_mask, + decode_only_last_token=True, + ) + return jnp.argmax(logits[:, 0, :], axis=-1), new_cache # (B,), cache + + +@functools.partial(nnx.jit, static_argnames=('n_steps',)) +def _decode_n_tokens(text_model, initial_token, initial_pos, initial_cache, + n_steps): + """Decode n_steps tokens entirely on-device via lax.scan. + + Calling nnx.jit once per token costs ~55ms Python overhead per step (NNX + module-tree extraction). lax.scan amortises that cost over all n_steps, + yielding ~10x higher throughput on A100. + + Args: + text_model: Gemma4TextModel NNX module (weights are frozen during decode). + initial_token: (B, 1) int32 — first token to feed (from prefill). + initial_pos: (B, 1) int32 — absolute position of initial_token. + initial_cache: KV-cache dict from _prefill. + n_steps: (static) number of tokens to generate. + + Returns: + all_tokens: (n_steps, B) int32 — generated token ids. + final_cache: updated KV-cache (useful if you want to extend later). + """ + def _step(carry, _): + token, pos, cache = carry + cache_len = next(iter(cache.values()))['v'].shape[1] + attn_mask = ( + jnp.arange(cache_len)[None, None, :] <= pos + ).astype(jnp.bool_) # (B, 1, cache_len) + logits, new_cache = text_model( + token, positions=pos, cache=cache, + attention_mask=attn_mask, decode_only_last_token=True, + ) + next_token = jnp.argmax(logits[:, 0, :], axis=-1) # (B,) + return (next_token[:, None], pos + 1, new_cache), next_token + + (_, _, final_cache), all_tokens = jax.lax.scan( + _step, + (initial_token, initial_pos, initial_cache), + xs=None, + length=n_steps, + ) + return all_tokens, final_cache # all_tokens: (n_steps, B) + + +# --------------------------------------------------------------------------- +# Public helpers +# --------------------------------------------------------------------------- + +def build_image_prompt_tokens(tokenizer, prompt_text, image_token_id, num_soft, + bos_id): + """[BOS] *num_soft . Returns (1, L) int32.""" + text_ids = tokenizer(prompt_text, add_special_tokens=False)['input_ids'] + ids = [bos_id] + [image_token_id] * num_soft + list(text_ids) + return np.asarray([ids], dtype=np.int32) + + +def _unpadded_patches(proc, image_chw, do_resize=True): + """Process one image and strip padding so #patches == num_soft * pool^2.""" + px, pos, num_soft = proc([image_chw], do_resize=do_resize) + n = num_soft[0] * proc.pooling_kernel_size ** 2 + return px[:, :n], pos[:, :n], num_soft[0] + + +def greedy_generate(model, tokens, pixel_values, pixel_position_ids, + max_new_tokens, eos_ids, tokenizer): + """JIT-compiled greedy decode with prefill/decode KV-cache split. + + First call JIT-compiles both stages (~1-2 min for 2B on A100). + Subsequent tokens decode in milliseconds. + """ + B, L = tokens.shape + tokens_j = jnp.asarray(tokens) + px_j = jnp.asarray(pixel_values) + pos_j = jnp.asarray(pixel_position_ids) + + # ---- Prefill ------------------------------------------------ + t0 = time.time() + print('Prefill (JIT-compiles on first run — patience)...', end=' ', + flush=True) + next_id_arr, cache = _prefill(model, tokens_j, px_j, pos_j, max_new_tokens) + jax.block_until_ready(next_id_arr) + print(f'done ({time.time() - t0:.1f}s)', flush=True) + + first_token = int(next_id_arr[0]) + if first_token in eos_ids: + return [first_token] + + # ---- Autoregressive decode (lax.scan — no per-step Python overhead) ------ + # All n_decode steps run on-device in a single JIT dispatch, avoiding the + # ~55ms Python overhead that nnx.jit incurs per call on a 2B model. + n_decode = max_new_tokens - 1 # first token already came from prefill + t_decode = time.time() + print(f'Decoding {n_decode} tokens via lax.scan ' + '(JIT-compiles on first run)...', end=' ', flush=True) + + all_tokens, _ = _decode_n_tokens( + model.text_model, + next_id_arr[:, None], # (B, 1) + jnp.full((B, 1), L, jnp.int32), # absolute pos of first decode token + cache, + n_decode, + ) + jax.block_until_ready(all_tokens) + + elapsed = time.time() - t_decode + tokens_list = all_tokens[:, 0].tolist() # (n_decode,) + rate = n_decode / max(elapsed, 1e-6) + print(f'done ({elapsed:.1f}s, {rate:.1f} tok/s)', flush=True) + + # Assemble output, stopping at first EOS + generated = [first_token] + for tok in tokens_list: + generated.append(tok) + if tok in eos_ids: + break + + return generated + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument('--ckpt', required=True, + help='Path to local gemma-4-*-it checkpoint directory') + ap.add_argument('--image', required=True, help='Image file path') + ap.add_argument('--prompt', default='Describe this image.') + ap.add_argument('--max-new-tokens', type=int, default=64) + ap.add_argument('--tp-size', type=int, default=1, + help='Tensor-parallel degree (default 1 = single device). ' + 'TP shards each weight across GPUs, but for a model that ' + 'fits on one GPU at batch=1 it is SLOWER than single-GPU ' + '(per-token all-reduces dominate): measured 46 tok/s at ' + 'TP=4 vs 136 tok/s at TP=1 for Gemma4 2B on A100. Only ' + 'raise this for models too large for one device.') + args = ap.parse_args() + + # ---- Tensor-parallel mesh ------------------------------------------------ + # Weights carry per-array sharding metadata (ShardingConfig); the safetensors + # loader reads it via nnx.get_named_sharding(mesh) and device_puts each tensor + # onto its target shards as it loads (no single-GPU bottleneck). XLA's SPMD + # partitioner then inserts the cross-device collectives automatically. + tp = args.tp_size if args.tp_size > 0 else jax.device_count() + mesh = jax.make_mesh( + (1, tp), ('fsdp', 'tp'), + axis_types=(jax.sharding.AxisType.Auto,) * 2, + ) + print(f'Mesh: {tp}-way tensor parallel over {jax.device_count()} device(s)', + flush=True) + + try: + from PIL import Image + from transformers import AutoTokenizer, Gemma4Config + from tunix.models.gemma4 import image_processing as ip + from tunix.models.gemma4 import model as model_lib + from tunix.models.gemma4 import multimodal as mm_lib + from tunix.models.gemma4 import vision_real + except ImportError as e: + print(f'ERROR: {e}\nSee the module docstring for required packages.', + file=sys.stderr) + sys.exit(2) + + print('Loading Gemma4Config...', flush=True) + hf_cfg = Gemma4Config.from_pretrained(args.ckpt) + vc = hf_cfg.vision_config + rope_theta = (vc.rope_parameters or {}).get('rope_theta', 100.0) + vision_cfg = vision_real.Gemma4VisionConfig( + hidden_size=vc.hidden_size, + intermediate_size=vc.intermediate_size, + num_hidden_layers=vc.num_hidden_layers, + num_attention_heads=vc.num_attention_heads, + num_key_value_heads=vc.num_key_value_heads, + head_dim=vc.head_dim, + rms_norm_eps=vc.rms_norm_eps, + patch_size=vc.patch_size, + position_embedding_size=vc.position_embedding_size, + pooling_kernel_size=vc.pooling_kernel_size, + rope_theta=rope_theta, + use_clipped_linears=vc.use_clipped_linears, + standardize=vc.standardize, + param_dtype=jnp.bfloat16, + dtype=jnp.bfloat16, + ) + # Auto-detect the Gemma4 text variant from the checkpoint, keyed by + # (num_hidden_layers, hidden_size), and pick the matching tunix preset. + tc = hf_cfg.text_config + variants = { + (35, 1536): ('E2B', model_lib.ModelConfig.gemma4_e2b), + (42, 2560): ('E4B', model_lib.ModelConfig.gemma4_e4b), + (60, 5376): ('31B', model_lib.ModelConfig.gemma4_31b), + (30, 2816): ('26B-A4B', model_lib.ModelConfig.gemma4_26b_a4b), + } + key = (tc.num_hidden_layers, tc.hidden_size) + if key not in variants: + print(f'ERROR: no tunix ModelConfig preset for this Gemma4 text variant ' + f'(num_hidden_layers={key[0]}, hidden_size={key[1]}). ' + f'Known variants: {[v[0] for v in variants.values()]}. ' + f'(The 12B model has no preset yet.)', file=sys.stderr) + sys.exit(2) + variant_name, cfg_factory = variants[key] + print(f'Detected Gemma4 text variant: {variant_name}', flush=True) + + text_cfg = cfg_factory() + # Run computation in bfloat16 to match weights and KV cache dtype. + text_cfg.dtype = jnp.bfloat16 + text_cfg.param_dtype = jnp.bfloat16 + # Inference (Megatron-style) sharding: residual stream replicated, inner + # attention/FFN dims sharded along 'tp'. is_sampling=True drops the FSDP + # axis used in training. + text_cfg.shd_config = model_lib.ShardingConfig.get_default_sharding( + is_sampling=True + ) + + print('Loading checkpoint (text + vision)...', flush=True) + t0 = time.time() + model = mm_lib.create_multimodal_from_safe_tensors( + args.ckpt, text_cfg, vision_cfg, + image_token_id=hf_cfg.image_token_id, mesh=mesh, dtype=jnp.bfloat16, + ) + print(f' checkpoint loaded in {time.time() - t0:.1f}s', flush=True) + tokenizer = AutoTokenizer.from_pretrained(args.ckpt) + + proc = ip.Gemma4ImageProcessor( + patch_size=vc.patch_size, pooling_kernel_size=vc.pooling_kernel_size + ) + img = np.transpose( + np.asarray(Image.open(args.image).convert('RGB')), (2, 0, 1) + ) + px, pos_ids, num_soft = _unpadded_patches(proc, img) + print(f'Image: {num_soft} soft tokens ({px.shape[1]} raw patches)', + flush=True) + + bos_id = tokenizer.bos_token_id or 2 + tokens = build_image_prompt_tokens( + tokenizer, args.prompt, hf_cfg.image_token_id, num_soft, bos_id + ) + # Collect all EOS token IDs (Gemma4-IT uses both token 1 and token 106 as stop signals) + raw_eos = tokenizer.eos_token_id + if isinstance(raw_eos, (list, tuple)): + eos_ids = set(raw_eos) + elif raw_eos is not None: + eos_ids = {raw_eos} + else: + eos_ids = {1} + eot = tokenizer.convert_tokens_to_ids('') + if eot and eot != tokenizer.unk_token_id: + eos_ids.add(eot) + print(f'Prompt: {args.prompt!r} (total prompt length: {tokens.shape[1]} tokens)', + flush=True) + + # Enter the mesh context so the activation-level shard() constraints inside + # the model forward resolve against our tensor-parallel mesh. + with jax.set_mesh(mesh): + out_ids = greedy_generate( + model, tokens, px, pos_ids, + max_new_tokens=args.max_new_tokens, + eos_ids=eos_ids, + tokenizer=tokenizer, + ) + + print('\n=== Caption ===') + print(tokenizer.decode(out_ids, skip_special_tokens=True)) + + +if __name__ == '__main__': + main() diff --git a/examples/gemma4/multimodal_parity_random_weights.py b/examples/gemma4/multimodal_parity_random_weights.py new file mode 100644 index 000000000..23d5dc737 --- /dev/null +++ b/examples/gemma4/multimodal_parity_random_weights.py @@ -0,0 +1,172 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Checkpoint-free full-model parity: HF Gemma4ForConditionalGeneration vs +Tunix `Gemma4Multimodal`. + +Builds a tiny random HF model (vision + text + projector + LM head), saves its +state dict as safetensors, loads into Tunix via +`create_multimodal_from_safe_tensors`, feeds identical (tokens, image) inputs, +and diffs per-position logits in fp32. Also captures the PLE values handed to +the text model on each side and compares them. + +How to read the output (see docs/gemma4_vision_port.md for the long version): + + * BOS position SHOULD be exact (~1e-7). It validates the embedding + + first-position attention + MLP path. + * If image positions diverge but the PLE-y-embed diff is zero, the divergence + is in the merge / per-layer-projection-of-merged-embeds path. + * If text positions diverge AND a pure-text parity (no image) also diverges, + you've hit the pre-existing text-model divergence in `tunix.models.gemma4` + -- out of scope for the vision PR. + +Run on any machine with torch + transformers (no weights, no GPU needed): + + pip install torch 'transformers>=5.9' safetensors + python examples/gemma4/multimodal_parity_random_weights.py +""" + +from __future__ import annotations + +import argparse +import os +import sys +import tempfile + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + try: + import numpy as np + import torch + from safetensors.numpy import save_file + from transformers.models.gemma4.configuration_gemma4 import ( + Gemma4Config, Gemma4TextConfig, Gemma4VisionConfig) + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4ForConditionalGeneration) + except ImportError as e: + print(f"ERROR: {e}\nRun: pip install torch 'transformers>=5.9' safetensors", + file=sys.stderr) + sys.exit(2) + + import jax.numpy as jnp + from tunix.models.gemma4 import model as tml + from tunix.models.gemma4 import multimodal as mm + from tunix.models.gemma4 import vision_real + + PAD, BOS, IMG_TOK, PER = 0, 2, 5, 8 + torch.manual_seed(args.seed) + tc = Gemma4TextConfig( + vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1, + global_head_dim=16, attention_k_eq_v=False, + num_attention_heads=4, num_key_value_heads=2, head_dim=16, + hidden_size_per_layer_input=PER, vocab_size_per_layer_input=64, + num_kv_shared_layers=0, final_logit_softcapping=None, + pad_token_id=PAD, bos_token_id=BOS, + ) + vc = Gemma4VisionConfig( + hidden_size=32, intermediate_size=64, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=4, head_dim=8, + patch_size=4, position_embedding_size=64, pooling_kernel_size=2, + ) + cfg = Gemma4Config(text_config=tc, vision_config=vc, audio_config=None, + image_token_id=IMG_TOK) + hf = Gemma4ForConditionalGeneration(cfg).eval() + + # Save HF state dict as safetensors with real key names. + d = tempfile.mkdtemp() + save_file({k: v.detach().to(torch.float32).numpy() + for k, v in hf.state_dict().items()}, + os.path.join(d, "model.safetensors")) + + # Load into Tunix. + text_cfg = tml.ModelConfig.gemma4_e2b() + text_cfg.num_layers = 1; text_cfg.num_embed = 64; text_cfg.embed_dim = 64 + text_cfg.hidden_dim = 128; text_cfg.num_heads = 4; text_cfg.num_kv_heads = 2 + text_cfg.head_dim = 16; text_cfg.per_layer_input_dim = PER + text_cfg.frac_shared_layers = 0.0; text_cfg.final_logit_softcap = None + text_cfg.param_dtype = jnp.float32; text_cfg.dtype = jnp.float32 + + vis_cfg = vision_real.Gemma4VisionConfig( + hidden_size=32, intermediate_size=64, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=4, head_dim=8, patch_size=4, + position_embedding_size=64, pooling_kernel_size=2, + param_dtype=jnp.float32, dtype=jnp.float32, + ) + model = mm.create_multimodal_from_safe_tensors( + d, text_cfg, vis_cfg, + image_token_id=IMG_TOK, pad_token_id=PAD, dtype=jnp.float32, + ) + + # Inputs: BOS + 4 image placeholders + 3 text tokens; 4x4 patch grid -> 4 soft tokens. + tokens = np.array([[BOS, IMG_TOK, IMG_TOK, IMG_TOK, IMG_TOK, 10, 11, 12]], + dtype=np.int64) + side = 4 + xs, ys = np.meshgrid(np.arange(side), np.arange(side), indexing="xy") + pos = np.stack([xs.reshape(-1), ys.reshape(-1)], -1)[None].astype(np.int64) + px = np.random.default_rng(args.seed + 1).random( + (1, side * side, 3 * vc.patch_size**2), dtype=np.float32) + + # Capture HF's per_layer_inputs (the y_embed handed to language_model). + ple_taps = {} + def hook(module, args_, kwargs_): + ple_taps["per_layer_inputs"] = kwargs_.get("per_layer_inputs") + hf.model.language_model.register_forward_pre_hook(hook, with_kwargs=True) + + with torch.no_grad(): + hf_out = hf(input_ids=torch.from_numpy(tokens), + pixel_values=torch.from_numpy(px), + image_position_ids=torch.from_numpy(pos)) + hf_logits = hf_out.logits.float().numpy() + hf_ple = ple_taps["per_layer_inputs"].float().numpy() + + # Tunix forward + PLE-y-embed for direct apples-to-apples comparison. + tk_j = jnp.asarray(tokens.astype(np.int32)) + llm_tokens = jnp.where(tk_j == IMG_TOK, PAD, tk_j) + emb_tab = model.text_model.embedder.per_layer_input_embedding.value + tunix_y = np.asarray( + (emb_tab[llm_tokens] + * jnp.sqrt(text_cfg.per_layer_input_dim).astype(emb_tab.dtype) + ).astype(jnp.float32)) + + tunix_logits = np.asarray( + model(tk_j, jnp.asarray(px), jnp.asarray(pos.astype(np.int32)))[0] + .astype(jnp.float32)) + + print("== PLE y_embed (HF.per_layer_inputs at language_model boundary) ==") + d_ple = np.abs(hf_ple - tunix_y) + print(f" max={d_ple.max():.3e} mean={d_ple.mean():.3e} " + f"-> {'OK (bit-exact)' if d_ple.max() < 1e-5 else 'DIVERGENT'}") + + print("\n== logits per position ==") + for i in range(tokens.shape[1]): + di = np.abs(hf_logits[0, i] - tunix_logits[0, i]) + tag = "IMG" if tokens[0, i] == IMG_TOK else ("TXT" if i > 4 else "BOS") + print(f" pos {i} ({tag}, tok={tokens[0,i]:3d}) max={di.max():.3e} " + f"mean={di.mean():.3e}") + + diff = np.abs(hf_logits - tunix_logits) + print(f"\noverall logit diff: max={diff.max():.3e} mean={diff.mean():.3e}") + print( + "\nNote: BOS (pos 0) parity isolates the text-model basics; non-zero diff\n" + "at later positions that also appears in pure-text parity is a\n" + "pre-existing Tunix-vs-HF text-model issue (out of scope for the vision\n" + "port).") + + +if __name__ == "__main__": + main() diff --git a/examples/gemma4/vision_parity_check.py b/examples/gemma4/vision_parity_check.py new file mode 100644 index 000000000..3fb637829 --- /dev/null +++ b/examples/gemma4/vision_parity_check.py @@ -0,0 +1,343 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Stage 3 — numeric parity harness for the gemma4_vision JAX port. + +Compares per-layer activations between: + * HF torch reference: ``transformers.models.gemma4.modeling_gemma4.Gemma4VisionModel`` + + ``Gemma4MultimodalEmbedder``, loaded from a real ``google/gemma-4-*-it`` + safetensors checkpoint. + * This port: ``tunix.models.gemma4.vision_real.Gemma4VisionStack``, loaded + from the same weights via + ``vision_params_safetensors.create_vision_stack_from_safe_tensors``. + +Both sides receive identical synthetic ``pixel_values`` + ``pixel_position_ids`` +with NO PADDING, so any divergence is attributable to the port (and not, e.g., +torchvision-vs-PIL resize). All comparisons are cast to fp32 for clean diffs. + +Run: + pip install torch transformers safetensors + python examples/gemma4/vision_parity_check.py --ckpt ~/gemma4-e2b + +Exit code 0 iff every checkpoint passes (max-abs diff < tolerance). +""" + +from __future__ import annotations + +import argparse +import glob +import os +import sys + +import numpy as np + + +# bf16 has ~3 decimal digits; after 16 sandwich layers compounded error around +# 1e-2 is normal. Set the per-checkpoint failure threshold accordingly. +_DEFAULT_TOL = 5e-2 + + +def _require_torch() -> bool: + try: + import torch # noqa: F401 + from safetensors.torch import load_file # noqa: F401 + from transformers import Gemma4Config # noqa: F401 + from transformers.models.gemma4.modeling_gemma4 import ( # noqa: F401 + Gemma4MultimodalEmbedder, + Gemma4VisionModel, + ) + return True + except ImportError as e: # pragma: no cover - environment check + print( + f"ERROR: missing dependency for the torch reference: {e}\n" + "Run: pip install torch transformers safetensors", + file=sys.stderr, + ) + return False + + +def make_inputs(num_soft: int, pool_k: int, patch_size: int, seed: int): + """Build a square, all-valid input grid (no padding).""" + side_pool = int(round(num_soft**0.5)) + if side_pool * side_pool != num_soft: + raise ValueError(f"--num-soft must be a perfect square (got {num_soft})") + side_patch = side_pool * pool_k # patches per side + num_patches = side_patch * side_patch + in_channels = 3 * patch_size * patch_size + + rng = np.random.default_rng(seed) + pixel_values = rng.random((1, num_patches, in_channels), dtype=np.float32) + xs, ys = np.meshgrid( + np.arange(side_patch), np.arange(side_patch), indexing="xy" + ) + positions = np.stack([xs.reshape(-1), ys.reshape(-1)], axis=-1)[None] + return pixel_values, positions.astype(np.int64) + + +# --------------------------------------------------------------------------- +# Torch reference +# --------------------------------------------------------------------------- +def torch_reference(ckpt_dir: str, pixel_values, pixel_position_ids): + """Loads vision + embed_vision from safetensors and runs a hooked forward.""" + import torch + from safetensors.torch import load_file + from transformers import Gemma4Config + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4MultimodalEmbedder, + Gemma4VisionModel, + ) + + ckpt = os.path.expanduser(ckpt_dir) + cfg = Gemma4Config.from_pretrained(ckpt) + vision_model = Gemma4VisionModel(cfg.vision_config).eval() + embed_vision = Gemma4MultimodalEmbedder(cfg.vision_config, cfg.text_config).eval() + + # Selective state-dict load: only model.vision_tower.* and model.embed_vision.*. + combined = {} + for fn in sorted(glob.glob(os.path.join(ckpt, "*.safetensors"))): + combined.update(load_file(fn)) + v_prefix = "model.vision_tower." + e_prefix = "model.embed_vision." + vision_state = { + k[len(v_prefix):]: v for k, v in combined.items() if k.startswith(v_prefix) + } + embed_state = { + k[len(e_prefix):]: v for k, v in combined.items() if k.startswith(e_prefix) + } + missing_v, unexpected_v = vision_model.load_state_dict(vision_state, strict=False) + missing_e, unexpected_e = embed_vision.load_state_dict(embed_state, strict=False) + if missing_v or unexpected_v: + # `_extra_state` / clip-buffer ±inf inits and config flags can show up here; + # surface counts but don't bail unless real weights are missing. + print( + f" [hf] vision_tower: missing={len(missing_v)} unexpected={len(unexpected_v)}", + file=sys.stderr, + ) + if missing_e or unexpected_e: + print( + f" [hf] embed_vision: missing={len(missing_e)} unexpected={len(unexpected_e)}", + file=sys.stderr, + ) + + param_dtype = next(vision_model.parameters()).dtype + taps: dict[str, np.ndarray] = {} + + def to_np(t): + return t.detach().to(torch.float32).cpu().numpy() + + def hook(name): + def fn(module, inp, out): + t = out + if isinstance(out, tuple): + t = out[0] + if hasattr(t, "last_hidden_state"): + t = t.last_hidden_state + taps[name] = to_np(t) + return fn + + vision_model.patch_embedder.register_forward_hook(hook("after_patch_embed")) + for i, layer in enumerate(vision_model.encoder.layers): + layer.register_forward_hook(hook(f"after_layer_{i:02d}")) + + def pool_hook(module, inp, out): + pooled, mask = out + taps["after_pool"] = to_np(pooled) + taps["after_pool_mask"] = mask.detach().cpu().numpy() + vision_model.pooler.register_forward_hook(pool_hook) + + with torch.no_grad(): + px_t = torch.from_numpy(pixel_values).to(param_dtype) + pos_t = torch.from_numpy(pixel_position_ids).long() + vis_out = vision_model(px_t, pos_t) + taps["after_vision_tower_gathered"] = to_np(vis_out.last_hidden_state) + proj = embed_vision(vis_out.last_hidden_state) + taps["after_projector"] = to_np(proj) + return taps + + +# --------------------------------------------------------------------------- +# JAX port path +# --------------------------------------------------------------------------- +def jax_port(ckpt_dir: str, pixel_values, pixel_position_ids): + """Instrumented forward through the ported Gemma4VisionStack.""" + import jax + import jax.numpy as jnp + from transformers import Gemma4Config + from tunix.models.gemma4 import vision_params_safetensors as vp + from tunix.models.gemma4 import vision_real + + ckpt = os.path.expanduser(ckpt_dir) + hf_cfg = Gemma4Config.from_pretrained(ckpt) + vc = hf_cfg.vision_config + rope_theta = (vc.rope_parameters or {}).get("rope_theta", 100.0) + + cfg = vision_real.Gemma4VisionConfig( + hidden_size=vc.hidden_size, + intermediate_size=vc.intermediate_size, + num_hidden_layers=vc.num_hidden_layers, + num_attention_heads=vc.num_attention_heads, + num_key_value_heads=vc.num_key_value_heads, + head_dim=vc.head_dim, + rms_norm_eps=vc.rms_norm_eps, + patch_size=vc.patch_size, + position_embedding_size=vc.position_embedding_size, + pooling_kernel_size=vc.pooling_kernel_size, + rope_theta=rope_theta, + use_clipped_linears=vc.use_clipped_linears, + standardize=vc.standardize, + param_dtype=jnp.bfloat16, + dtype=jnp.bfloat16, + ) + stack = vp.create_vision_stack_from_safe_tensors( + file_dir=ckpt, + config=cfg, + text_hidden_size=hf_cfg.text_config.hidden_size, + dtype=jnp.bfloat16, + ) + + taps: dict[str, np.ndarray] = {} + px = jnp.asarray(pixel_values, dtype=jnp.bfloat16) + pos = jnp.asarray(pixel_position_ids, dtype=jnp.int32) + padding_positions = jnp.all(pos == -1, axis=-1) + valid_mask = jnp.logical_not(padding_positions) + + def cap(name, arr): + taps[name] = np.asarray(arr.astype(jnp.float32)) + + # 1. patch embedder + pe = stack.vision_tower.patch_embedder(px, pos, padding_positions) + cap("after_patch_embed", pe) + + # 2. encoder, layer-by-layer + attn_bias = jnp.where( + valid_mask[:, None, None, :], 0.0, jnp.finfo(jnp.float32).min + ) + cos, sin = stack.vision_tower.encoder.rotary_emb(pos) + x = pe + for i, layer in enumerate(stack.vision_tower.encoder.layers): + x = layer(x, cos, sin, pos, attn_bias) + cap(f"after_layer_{i:02d}", x) + + # 3. pooler + pk = cfg.pooling_kernel_size + output_length = px.shape[-2] // (pk * pk) + pooled, mask = stack.vision_tower.pooler( + x, pos, padding_positions, output_length + ) + cap("after_pool", pooled) + taps["after_pool_mask"] = np.asarray(mask) + + # 4. gather valid rows (mirrors HF `hidden_states[pooler_mask]`). + flat_pool = pooled.reshape(-1, pooled.shape[-1]) + flat_mask = mask.reshape(-1) + cap("after_vision_tower_gathered", flat_pool[flat_mask]) + + # 5. projector (embed_vision) on the gathered tensor, to match HF call site. + proj = stack.embed_vision(flat_pool[flat_mask]) + cap("after_projector", proj) + return taps + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- +def _diff_stats(a: np.ndarray, b: np.ndarray) -> tuple[float, float, float]: + d = np.abs(a.astype(np.float32) - b.astype(np.float32)) + denom = np.maximum(np.abs(a.astype(np.float32)), 1e-6) + return float(d.max()), float(d.mean()), float(np.median(d / denom)) + + +def report(name: str, hf: np.ndarray, jx: np.ndarray, tol: float) -> bool: + if hf.shape != jx.shape: + print(f" {name:36s} SHAPE MISMATCH hf={hf.shape} jx={jx.shape}") + return False + max_abs, mean_abs, median_rel = _diff_stats(hf, jx) + ok = max_abs < tol + print( + f" {name:36s} {'OK ' if ok else 'FAIL'}" + f" max_abs={max_abs:.3e} mean_abs={mean_abs:.3e}" + f" median_rel={median_rel:.3e}" + ) + return ok + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--ckpt", required=True, help="Gemma-4 checkpoint directory") + ap.add_argument( + "--num-soft", + type=int, + default=4, + help="Soft tokens per image (perfect square; 4 means 6x6 patches).", + ) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument( + "--tol", type=float, default=_DEFAULT_TOL, help="Per-checkpoint max-abs threshold." + ) + args = ap.parse_args() + + if not _require_torch(): + sys.exit(2) + + from transformers import Gemma4Config + + hf_cfg = Gemma4Config.from_pretrained(os.path.expanduser(args.ckpt)) + pool_k = hf_cfg.vision_config.pooling_kernel_size + patch = hf_cfg.vision_config.patch_size + print( + f"Inputs: num_soft={args.num_soft}, pool_k={pool_k}, patch={patch}, " + f"layers={hf_cfg.vision_config.num_hidden_layers}, seed={args.seed}" + ) + px, pos = make_inputs(args.num_soft, pool_k, patch, args.seed) + print(f" pixel_values={px.shape} pixel_position_ids={pos.shape}\n") + + print("== HF torch reference ==", flush=True) + hf = torch_reference(args.ckpt, px, pos) + print(f" captured {len(hf)} taps\n") + + print("== JAX port ==", flush=True) + jx = jax_port(args.ckpt, px, pos) + print(f" captured {len(jx)} taps\n") + + print(f"== Parity report (tol={args.tol:.0e}, fp32 cast) ==") + passed = True + ordered = ( + ["after_patch_embed"] + + [f"after_layer_{i:02d}" for i in range(hf_cfg.vision_config.num_hidden_layers)] + + ["after_pool", "after_vision_tower_gathered", "after_projector"] + ) + for k in ordered: + if k not in hf or k not in jx: + print(f" {k:36s} MISSING from {'hf' if k not in hf else 'jx'}") + passed = False + continue + passed &= report(k, hf[k], jx[k], args.tol) + + # Mask sanity (must be identical): + if "after_pool_mask" in hf and "after_pool_mask" in jx: + same = np.array_equal(hf["after_pool_mask"].reshape(-1), jx["after_pool_mask"].reshape(-1)) + print(f" {'pool_mask':36s} {'OK ' if same else 'FAIL'} identical={same}") + passed &= same + + print() + if passed: + print("PARITY PASSED -- the JAX port matches HF within tolerance.") + sys.exit(0) + else: + print("PARITY FAILED -- see the first FAIL row to localize the bug.") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/gemma4/vision_parity_random_weights.py b/examples/gemma4/vision_parity_random_weights.py new file mode 100644 index 000000000..dfa84f4ff --- /dev/null +++ b/examples/gemma4/vision_parity_random_weights.py @@ -0,0 +1,167 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Checkpoint-free numeric parity for the gemma4_vision JAX port. + +Unlike `vision_parity_check.py` (which needs a real checkpoint), this builds a +small HF `Gemma4VisionModel` + `Gemma4MultimodalEmbedder` with RANDOM weights, +serializes them to a temporary safetensors file with real checkpoint key names, +loads them into the JAX `Gemma4VisionStack` via the production loader, and +compares per-layer activations. It therefore validates the PORT MATH + the key +mapping + the loader in one shot, on any machine with torch + a transformers +that ships Gemma 4 -- no model download required. + +Everything runs in fp32 so the tolerance is tight (default 1e-3). + +Run: + pip install torch 'transformers>=5.9' # must include models/gemma4 + python examples/gemma4/vision_parity_random_weights.py + +NOTE: if `transformers` fails to import gemma4 with a `torch._dynamo` +"Duplicate dispatch rule" error, your transformers/torch combo trips a torch +dynamo packaging bug. `pip install 'transformers==5.9.0'` is a known-good combo. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import tempfile + + +def main(): + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--layers", type=int, default=2) + ap.add_argument("--tol", type=float, default=1e-3) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + try: + import numpy as np + import torch + from safetensors.numpy import save_file + from transformers.models.gemma4.configuration_gemma4 import ( + Gemma4TextConfig, + Gemma4VisionConfig, + ) + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4MultimodalEmbedder, + Gemma4VisionModel, + ) + except ImportError as e: + print( + f"ERROR: {e}\nRun: pip install torch 'transformers>=5.9' safetensors\n" + "(If gemma4 import hits a torch._dynamo 'Duplicate dispatch rule' error, " + "use transformers==5.9.0.)", + file=sys.stderr, + ) + sys.exit(2) + + import jax.numpy as jnp + from tunix.models.gemma4 import vision_params_safetensors as vp + from tunix.models.gemma4 import vision_real + + torch.manual_seed(args.seed) + # Small but structurally identical to the real tower. + patch = 4 + vc = Gemma4VisionConfig( + hidden_size=64, intermediate_size=128, num_hidden_layers=args.layers, + num_attention_heads=4, num_key_value_heads=4, head_dim=16, patch_size=patch, + position_embedding_size=256, pooling_kernel_size=2, rms_norm_eps=1e-6, + ) + tc = Gemma4TextConfig(hidden_size=48) + vm = Gemma4VisionModel(vc).eval() + ev = Gemma4MultimodalEmbedder(vc, tc).eval() + + # Serialize random weights under real checkpoint key names. + ckpt = {} + for k, v in vm.state_dict().items(): + ckpt[f"model.vision_tower.{k}"] = v.detach().to(torch.float32).numpy() + for k, v in ev.state_dict().items(): + ckpt[f"model.embed_vision.{k}"] = v.detach().to(torch.float32).numpy() + d = tempfile.mkdtemp() + save_file(ckpt, os.path.join(d, "model.safetensors")) + + # No-padding square input -> HF gather is identity. + side = 4 + xs, ys = np.meshgrid(np.arange(side), np.arange(side), indexing="xy") + pos = np.stack([xs.reshape(-1), ys.reshape(-1)], axis=-1)[None].astype(np.int64) + px = np.random.default_rng(args.seed + 1).random( + (1, side * side, 3 * patch**2), dtype=np.float32 + ) + + # --- HF reference with per-layer hooks --- + hf = {} + vm.patch_embedder.register_forward_hook( + lambda m, i, o: hf.update(after_patch_embed=o.detach().float().numpy()) + ) + for idx, layer in enumerate(vm.encoder.layers): + layer.register_forward_hook( + (lambda j: (lambda m, i, o: hf.update( + **{f"after_layer_{j:02d}": (o[0] if isinstance(o, tuple) else o).detach().float().numpy()} + )))(idx) + ) + with torch.no_grad(): + out = vm(torch.from_numpy(px), torch.from_numpy(pos)) + hf["tower"] = out.last_hidden_state.float().numpy() + hf["proj"] = ev(out.last_hidden_state).float().numpy() + + # --- JAX port through the production loader --- + cfg = vision_real.Gemma4VisionConfig( + hidden_size=64, intermediate_size=128, num_hidden_layers=args.layers, + num_attention_heads=4, num_key_value_heads=4, head_dim=16, patch_size=patch, + position_embedding_size=256, pooling_kernel_size=2, rms_norm_eps=1e-6, + param_dtype=jnp.float32, dtype=jnp.float32, + ) + stack = vp.create_vision_stack_from_safe_tensors(d, cfg, text_hidden_size=48, dtype=jnp.float32) + pos_j, px_j = jnp.asarray(pos, jnp.int32), jnp.asarray(px, jnp.float32) + pad = jnp.all(pos_j == -1, axis=-1) + valid = jnp.logical_not(pad) + + jx = {} + pe = stack.vision_tower.patch_embedder(px_j, pos_j, pad) + jx["after_patch_embed"] = np.asarray(pe.astype(jnp.float32)) + bias = jnp.where(valid[:, None, None, :], 0.0, jnp.finfo(jnp.float32).min) + cos, sin = stack.vision_tower.encoder.rotary_emb(pos_j) + x = pe + for idx, layer in enumerate(stack.vision_tower.encoder.layers): + x = layer(x, cos, sin, pos_j, bias) + jx[f"after_layer_{idx:02d}"] = np.asarray(x.astype(jnp.float32)) + output_length = px_j.shape[-2] // (cfg.pooling_kernel_size**2) + pooled, mask = stack.vision_tower.pooler(x, pos_j, pad, output_length) + flat = pooled.reshape(-1, pooled.shape[-1])[mask.reshape(-1)] + jx["tower"] = np.asarray(flat.astype(jnp.float32)) + jx["proj"] = np.asarray(stack.embed_vision(flat).astype(jnp.float32)) + + # --- report --- + print(f"== checkpoint-free parity (fp32, tol={args.tol:.0e}) ==") + order = ( + ["after_patch_embed"] + + [f"after_layer_{i:02d}" for i in range(args.layers)] + + ["tower", "proj"] + ) + ok = True + for k in order: + a, b = hf[k], jx[k] + m = float(np.abs(a - b).max()) + passed = a.shape == b.shape and m < args.tol + ok &= passed + print(f" {k:20s} {'OK ' if passed else 'FAIL'} max_abs={m:.3e} shapes hf{a.shape} jx{b.shape}") + print("\nPARITY PASSED" if ok else "\nPARITY FAILED") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/models/gemma4/image_processing_test.py b/tests/models/gemma4/image_processing_test.py new file mode 100644 index 000000000..09bb3af78 --- /dev/null +++ b/tests/models/gemma4/image_processing_test.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Tests for the Gemma 4 image processor port (deterministic parts only).""" + +from __future__ import annotations + +import jax.numpy as jnp +import numpy as np +from absl.testing import absltest +from flax import nnx +from tunix.models.gemma4 import image_processing as ip +from tunix.models.gemma4 import vision_real + + +class ImageProcessingTest(absltest.TestCase): + + def test_target_size_divisible_and_within_budget(self): + th, tw = ip.get_aspect_ratio_preserving_size( + height=480, width=640, patch_size=16, max_patches=2520, pooling_kernel_size=3 + ) + self.assertEqual(th % 48, 0) + self.assertEqual(tw % 48, 0) + self.assertLessEqual((th // 16) * (tw // 16), 2520) + + def test_patchify_shape_and_order(self): + # 2x3 patch grid, patch_size=2, single channel ramp to check ordering. + patch_size = 2 + img = np.arange(1 * 4 * 6).reshape(1, 4, 6).astype(np.float32) # (C,H,W) + patches = ip.convert_image_to_patches(img, patch_size) + self.assertEqual(patches.shape, (2 * 3, patch_size * patch_size * 1)) + # First patch is the top-left 2x2 block: rows [0,1] cols [0,1]. + np.testing.assert_array_equal(patches[0], np.array([0, 1, 6, 7])) + + def test_position_ids_are_xy_and_match_patch_order(self): + pos = ip.build_position_ids(patch_height=2, patch_width=3) + self.assertEqual(pos.shape, (6, 2)) + # Row-major over (row, col); stored as (x=col, y=row). + np.testing.assert_array_equal( + pos, np.array([[0, 0], [1, 0], [2, 0], [0, 1], [1, 1], [2, 1]]) + ) + + def test_padding_uses_minus_one_for_positions(self): + patches = np.ones((4, 12), dtype=np.float32) + positions = np.zeros((4, 2), dtype=np.int64) + patches, positions = ip.pad_along_first_dim(patches, positions, target_length=6) + self.assertEqual(patches.shape, (6, 12)) + np.testing.assert_array_equal(patches[4:], 0) + np.testing.assert_array_equal(positions[4:], -1) + + def test_end_to_end_into_vision_stack(self): + # A 48x48 image -> 3x3 patches (9 patches) -> pool 3x3 -> 1 soft token. + proc = ip.Gemma4ImageProcessor(patch_size=16, pooling_kernel_size=3, max_soft_tokens=32) + img = (np.random.rand(3, 48, 48) * 255).astype(np.uint8) + px, pos, num_soft = proc([img], do_resize=False) + + self.assertEqual(px.shape, (1, proc.max_patches, 3 * 16**2)) + self.assertEqual(pos.shape, (1, proc.max_patches, 2)) + self.assertEqual(num_soft, [1]) + + cfg = vision_real.Gemma4VisionConfig(num_hidden_layers=1) + stack = vision_real.Gemma4VisionStack(cfg, text_hidden_size=1536, rngs=nnx.Rngs(0)) + soft, mask = stack(jnp.asarray(px), jnp.asarray(pos)) + # max_patches=288 -> /9 = 32 pooled tokens; mask marks the 1 real soft token. + self.assertEqual(soft.shape, (1, proc.max_patches // 9, 1536)) + self.assertEqual(int(mask.sum()), 1) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/models/gemma4/model_test.py b/tests/models/gemma4/model_test.py index fdf8b457e..3449361bc 100644 --- a/tests/models/gemma4/model_test.py +++ b/tests/models/gemma4/model_test.py @@ -16,10 +16,13 @@ from __future__ import annotations +import dataclasses + from absl.testing import absltest from flax import nnx import jax import jax.numpy as jnp +from tunix.models.gemma3 import vision as vision_lib from tunix.models.gemma4 import model as model_lib @@ -195,6 +198,144 @@ def body_fn(step, _): _, logits = compiled_decode(state) self.assertEqual(logits.shape, (2, 32, config.num_embed)) + def test_text_only_no_vision_encoder(self): + config = model_lib.ModelConfig.gemma4_e2b() + self.assertIsNone(config.vision_config) + config.num_layers = 1 + config.embed_dim = 256 + config.hidden_dim = 512 + config.num_heads = 4 + config.head_dim = 64 + config.num_kv_heads = 1 + config.frac_shared_layers = 0.0 + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs) + self.assertIsNone(model.vision_encoder) + self.assertFalse(hasattr(model.embedder, "mm_input_projection")) + + def test_forward_pass_multimodal(self): + # Use a tiny SigLIP config so the test stays fast/light. + small_vision_config = vision_lib.SigLIPConfig( + num_mm_tokens_per_image_prepool=16, + num_mm_tokens_per_image=4, + image_height=32, + image_width=32, + image_channels=3, + soft_token_placeholder=219, + patch_size=(8, 8), + width=32, + depth=1, + mlp_dim=64, + num_heads=4, + ) + base_config = model_lib.ModelConfig.gemma4_e2b() + config = dataclasses.replace( + base_config, + num_layers=1, + embed_dim=256, + hidden_dim=512, + num_heads=4, + head_dim=64, + num_kv_heads=1, + frac_shared_layers=0.0, + vision_config=small_vision_config, + ) + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs) + self.assertIsNotNone(model.vision_encoder) + self.assertTrue(hasattr(model.embedder, "mm_input_projection")) + self.assertTrue(hasattr(model.embedder, "mm_soft_embedding_norm")) + + batch_size = 2 + seq_len = 16 + num_images = 1 + num_mm = small_vision_config.num_mm_tokens_per_image + + # Place the soft-image placeholder in positions [1:1+num_mm] of each row. + tokens = jnp.ones((batch_size, seq_len), dtype=jnp.int32) + tokens = tokens.at[:, 1 : 1 + num_mm].set( + small_vision_config.soft_token_placeholder + ) + + positions = jnp.tile( + jnp.arange(seq_len)[None, :], (batch_size, 1) + ) + attn_mask = model.get_attention_mask(tokens) + + images = jnp.zeros( + ( + batch_size, + num_images, + small_vision_config.image_height, + small_vision_config.image_width, + small_vision_config.image_channels, + ), + dtype=jnp.float32, + ) + + logits, _ = model( + tokens, + positions=positions, + attention_mask=attn_mask, + images=images, + ) + self.assertEqual(logits.shape, (batch_size, seq_len, config.num_embed)) + + def test_multimodal_call_without_vision_encoder_raises(self): + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 1 + config.embed_dim = 256 + config.hidden_dim = 512 + config.num_heads = 4 + config.head_dim = 64 + config.num_kv_heads = 1 + config.frac_shared_layers = 0.0 + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs) + + tokens = jnp.ones((1, 4), dtype=jnp.int32) + positions = jnp.arange(4)[None, :] + attn_mask = jnp.tril(jnp.ones((4, 4), dtype=jnp.bool_))[None, ...] + images = jnp.zeros((1, 1, 32, 32, 3), dtype=jnp.float32) + + with self.assertRaises(ValueError): + model( + tokens, + positions=positions, + attention_mask=attn_mask, + images=images, + ) + + def test_get_attention_mask_text_only(self): + config = model_lib.ModelConfig.gemma4_e2b() + config.num_layers = 1 + config.embed_dim = 256 + config.hidden_dim = 512 + config.num_heads = 4 + config.head_dim = 64 + config.num_kv_heads = 1 + config.frac_shared_layers = 0.0 + + rngs = nnx.Rngs(0) + model = model_lib.Gemma4(config, rngs=rngs) + # No vision config => no bidirectional span; mask should be purely causal. + tokens = jnp.array([[1, 2, 3, 0, 0]]) + mask = model.get_attention_mask(tokens) + expected = jnp.array( + [[ + [1, 0, 0, 0, 0], + [1, 1, 0, 0, 0], + [1, 1, 1, 0, 0], + [1, 1, 1, 0, 0], + [1, 1, 1, 0, 0], + ]], + dtype=jnp.bool_, + ) + self.assertTrue(bool(jnp.all(mask == expected))) + if __name__ == "__main__": absltest.main() diff --git a/tests/models/gemma4/multimodal_test.py b/tests/models/gemma4/multimodal_test.py new file mode 100644 index 000000000..c52db03d8 --- /dev/null +++ b/tests/models/gemma4/multimodal_test.py @@ -0,0 +1,169 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Stage 4 tests: Gemma4Multimodal wiring (merge + forward). + +JAX-only (no torch/checkpoint). Validates that vision soft tokens are scattered +into image-token positions, that text positions are untouched, and that the +image actually influences the language-model output. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +from absl.testing import absltest +from flax import nnx +from tunix.models.gemma4 import model as model_lib +from tunix.models.gemma4 import multimodal as mm_lib +from tunix.models.gemma4 import vision_real + +_IMAGE_TOKEN_ID = 5 +_NUM_SOFT = 4 # 4x4 patches, pool 2x2 + + +def _build(): + text_cfg = model_lib.ModelConfig.gemma4_e2b() # text_only -> no SigLIP + text_cfg.num_layers = 1 + text_cfg.embed_dim = 64 + text_cfg.hidden_dim = 128 + text_cfg.num_heads = 4 + text_cfg.head_dim = 16 + text_cfg.num_kv_heads = 1 + text_cfg.frac_shared_layers = 0.0 + text_model = model_lib.Gemma4(text_cfg, rngs=nnx.Rngs(0)) + + vcfg = vision_real.Gemma4VisionConfig( + hidden_size=32, intermediate_size=64, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=4, head_dim=8, patch_size=4, + position_embedding_size=64, pooling_kernel_size=2, + ) + vision = vision_real.Gemma4VisionStack( + vcfg, text_hidden_size=text_cfg.embed_dim, rngs=nnx.Rngs(1) + ) + model = mm_lib.Gemma4Multimodal( + text_model, vision, image_token_id=_IMAGE_TOKEN_ID + ) + return model, vcfg + + +def _square_positions(side: int) -> jnp.ndarray: + xs, ys = jnp.meshgrid(jnp.arange(side), jnp.arange(side), indexing="xy") + return jnp.stack([xs.reshape(-1), ys.reshape(-1)], axis=-1)[None] + + +def _image_inputs(vcfg, seed): + side = 4 + pos = _square_positions(side) + px = jax.random.uniform( + jax.random.PRNGKey(seed), (1, side * side, 3 * vcfg.patch_size**2) + ) + return px, pos.astype(jnp.int32) + + +class Gemma4MultimodalTest(absltest.TestCase): + + def test_merge_places_soft_tokens_at_image_positions(self): + model, vcfg = _build() + px, pos = _image_inputs(vcfg, seed=2) + # 4 image placeholders at positions 1..4; text elsewhere. + tokens = jnp.array([[2, 5, 5, 5, 5, 10, 11, 12]], dtype=jnp.int32) + + merged = model.encode_multimodal_inputs(tokens, px, pos) + text_only = model.text_model.embedder.encode(tokens) + soft, _ = model.vision(px, pos) + + self.assertEqual(merged.shape, text_only.shape) + # Image positions must hold the projected soft tokens... + np.testing.assert_allclose( + np.asarray(merged[0, 1:5]), np.asarray(soft[0, :4]), rtol=1e-5, atol=1e-5 + ) + # ...and text positions must be unchanged. + np.testing.assert_allclose( + np.asarray(merged[0, 5:]), np.asarray(text_only[0, 5:]), rtol=1e-5, atol=1e-5 + ) + + def test_forward_shape_and_image_affects_logits(self): + model, vcfg = _build() + tokens = jnp.array([[2, 5, 5, 5, 5, 10, 11, 12]], dtype=jnp.int32) + px_a, pos = _image_inputs(vcfg, seed=2) + px_b, _ = _image_inputs(vcfg, seed=99) + + logits_a, _ = model(tokens, px_a, pos) + logits_b, _ = model(tokens, px_b, pos) + + self.assertEqual(logits_a.shape, (1, 8, model.text_model.config.num_embed)) + self.assertTrue(bool(jnp.all(jnp.isfinite(logits_a)))) + # A text position AFTER the image must react to the image content (causal + # attention from the soft tokens), so different images => different logits. + self.assertFalse( + np.allclose(np.asarray(logits_a[0, -1]), np.asarray(logits_b[0, -1]), + rtol=1e-4, atol=1e-4) + ) + + def test_per_layer_inputs_substitutes_pad_at_image_positions(self): + """HF Gemma4Model treats image positions as `pad_token_id` for the PLE + token-identity lookup. This locks in the matching substitution.""" + model, vcfg = _build() + px, pos = _image_inputs(vcfg, seed=3) + tokens = jnp.array([[2, 5, 5, 5, 5, 10, 11, 12]], dtype=jnp.int32) + + merged = model.encode_multimodal_inputs(tokens, px, pos) + ple = model._compute_per_layer_inputs(tokens, merged) + + # Reference: encode_per_layer_input applied with image_token_id->pad_id ids. + pad = model.pad_token_id + is_image = (tokens == model.image_token_id) + pad_tokens = jnp.where(is_image, pad, tokens) + expected = model.text_model.embedder.encode_per_layer_input(merged, pad_tokens) + np.testing.assert_allclose(np.asarray(ple), np.asarray(expected), + rtol=1e-6, atol=1e-6) + + # Sanity: image positions of `ple` must differ from a PLE computed without + # the substitution (so the fix is doing something). + without = model.text_model.embedder.encode_per_layer_input(merged, tokens) + self.assertFalse(np.allclose(np.asarray(ple[0, 1]), np.asarray(without[0, 1]))) + + def test_bidirectional_image_span_flag_changes_mask(self): + """Flag default (False) -> plain causal; True -> bidirectional over images.""" + model_causal, _ = _build() # default bidirectional_image_span=False + text_model, vision = model_causal.text_model, model_causal.vision + model_bidir = mm_lib.Gemma4Multimodal( + text_model, vision, image_token_id=_IMAGE_TOKEN_ID, + bidirectional_image_span=True, + ) + tokens = jnp.array([[2, 5, 5, 5, 5, 10]], dtype=jnp.int32) + m_c = model_causal.get_attention_mask(tokens)[0] + m_b = model_bidir.get_attention_mask(tokens)[0] + # In bidirectional mode, the 1st image position should see the 4th (future). + self.assertFalse(bool(m_c[1, 4])) + self.assertTrue(bool(m_b[1, 4])) + + def test_attention_mask_bidirectional_when_flag_enabled(self): + text_model, vision = _build()[0].text_model, _build()[0].vision + model = mm_lib.Gemma4Multimodal( + text_model, vision, image_token_id=_IMAGE_TOKEN_ID, + bidirectional_image_span=True, + ) + tokens = jnp.array([[2, 5, 5, 5, 5, 10]], dtype=jnp.int32) + mask = model.get_attention_mask(tokens)[0] + # First image soft-token (pos 1) should attend forward to the rest of the + # image span (pos 2..4) -> bidirectional within the image. + self.assertTrue(bool(mask[1, 4])) + self.assertTrue(bool(mask[5, 4])) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/models/gemma4/vision_params_safetensors_test.py b/tests/models/gemma4/vision_params_safetensors_test.py new file mode 100644 index 000000000..f8c519ee4 --- /dev/null +++ b/tests/models/gemma4/vision_params_safetensors_test.py @@ -0,0 +1,142 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Key-coverage tests for the Gemma 4 vision safetensors mapping. + +Verifies WITHOUT real weights that every real checkpoint vision key maps to an +existing model param, that audio keys and clip buffers are skipped, and that +every loadable model param is covered by exactly one checkpoint key. This is the +check that catches the bugs that broke the SigLIP loader (wrong prefix, the +`.linear.` nesting, missing transposes, uninitialised params). +""" + +from __future__ import annotations + +import jax +from absl.testing import absltest +from flax import nnx +from tunix.models.gemma4 import vision_params_safetensors as vp +from tunix.models.gemma4 import vision_real +from tunix.utils import torch_utils + +_NUM_LAYERS = 2 # structurally identical to the real 16-layer tower + + +def _model_param_keys() -> set[str]: + cfg = vision_real.Gemma4VisionConfig(num_hidden_layers=_NUM_LAYERS) + model = vision_real.Gemma4VisionStack(cfg, text_hidden_size=1536, rngs=nnx.Rngs(0)) + _, state = nnx.split(model) + keys = set() + for path, _ in jax.tree_util.tree_flatten_with_path(state)[0]: + parts = [] + for k in path: + name = getattr(k, "name", None) + if name is None: + idx = getattr(k, "idx", None) + name = idx if idx is not None else getattr(k, "key", k) + if name == "value": # nnx.Param leaf suffix + continue + parts.append(str(name)) + keys.add(".".join(parts)) + return keys + + +def _real_checkpoint_keys(): + """Generate the real `google/gemma-4-*-it` vision/audio/clip key names.""" + vision, skip = [], [] + vision.append("model.vision_tower.patch_embedder.input_proj.weight") + vision.append("model.vision_tower.patch_embedder.position_embedding_table") + vision.append("model.embed_vision.embedding_projection.weight") + for n in range(_NUM_LAYERS): + base = f"model.vision_tower.encoder.layers.{n}" + for proj in ("q_proj", "k_proj", "v_proj", "o_proj"): + vision.append(f"{base}.self_attn.{proj}.linear.weight") + # e2b clip buffers riding alongside each linear -> must be skipped. + for buf in ("input_min", "input_max", "output_min", "output_max"): + skip.append(f"{base}.self_attn.{proj}.{buf}") + for norm in ("q_norm", "k_norm"): + vision.append(f"{base}.self_attn.{norm}.weight") + for proj in ("gate_proj", "up_proj", "down_proj"): + vision.append(f"{base}.mlp.{proj}.linear.weight") + for buf in ("input_min", "input_max", "output_min", "output_max"): + skip.append(f"{base}.mlp.{proj}.{buf}") + for norm in ( + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", + ): + vision.append(f"{base}.{norm}.weight") + # Out-of-scope towers that must be skipped, not errored. + skip.append("model.audio_tower.layers.0.feed_forward1.ffw_layer_1.linear.weight") + skip.append("model.embed_audio.embedding_projection.weight") + return vision, skip + + +class VisionMappingTest(absltest.TestCase): + + def setUp(self): + super().setUp() + self.cfg = vision_real.Gemma4VisionConfig(num_hidden_layers=_NUM_LAYERS) + self.mapping = vp.vision_key_mapping(self.cfg) + self.param_keys = _model_param_keys() + self.vision_keys, self.skip_keys = _real_checkpoint_keys() + + def test_every_vision_key_maps_to_an_existing_param(self): + for k in self.vision_keys: + jax_key, _ = torch_utils.torch_key_to_jax_key(self.mapping, k) + self.assertIn( + jax_key, + self.param_keys, + f"checkpoint key {k!r} -> {jax_key!r} which is not a model param", + ) + + def test_linears_transpose_norms_do_not(self): + jk, t = torch_utils.torch_key_to_jax_key( + self.mapping, "model.vision_tower.encoder.layers.0.mlp.gate_proj.linear.weight" + ) + self.assertEqual(jk, "vision_tower.encoder.layers.0.mlp.gate_proj.linear.kernel") + self.assertEqual(t, ((1, 0), None)) # transpose + jk, t = torch_utils.torch_key_to_jax_key( + self.mapping, "model.vision_tower.encoder.layers.0.input_layernorm.weight" + ) + self.assertEqual(jk, "vision_tower.encoder.layers.0.input_layernorm.scale") + self.assertIsNone(t) # no transform + + def test_audio_and_clip_buffers_are_skipped(self): + # The generic loader skips keys that match no pattern (ValueError -> skip). + for k in self.skip_keys: + with self.assertRaises(ValueError): + torch_utils.torch_key_to_jax_key(self.mapping, k) + + def test_no_double_match(self): + # Each real vision key must match exactly one pattern (loader requirement). + for k in self.vision_keys: + torch_utils.torch_key_to_jax_key(self.mapping, k) # raises if != 1 match + + def test_every_loadable_param_is_covered(self): + # No uninitialised params: every model param (except the derived rope buffer) + # must be the target of exactly one checkpoint key. + mapped = set() + for k in self.vision_keys: + jax_key, _ = torch_utils.torch_key_to_jax_key(self.mapping, k) + mapped.add(jax_key) + uncovered = self.param_keys - mapped + # rotary inv_freq is computed at call time, not stored -> shouldn't appear. + uncovered = {k for k in uncovered if "inv_freq" not in k} + self.assertEqual(uncovered, set(), f"params with no checkpoint key: {uncovered}") + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/models/gemma4/vision_real_test.py b/tests/models/gemma4/vision_real_test.py new file mode 100644 index 000000000..f3615ee6f --- /dev/null +++ b/tests/models/gemma4/vision_real_test.py @@ -0,0 +1,123 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Stage-1 tests for the real Gemma 4 vision tower port. + +These verify WIRING and SHAPES and that the nnx module tree maps 1:1 onto the +real `google/gemma-4-e2b-it` checkpoint key names. They do NOT verify numeric +parity with HuggingFace — that requires the real weights + torch and is a +separate stage (see docs/gemma4_vision_port.md). +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +from absl.testing import absltest +from flax import nnx +from tunix.models.gemma4 import vision_real as vr + + +def _state_keys(module) -> set[str]: + """Flatten an nnx module's params to loader-style dotted keys.""" + _, state = nnx.split(module) + keys = set() + for path, _ in jax.tree_util.tree_flatten_with_path(state)[0]: + parts = [] + for k in path: + name = getattr(k, "name", None) + if name is None: + idx = getattr(k, "idx", None) + name = idx if idx is not None else getattr(k, "key", k) + if name == "value": # nnx.Param leaf suffix + continue + parts.append(str(name)) + keys.add(".".join(parts)) + return keys + + +def _square_positions(side: int) -> jnp.ndarray: + xs, ys = jnp.meshgrid(jnp.arange(side), jnp.arange(side), indexing="xy") + return jnp.stack([xs.reshape(-1), ys.reshape(-1)], axis=-1)[None] + + +class VisionRealTest(absltest.TestCase): + + def test_forward_shapes_and_pooling(self): + cfg = vr.Gemma4VisionConfig(num_hidden_layers=2) + model = vr.Gemma4VisionModel(cfg, rngs=nnx.Rngs(0)) + + side = 12 # 144 patches -> pool by 3x3 -> 16 soft tokens. + p = side * side + pos = _square_positions(side) + px = jax.random.uniform( + jax.random.PRNGKey(1), (1, p, 3 * cfg.patch_size**2) + ) + + soft, mask = model(px, pos) + self.assertEqual(soft.shape, (1, 16, cfg.hidden_size)) + self.assertEqual(mask.shape, (1, 16)) + self.assertTrue(bool(jnp.all(mask))) + self.assertTrue(bool(jnp.all(jnp.isfinite(soft)))) + + def test_projector_maps_to_text_dim(self): + proj = vr.Gemma4MultimodalEmbedder( + vision_hidden_size=768, text_hidden_size=1536, eps=1e-6, rngs=nnx.Rngs(0) + ) + out = proj(jnp.ones((1, 16, 768))) + self.assertEqual(out.shape, (1, 16, 1536)) + + def test_module_tree_matches_checkpoint_keys(self): + """Every param path must correspond to a real checkpoint key name. + + The real keys (verified from google/gemma-4-e2b-it) are, e.g.: + model.vision_tower.patch_embedder.input_proj.weight + model.vision_tower.patch_embedder.position_embedding_table + model.vision_tower.encoder.layers.N.self_attn.q_proj.linear.weight + model.vision_tower.encoder.layers.N.self_attn.q_norm.weight + model.vision_tower.encoder.layers.N.mlp.gate_proj.linear.weight + model.vision_tower.encoder.layers.N.input_layernorm.weight + A loader maps `...weight`->`...kernel` (linears, transposed) and + `...weight`->`...scale` (norms). This test pins the nnx side of that map. + """ + cfg = vr.Gemma4VisionConfig(num_hidden_layers=1) + keys = _state_keys(vr.Gemma4VisionModel(cfg, rngs=nnx.Rngs(0))) + for expected in [ + "patch_embedder.input_proj.kernel", + "patch_embedder.position_embedding_table", + "encoder.layers.0.self_attn.q_proj.linear.kernel", + "encoder.layers.0.self_attn.k_proj.linear.kernel", + "encoder.layers.0.self_attn.v_proj.linear.kernel", + "encoder.layers.0.self_attn.o_proj.linear.kernel", + "encoder.layers.0.self_attn.q_norm.scale", + "encoder.layers.0.self_attn.k_norm.scale", + "encoder.layers.0.mlp.gate_proj.linear.kernel", + "encoder.layers.0.mlp.up_proj.linear.kernel", + "encoder.layers.0.mlp.down_proj.linear.kernel", + "encoder.layers.0.input_layernorm.scale", + "encoder.layers.0.post_attention_layernorm.scale", + "encoder.layers.0.pre_feedforward_layernorm.scale", + "encoder.layers.0.post_feedforward_layernorm.scale", + ]: + self.assertIn(expected, keys, f"missing param path: {expected}") + + def test_v_norm_has_no_scale(self): + # HF v_norm uses with_scale=False -> no learnable param. + cfg = vr.Gemma4VisionConfig(num_hidden_layers=1) + keys = _state_keys(vr.Gemma4VisionModel(cfg, rngs=nnx.Rngs(0))) + self.assertNotIn("encoder.layers.0.self_attn.v_norm.scale", keys) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/rl/common_test.py b/tests/rl/common_test.py index b089e6995..03ab51427 100644 --- a/tests/rl/common_test.py +++ b/tests/rl/common_test.py @@ -292,6 +292,101 @@ def test_compute_per_token_logps( logits.shape, (expected_logps.shape[0], expected_logps.shape[1], 256) ) + def test_model_accepts_segment_ids_helper(self): + # ToyTransformer.__call__ accepts segment_ids explicitly. + model_with = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0)) + self.assertTrue(common._model_accepts_segment_ids(model_with)) + + class _NoSegIdsModel(nnx.Module): + + def __call__(self, tokens, positions=None, cache=None, + attention_mask=None): + del tokens, positions, cache, attention_mask + return None, None + + self.assertFalse(common._model_accepts_segment_ids(_NoSegIdsModel())) + + class _VarKwargsModel(nnx.Module): + + def __call__(self, tokens, **kwargs): + del tokens, kwargs + return None, None + + # **kwargs catch-all should be treated as accepting segment_ids. + self.assertTrue(common._model_accepts_segment_ids(_VarKwargsModel())) + + def test_compute_per_token_logps_model_without_segment_ids(self): + # Reproduces google/tunix#1539: when the model's __call__ does NOT accept + # segment_ids, compute_per_token_logps must not pass it through (it would + # raise TypeError on Gemma3 etc.). + class _NoSegIdsTransformer(nnx.Module): + + def __init__(self, vocab_size: int, rngs: nnx.Rngs): + self.vocab_size = vocab_size + self.emb = nnx.Embed(vocab_size, 8, rngs=rngs) + self.head = nnx.Linear(8, vocab_size, rngs=rngs) + + def __call__( + self, + tokens, + positions=None, + cache=None, + attention_mask=None, + ): + del positions, cache, attention_mask + x = self.emb(tokens) + return self.head(x), None + + model = _NoSegIdsTransformer(vocab_size=8, rngs=nnx.Rngs(0)) + graphdef, state = nnx.split(model) + + prompt_tokens = jnp.array([[1, 2, 3, 4], [0, 0, 1, 2]], dtype=jnp.int32) + completion_tokens = jnp.array( + [[5, 6, 7, 0], [3, 4, 5, 6]], dtype=jnp.int32 + ) + + # Should not raise even though input_seg_ids would be derived from + # process_ids — the gating helper must suppress passing segment_ids. + per_token_logps = common.compute_per_token_logps( + graphdef, + state, + prompt_tokens, + completion_tokens, + pad_id=0, + eos_id=-1, + return_logits=False, + ) + self.assertEqual(per_token_logps.shape, (2, 4)) + + def test_compute_score_model_without_segment_ids(self): + # Same regression but for compute_score, which also unconditionally passed + # the derived segment_ids before this fix. + class _NoSegIdsScorer(nnx.Module): + + def __init__(self, rngs: nnx.Rngs): + self.emb = nnx.Embed(8, 8, rngs=rngs) + self.head = nnx.Linear(8, 1, rngs=rngs) + + def __call__(self, tokens, positions=None, cache=None, + attention_mask=None): + del positions, cache, attention_mask + return self.head(self.emb(tokens)) + + model = _NoSegIdsScorer(rngs=nnx.Rngs(0)) + + prompt_tokens = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + completion_tokens = jnp.array([[5, 6, 7, 0]], dtype=jnp.int32) + + scores = common.compute_score( + model, + prompt_tokens, + completion_tokens, + pad_id=0, + eos_id=-1, + ) + # [B, T] after the squeeze inside compute_score. + self.assertEqual(scores.shape, (1, 8)) + def test_np_make_completion_mask(self): completion_ids = np.array( [ diff --git a/tunix/models/gemma4/image_processing.py b/tunix/models/gemma4/image_processing.py new file mode 100644 index 000000000..c04e531b7 --- /dev/null +++ b/tunix/models/gemma4/image_processing.py @@ -0,0 +1,176 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""NumPy port of HF `Gemma4ImageProcessor` (patchify + 2D position ids). + +Produces the two arrays the real vision tower consumes: + * ``pixel_values`` [B, max_patches, 3*patch^2] (flattened patches) + * ``pixel_position_ids`` [B, max_patches, 2] ((x, y); padding = -1) + +PARITY CAVEAT: the deterministic parts here (target-size computation, +patchification order, position-id grid, padding) are faithful to HF. The +*resize* step in HF uses torchvision's antialiased bilinear, which is hard to +reproduce bit-exactly in PIL/NumPy. For numeric parity validation (Stage 3), +feed BOTH stacks the SAME ``pixel_values``/``pixel_position_ids`` (e.g. produced +by HF's processor) so vision-tower parity is isolated from resize differences. +""" + +from __future__ import annotations + +import dataclasses +import math + +import numpy as np + +# HF supports a fixed set of soft-token budgets. +_SUPPORTED_SOFT_TOKENS = (32, 64, 256, 280) + + +def get_aspect_ratio_preserving_size( + height: int, + width: int, + patch_size: int, + max_patches: int, + pooling_kernel_size: int, +) -> tuple[int, int]: + """Largest (h, w) that (1) yields <= max_patches patches and (2) is divisible + by ``pooling_kernel_size * patch_size``. Faithful to HF.""" + total_px = height * width + target_px = max_patches * (patch_size**2) + factor = math.sqrt(target_px / total_px) + ideal_height = factor * height + ideal_width = factor * width + side_mult = pooling_kernel_size * patch_size + + target_height = int(math.floor(ideal_height / side_mult)) * side_mult + target_width = int(math.floor(ideal_width / side_mult)) * side_mult + + if target_height == 0 and target_width == 0: + raise ValueError( + "Resizing to 0x0; height/width must be divisible by " + f"pooling_kernel_size*patch_size={side_mult}." + ) + max_side_length = (max_patches // pooling_kernel_size**2) * side_mult + if target_height == 0: + target_height = side_mult + target_width = min(int(math.floor(width / height)) * side_mult, max_side_length) + elif target_width == 0: + target_width = side_mult + target_height = min(int(math.floor(height / width)) * side_mult, max_side_length) + return target_height, target_width + + +def convert_image_to_patches(image_chw: np.ndarray, patch_size: int) -> np.ndarray: + """(C, H, W) -> (num_patches, patch_size*patch_size*C), channel-last in patch. + + Mirrors HF/SigLIP2 `convert_image_to_patches`: + reshape (C, nH, pH, nW, pW) -> permute (nH, nW, pH, pW, C) -> (nH*nW, -1). + """ + c, h, w = image_chw.shape + nh, nw = h // patch_size, w // patch_size + x = image_chw.reshape(c, nh, patch_size, nw, patch_size) + x = np.transpose(x, (1, 3, 2, 4, 0)) # (nH, nW, pH, pW, C) + return x.reshape(nh * nw, -1) + + +def build_position_ids(patch_height: int, patch_width: int) -> np.ndarray: + """(x, y) position per patch, row-major over (nH, nW) to match patchify order.""" + xs, ys = np.meshgrid( + np.arange(patch_width), np.arange(patch_height), indexing="xy" + ) + return np.stack([xs, ys], axis=-1).reshape(patch_height * patch_width, 2) + + +def pad_along_first_dim( + patches: np.ndarray, positions: np.ndarray, target_length: int +) -> tuple[np.ndarray, np.ndarray]: + """Pad patches with 0 and positions with -1 up to ``target_length``.""" + cur = patches.shape[0] + if target_length > cur: + pad = target_length - cur + patches = np.pad(patches, ((0, pad), (0, 0)), constant_values=0) + positions = np.pad(positions, ((0, pad), (0, 0)), constant_values=-1) + return patches, positions + + +@dataclasses.dataclass +class Gemma4ImageProcessor: + """Minimal inference image processor for the real Gemma 4 vision tower.""" + + patch_size: int = 16 + pooling_kernel_size: int = 3 + max_soft_tokens: int = 280 + rescale_factor: float = 1.0 / 255.0 + + def __post_init__(self): + if self.max_soft_tokens not in _SUPPORTED_SOFT_TOKENS: + raise ValueError( + f"max_soft_tokens must be one of {_SUPPORTED_SOFT_TOKENS}, " + f"got {self.max_soft_tokens}." + ) + + @property + def max_patches(self) -> int: + return self.max_soft_tokens * self.pooling_kernel_size**2 + + def _resize(self, image_chw: np.ndarray) -> np.ndarray: + """Aspect-ratio-preserving resize via PIL bilinear+antialias. + + NOTE: not bit-exact with torchvision (see module docstring). Skipped if the + image already has target dimensions. + """ + c, h, w = image_chw.shape + th, tw = get_aspect_ratio_preserving_size( + h, w, self.patch_size, self.max_patches, self.pooling_kernel_size + ) + if (th, tw) == (h, w): + return image_chw + try: + from PIL import Image # pylint: disable=g-import-not-at-top + except ImportError as e: + raise ImportError( + "Pillow is required to resize images in Gemma4ImageProcessor; " + "install pillow or pass pre-resized CHW arrays." + ) from e + hwc = np.transpose(image_chw, (1, 2, 0)) + pil = Image.fromarray(hwc.astype(np.uint8)) if hwc.dtype != np.uint8 else \ + Image.fromarray(hwc) + pil = pil.resize((tw, th), resample=Image.BILINEAR) + return np.transpose(np.asarray(pil), (2, 0, 1)) + + def __call__( + self, images, do_resize: bool = True + ) -> tuple[np.ndarray, np.ndarray, list[int]]: + """images: list of (C, H, W) uint8/float arrays. Returns + (pixel_values [B, max_patches, 3*p^2], pixel_position_ids [B, max_patches, 2], + num_soft_tokens_per_image).""" + pixel_values, position_ids, num_soft = [], [], [] + for image in images: + image = np.asarray(image) + if do_resize: + image = self._resize(image) + image = image.astype(np.float32) * self.rescale_factor # -> [0, 1] + patch_h = image.shape[-2] // self.patch_size + patch_w = image.shape[-1] // self.patch_size + patches = convert_image_to_patches(image, self.patch_size) + num_soft.append(patches.shape[0] // self.pooling_kernel_size**2) + positions = build_position_ids(patch_h, patch_w) + patches, positions = pad_along_first_dim(patches, positions, self.max_patches) + pixel_values.append(patches) + position_ids.append(positions) + return ( + np.stack(pixel_values, axis=0), + np.stack(position_ids, axis=0), + num_soft, + ) diff --git a/tunix/models/gemma4/model.py b/tunix/models/gemma4/model.py index c64a32522..0615b412c 100644 --- a/tunix/models/gemma4/model.py +++ b/tunix/models/gemma4/model.py @@ -18,7 +18,8 @@ import enum from functools import partial import itertools -from typing import Tuple +from typing import Tuple, Union +import einops import flax from flax import nnx import jax @@ -31,6 +32,9 @@ from jax.sharding import PartitionSpec as P import jaxtyping from tunix.generate.mappings import BackendMappingMixin +from tunix.models.gemma3 import merge_embeddings as merge_embeddings_lib +from tunix.models.gemma3 import utils as mm_utils +from tunix.models.gemma3 import vision from tunix.models.gemma4 import moe from tunix.utils import compat from tunix.utils import env_utils @@ -75,6 +79,8 @@ class ShardingConfig: per_layer_input_gate: Tuple[str | None, ...] per_layer_projection: Tuple[str | None, ...] per_layer_input_embedding: Tuple[str | None, ...] + # SigLIP vision encoder sharding. + siglip: vision.SigLIPShardingConfig | None = None @staticmethod def get_default_sharding(is_sampling: bool = False): @@ -100,6 +106,7 @@ def get_default_sharding(is_sampling: bool = False): per_layer_input_gate=(fsdp, 'tp'), per_layer_projection=('tp', fsdp), per_layer_input_embedding=('tp', None, fsdp), + siglip=vision.SigLIPShardingConfig.get_default_sharding(is_sampling), ) @@ -131,6 +138,10 @@ class ModelConfig: local_scale_factor: float = 1.0 global_scale_factor: float = 1.0 + # Vision config. If set, the model includes a SigLIP vision encoder and + # accepts `images` in the forward pass for multi-modal inputs. + vision_config: vision.SigLIPConfig | None = None + shd_config: ShardingConfig = ShardingConfig.get_default_sharding() remat_config: RematConfig = RematConfig.NONE param_dtype: jnp.dtype = jnp.float32 @@ -157,6 +168,7 @@ def __post_init__(self): def gemma4_e2b( cls, sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(), + text_only: bool = True, ) -> 'ModelConfig': return cls( num_layers=35, @@ -171,6 +183,7 @@ def gemma4_e2b( per_layer_input_dim=256, frac_shared_layers=20.0 / 35, override_kv_shared_ffw_hidden=int(1536 * 4 * 2), + vision_config=None if text_only else vision.SigLIPConfig(), attention_pattern=( AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, @@ -184,6 +197,7 @@ def gemma4_e2b( def gemma4_e4b( cls, sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(), + text_only: bool = True, ) -> 'ModelConfig': return cls( num_layers=42, @@ -197,6 +211,7 @@ def gemma4_e4b( shd_config=sharding_config, per_layer_input_dim=256, frac_shared_layers=18.0 / 42, + vision_config=None if text_only else vision.SigLIPConfig(), attention_pattern=( AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, @@ -211,6 +226,7 @@ def gemma4_e4b( def gemma4_31b( cls, sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(), + text_only: bool = True, ) -> 'ModelConfig': return cls( num_layers=60, @@ -224,6 +240,7 @@ def gemma4_31b( sliding_window_size=1024, shd_config=sharding_config, k_eq_v_global=True, + vision_config=None if text_only else vision.SigLIPConfig(), attention_pattern=( AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, @@ -238,6 +255,7 @@ def gemma4_31b( def gemma4_26b_a4b( cls, sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(), + text_only: bool = True, ) -> 'ModelConfig': return cls( num_layers=30, @@ -257,6 +275,7 @@ def gemma4_26b_a4b( moe_dense_hidden_dim=2112, k_eq_v_global=True, global_rope_proportion=0.25, + vision_config=None if text_only else vision.SigLIPConfig(), attention_pattern=( AttentionType.LOCAL_SLIDING, AttentionType.LOCAL_SLIDING, @@ -275,6 +294,7 @@ def __init__( self, config: ModelConfig, rngs: nnx.Rngs, + vision_proj_dim: int | None = None, ): self.config = config self.vocab_size = config.num_embed @@ -314,6 +334,23 @@ def __init__( sharding=config.shd_config.per_layer_input_embedding, ) + if vision_proj_dim: + self.mm_soft_embedding_norm = RMSNorm( + vision_proj_dim, + rngs=rngs, + sharding=config.shd_config.vision_soft_emb_norm_weight, + dtype=self.config.dtype, + param_dtype=self.param_dtype, + ) + self.mm_input_projection = Einsum( + einsum_str='...TM,MD->...TD', + shape=(vision_proj_dim, self.embed_dim), + rngs=rngs, + sharding=config.shd_config.vision_proj, + dtype=self.config.dtype, + param_dtype=self.param_dtype, + ) + def encode(self, x: jaxtyping.ArrayLike) -> jaxtyping.Array: x = self.input_embedding[(x,)] x *= jnp.sqrt(x.shape[-1]).astype(x.dtype) @@ -333,6 +370,12 @@ def encode_per_layer_input( y *= jnp.sqrt(self.config.per_layer_input_dim).astype(y.dtype) return (x + y) * jax.lax.rsqrt(2.0).astype(x.dtype) + def encode_vision(self, x: jaxtyping.ArrayLike) -> jaxtyping.Array: + """Projects vision encoder outputs into the language model embedding space.""" + x = self.mm_soft_embedding_norm(x) + x = self.mm_input_projection(x) + return x + def decode(self, x: jaxtyping.ArrayLike) -> jaxtyping.Array: x = jnp.astype(x, self.config.dtype) w = jnp.astype(self.input_embedding.value, self.config.dtype) @@ -428,13 +471,20 @@ def __init__( dim: int, *, rngs: nnx.Rngs, - sharding: ShardingConfig = ShardingConfig.get_default_sharding(), + sharding: Union[ShardingConfig, Tuple[str | None, ...]] = ( + ShardingConfig.get_default_sharding() + ), dtype: jnp.dtype, param_dtype: jnp.dtype, ): + scale_sharding = ( + sharding.rms_norm_weight + if isinstance(sharding, ShardingConfig) + else sharding + ) self.scale = nnx.Param( nnx.initializers.ones_init()(rngs.params(), dim).astype(param_dtype), - sharding=sharding.rms_norm_weight, + sharding=scale_sharding, ) self.dtype = dtype @@ -1186,7 +1236,25 @@ class Gemma4(BackendMappingMixin, nnx.Module): def __init__(self, config: ModelConfig, *, rngs: nnx.Rngs): self.config = config - self.embedder = Embedder(config, rngs=rngs) + + if config.vision_config is not None: + self.vision_encoder = vision.SigLiP( + config=config.vision_config, + shd_config=config.shd_config.siglip, + rngs=rngs, + ) + else: + self.vision_encoder = None + + self.embedder = Embedder( + config, + rngs=rngs, + vision_proj_dim=( + self.vision_encoder.siglip_encoder.width + if self.vision_encoder is not None + else None + ), + ) pattern = ( config.attention_pattern @@ -1241,16 +1309,28 @@ def __call__( cache=None, attention_mask=None, decode_only_last_token=False, + *, + images: jaxtyping.Array | None = None, + input_embeddings: jaxtyping.Array | None = None, + per_layer_inputs: jaxtyping.Array | None = None, ): if positions is None: B, T = tokens.shape # pylint: disable=invalid-name positions = jnp.tile(jnp.arange(T)[None, :], (B, 1)) new_cache = {} - x = self.embedder.encode(tokens) + # `input_embeddings` lets a multimodal wrapper inject already-merged + # (text + vision soft-token) embeddings; otherwise encode tokens here. + if input_embeddings is not None: + x = input_embeddings + else: + x = self._encode_and_get_inputs(tokens=tokens, images=images) - per_layer_inputs = None - if self.config.per_layer_input_dim > 0: + # `per_layer_inputs` lets a multimodal wrapper precompute PLE from + # pad-substituted tokens/embeddings (matching HF Gemma4Model.forward), + # since the merged embeddings include vision soft tokens at image-token + # positions and shouldn't feed the per-layer projection there. + if per_layer_inputs is None and self.config.per_layer_input_dim > 0: per_layer_inputs = self.embedder.encode_per_layer_input(x, tokens) # Stores the raw KV projections for the current forward pass. Used for @@ -1315,3 +1395,111 @@ def init_cache(self, batch_size, max_seq_len, dtype): continue # Skip shared layers. cache[f'layer_{i}'] = layer.init_cache(batch_size, max_seq_len, dtype) return cache + + def _encode_and_get_inputs( + self, + *, + tokens: jaxtyping.Array, # (B, L) + images: jaxtyping.Array | None = None, # (B, H, W, C) or (B, N, H, W, C) + ) -> jaxtyping.Array: + """Encode the text tokens, eventually including the vision embeddings.""" + if self.config.vision_config is not None and images is not None: + self._assert_support_mm() + if len(images.shape) == 4: # If num_images is 1, add an axis. + images = einops.rearrange(images, 'b h w c -> b 1 h w c') + + x = self.embedder.encode(tokens) + + if images is not None: + x = self._merge_mm_embeddings(tokens=tokens, embeddings=x, images=images) + return x + + def _assert_support_mm(self) -> None: + if self.vision_encoder is None: + raise ValueError( + f'The model {type(self).__name__!r} does not have a vision encoder,' + ' yet images are provided. Construct the model with a `vision_config`' + ' (e.g. `ModelConfig.gemma4_e4b(text_only=False)`) to enable' + ' multi-modal inputs.' + ) + + def _merge_mm_embeddings( + self, + *, + tokens: jaxtyping.ArrayLike, # (B, L) + embeddings: jaxtyping.ArrayLike, # (B, L, D) + images: jaxtyping.ArrayLike, # (B, N, H, W, C) + ) -> jaxtyping.ArrayLike: + """Update the text embeddings to include the vision soft tokens.""" + soft_embeddings = self._encode_vision(images) + if self.config.vision_config is None: + raise ValueError( + '`vision_config` is required for `_merge_mm_embeddings`.' + ) + return merge_embeddings_lib.merge_embeddings( + text_embeddings=embeddings, + vision_embeddings=soft_embeddings, + mask=tokens == self.config.vision_config.soft_token_placeholder, + ) + + def _encode_vision( + self, images: jaxtyping.ArrayLike # (B, N, H, W, C) + ) -> jaxtyping.ArrayLike: # (B, N, P, D) + """Encode the images into the same space as the text embeddings.""" + if self.vision_encoder is None: + raise ValueError('`vision_encoder` is None, cannot encode images.') + soft_embeddings = self.vision_encoder(images=images) + soft_embeddings = self.embedder.encode_vision(soft_embeddings) + return soft_embeddings + + def get_attention_mask( + self, + tokens: jaxtyping.ArrayLike, # (B, L) + *, + inputs_mask: jaxtyping.ArrayLike | None = None, # (B, L) + ): + """Returns the attention mask for the transformer. + + For multi-modal inputs, the mask is bidirectional over image soft-token + spans (so all soft tokens of an image attend to each other) while remaining + causal over text tokens, matching the Gemma 3 behavior. + """ + token_placeholder_id = ( + None + if self.config.vision_config is None + else self.config.vision_config.soft_token_placeholder + ) + return mm_utils.get_attention_mask( + tokens, + inputs_mask=inputs_mask, + token_placeholder_id=token_placeholder_id, + ) + + def get_model_input(self): + """Returns a dummy model input for the transformer. + + Used to trace the graph for LoRA application and similar transforms. + """ + dummy_batch_size = 2 + dummy_seq_len = 1 + inputs = { + 'tokens': jnp.ones( + (dummy_batch_size, dummy_seq_len), dtype=jnp.int32 + ), + 'positions': jnp.ones( + (dummy_batch_size, dummy_seq_len), dtype=jnp.int32 + ), + 'cache': None, + 'attention_mask': jnp.ones( + (dummy_batch_size, 1, dummy_seq_len), dtype=jnp.bool + ), + } + + if self.vision_encoder is not None: + vc = self.config.vision_config + inputs['images'] = jnp.ones( + (dummy_batch_size, 1, vc.image_height, vc.image_width, + vc.image_channels), + dtype=jnp.float32, + ) + return inputs diff --git a/tunix/models/gemma4/multimodal.py b/tunix/models/gemma4/multimodal.py new file mode 100644 index 000000000..696e3dff5 --- /dev/null +++ b/tunix/models/gemma4/multimodal.py @@ -0,0 +1,181 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Stage 4: wire the real gemma4_vision stack into the Gemma 4 text model. + +`Gemma4Multimodal` composes the existing text `Gemma4` with the ported +`Gemma4VisionStack` (vision tower + embed_vision projector). The forward: + + 1. embed text tokens (scaled, as usual); + 2. run the vision stack to get soft tokens already projected into text space; + 3. scatter the soft tokens into the positions where ``tokens == image_token_id`` + (the HF `masked_scatter` equivalent), via ``merge_embeddings``; + 4. run the text transformer on the merged embeddings. + +This deliberately wraps `Gemma4` (passing merged embeddings through the new +``input_embeddings`` arg) rather than editing its layer loop, and leaves the +legacy SigLIP scaffolding untouched. + +Scope note: the merge currently assumes the number of valid soft tokens equals +the number of ``image_token_id`` placeholders in the sequence (true for a +single, non-padded image, which is what the parity tests cover). Variable-size +images / multiple images per sequence need a gather of valid pooled tokens +first; that is a follow-up. +""" + +from __future__ import annotations + +import jax +from flax import nnx +import jax.numpy as jnp +import jaxtyping +from tunix.models.gemma3 import merge_embeddings as merge_embeddings_lib +from tunix.models.gemma3 import utils as mm_utils +from tunix.models.gemma4 import model as model_lib +from tunix.models.gemma4 import params_safetensors as text_params +from tunix.models.gemma4 import vision_params_safetensors as vision_params +from tunix.models.gemma4 import vision_real + + +class Gemma4Multimodal(nnx.Module): + """Text `Gemma4` + real `Gemma4VisionStack`, merged at image-token positions.""" + + def __init__( + self, + text_model: model_lib.Gemma4, + vision_stack: vision_real.Gemma4VisionStack, + *, + image_token_id: int, + pad_token_id: int = 0, + bidirectional_image_span: bool = False, + ): + self.text_model = text_model + self.vision = vision_stack + self.image_token_id = image_token_id + # Per-layer-input embeddings must be computed from pad-substituted + # tokens/embeddings at image positions (HF Gemma4Model.forward behavior); + # otherwise PLE leaks vision soft tokens and image_token_id lookups into + # positions that HF treats as "no language token present". + self.pad_token_id = pad_token_id + # Larger Gemma 4 checkpoints (config `use_bidirectional_attention == "vision"`) + # use the Gemma-3-style bidirectional mask over each image's soft-token span; + # smaller ones use a plain causal mask. Default matches HF small-model + # behaviour; flip True for checkpoints that set `use_bidirectional_attention`. + self.bidirectional_image_span = bidirectional_image_span + + def encode_multimodal_inputs( + self, + tokens: jaxtyping.Array, # (B, L) + pixel_values: jaxtyping.Array, # (B, P, 3*patch^2) + pixel_position_ids: jaxtyping.Array, # (B, P, 2) + ) -> jaxtyping.Array: # (B, L, D) + """Text embeddings with vision soft tokens scattered into image slots.""" + text_embeddings = self.text_model.embedder.encode(tokens) # (B, L, D) + soft_tokens, _mask = self.vision(pixel_values, pixel_position_ids) # (B, S, D) + # merge_embeddings expects vision as (B, num_images, tokens_per_image, D). + return merge_embeddings_lib.merge_embeddings( + text_embeddings=text_embeddings, + vision_embeddings=soft_tokens[:, None, :, :], + mask=(tokens == self.image_token_id), + ) + + def _compute_per_layer_inputs( + self, tokens: jaxtyping.Array, merged_embeddings: jaxtyping.Array + ) -> jaxtyping.Array: + """HF-style PLE: ONLY the token-identity branch is pad-substituted. + + HF `Gemma4Model.forward` calls `get_per_layer_inputs(llm_input_ids, …)` + which ignores `inputs_embeds` and just looks up `embed_tokens_per_layer` + on the pad-substituted ids. The context-projection branch (inside + `project_per_layer_inputs`) then runs on the **merged** `inputs_embeds` + (vision features at image positions). So we pass merged embeddings to the + `x` argument and pad-substituted ids to `t`. + """ + llm_tokens = jnp.where( + tokens == self.image_token_id, self.pad_token_id, tokens + ) + return self.text_model.embedder.encode_per_layer_input( + merged_embeddings, llm_tokens + ) + + def get_attention_mask( + self, + tokens: jaxtyping.Array, # (B, L) + *, + inputs_mask: jaxtyping.Array | None = None, + ): + """Causal by default; bidirectional over image spans iff requested.""" + placeholder = ( + self.image_token_id if self.bidirectional_image_span else None + ) + return mm_utils.get_attention_mask( + tokens, inputs_mask=inputs_mask, token_placeholder_id=placeholder, + ) + + def __call__( + self, + tokens: jaxtyping.Array, # (B, L) + pixel_values: jaxtyping.Array, # (B, P, 3*patch^2) + pixel_position_ids: jaxtyping.Array, # (B, P, 2) + positions=None, + cache=None, + attention_mask=None, + decode_only_last_token: bool = False, + ): + x = self.encode_multimodal_inputs(tokens, pixel_values, pixel_position_ids) + if attention_mask is None: + attention_mask = self.get_attention_mask(tokens) + ple = None + if self.text_model.config.per_layer_input_dim > 0: + ple = self._compute_per_layer_inputs(tokens, x) + return self.text_model( + tokens, + positions=positions, + cache=cache, + attention_mask=attention_mask, + decode_only_last_token=decode_only_last_token, + input_embeddings=x, + per_layer_inputs=ple, + ) + + +def create_multimodal_from_safe_tensors( + file_dir: str, + text_config: model_lib.ModelConfig, + vision_config: vision_real.Gemma4VisionConfig, + *, + image_token_id: int, + pad_token_id: int = 0, + bidirectional_image_span: bool = False, + mesh: jax.sharding.Mesh | None = None, + dtype=None, +) -> Gemma4Multimodal: + """Loads a real `google/gemma-4-*-it` checkpoint into a `Gemma4Multimodal`. + + The text and vision loaders each read the same directory and skip the other's + keys (text skips `vision_tower`/`audio_tower`/`embed_*`; vision skips + `language_model`/`audio_tower`). The vision projector output dim must equal the + text model's `embed_dim`. + """ + text_model = text_params.create_model_from_safe_tensors( + file_dir, text_config, mesh, dtype=dtype + ) + vision_stack = vision_params.create_vision_stack_from_safe_tensors( + file_dir, vision_config, text_config.embed_dim, mesh=mesh, dtype=dtype + ) + return Gemma4Multimodal( + text_model, vision_stack, + image_token_id=image_token_id, pad_token_id=pad_token_id, + bidirectional_image_span=bidirectional_image_span, + ) diff --git a/tunix/models/gemma4/params_safetensors.py b/tunix/models/gemma4/params_safetensors.py index 9eae3f9a8..e1b5bc254 100644 --- a/tunix/models/gemma4/params_safetensors.py +++ b/tunix/models/gemma4/params_safetensors.py @@ -246,6 +246,104 @@ def _get_key_and_transform_mapping(cfg: model_lib.ModelConfig): ), } + # Vision Tower (SigLIP) and multi-modal projector. + if cfg.vision_config is not None: + mapping.update({ + r"vision_tower\.vision_model\.embeddings\.patch_embedding\.weight": ( + "vision_encoder.siglip_encoder.embedding.kernel", + ((2, 3, 1, 0), None), + ), + r"vision_tower\.vision_model\.embeddings\.patch_embedding\.bias": ( + "vision_encoder.siglip_encoder.embedding.bias", + None, + ), + r"vision_tower\.vision_model\.embeddings\.position_embedding\.weight": ( + "vision_encoder.siglip_encoder.pos_embedding", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.q_proj\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.query_proj.kernel", + ((1, 0), None), + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.k_proj\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.key_proj.kernel", + ((1, 0), None), + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.v_proj\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.value_proj.kernel", + ((1, 0), None), + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.out_proj\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.out_proj.kernel", + ((1, 0), None), + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.q_proj\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.query_proj.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.k_proj\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.key_proj.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.v_proj\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.value_proj.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.self_attn\.out_proj\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.attn.out_proj.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm1\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.ln1.scale", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm1\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.ln1.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm2\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.ln2.scale", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.layer_norm2\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.ln2.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc1\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.mlp.fc1.kernel", + ((1, 0), None), + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc2\.weight": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.mlp.fc2.kernel", + ((1, 0), None), + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc1\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.mlp.fc1.bias", + None, + ), + r"vision_tower\.vision_model\.encoder\.layers\.([0-9]+)\.mlp\.fc2\.bias": ( + r"vision_encoder.siglip_encoder.transformer.blocks.\1.mlp.fc2.bias", + None, + ), + r"vision_tower\.vision_model\.post_layernorm\.weight": ( + "vision_encoder.siglip_encoder.transformer.encoder_norm.scale", + None, + ), + r"vision_tower\.vision_model\.post_layernorm\.bias": ( + "vision_encoder.siglip_encoder.transformer.encoder_norm.bias", + None, + ), + # Multi-modal Projector + r"multi_modal_projector\.mm_input_projection_weight": ( + "embedder.mm_input_projection.w", + None, + ), + r"multi_modal_projector\.mm_soft_emb_norm\.weight": ( + "embedder.mm_soft_embedding_norm.scale", + None, + ), + }) + return mapping diff --git a/tunix/models/gemma4/vision_params_safetensors.py b/tunix/models/gemma4/vision_params_safetensors.py new file mode 100644 index 000000000..dd2dc3db0 --- /dev/null +++ b/tunix/models/gemma4/vision_params_safetensors.py @@ -0,0 +1,110 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Safetensors loader for the real Gemma 4 vision stack. + +Maps `google/gemma-4-*-it` checkpoint keys (`model.vision_tower.*`, +`model.embed_vision.*`) onto a `vision_real.Gemma4VisionStack`. Audio keys +(`model.audio_tower.*`, `model.embed_audio.*`) and the e2b clip buffers +(`...{proj}.input_min/max`, `...{proj}.output_min/max`) are intentionally +unmapped, so the generic loader skips them (logged, not loaded). + +IMPORTANT: the real keys are prefixed `model.`, and the loader matches with +`re.match` (anchored at the start). Every pattern therefore starts with an +optional `(?:model\\.)?` and is `$`-anchored so it matches exactly one tensor. +""" + +from __future__ import annotations + +import jax +from jax import numpy as jnp +from tunix.models import safetensors_loader +from tunix.models.gemma4 import vision_real + + +# Per-layer linears that are bias-free and stored as `...{name}.linear.weight` +# (torch [out, in]); nnx holds `...{name}.linear.kernel` [in, out] -> transpose. +_ATTN_LINEARS = ("q_proj", "k_proj", "v_proj", "o_proj") +_MLP_LINEARS = ("gate_proj", "up_proj", "down_proj") +# Scaled RMSNorms (q_norm/k_norm + the four sandwich norms). v_norm is unscaled +# (no weight in the checkpoint) and the projector pre-norm is unscaled too. +_ATTN_NORMS = ("q_norm", "k_norm") +_LAYER_NORMS = ( + "input_layernorm", + "post_attention_layernorm", + "pre_feedforward_layernorm", + "post_feedforward_layernorm", +) + +_TRANSPOSE = ((1, 0), None) # (permute, reshape) for a 2D weight. +_IDENTITY = None + + +def vision_key_mapping(config: vision_real.Gemma4VisionConfig): + """Returns {torch_key_regex: (nnx_key_repl, transform)} for the vision stack.""" + p = r"(?:model\.)?" # optional checkpoint `model.` prefix + mapping = { + # --- patch embedder --- + p + r"vision_tower\.patch_embedder\.input_proj\.weight$": ( + "vision_tower.patch_embedder.input_proj.kernel", + _TRANSPOSE, + ), + p + r"vision_tower\.patch_embedder\.position_embedding_table$": ( + "vision_tower.patch_embedder.position_embedding_table", + _IDENTITY, + ), + # --- projector (embed_vision); pre-projection norm is unscaled (no key) --- + p + r"embed_vision\.embedding_projection\.weight$": ( + "embed_vision.embedding_projection.kernel", + _TRANSPOSE, + ), + } + # --- encoder layers --- + for name in _ATTN_LINEARS: + mapping[ + p + rf"vision_tower\.encoder\.layers\.(\d+)\.self_attn\.{name}\.linear\.weight$" + ] = (rf"vision_tower.encoder.layers.\1.self_attn.{name}.linear.kernel", _TRANSPOSE) + for name in _ATTN_NORMS: + mapping[ + p + rf"vision_tower\.encoder\.layers\.(\d+)\.self_attn\.{name}\.weight$" + ] = (rf"vision_tower.encoder.layers.\1.self_attn.{name}.scale", _IDENTITY) + for name in _MLP_LINEARS: + mapping[ + p + rf"vision_tower\.encoder\.layers\.(\d+)\.mlp\.{name}\.linear\.weight$" + ] = (rf"vision_tower.encoder.layers.\1.mlp.{name}.linear.kernel", _TRANSPOSE) + for name in _LAYER_NORMS: + mapping[ + p + rf"vision_tower\.encoder\.layers\.(\d+)\.{name}\.weight$" + ] = (rf"vision_tower.encoder.layers.\1.{name}.scale", _IDENTITY) + return mapping + + +def create_vision_stack_from_safe_tensors( + file_dir: str, + config: vision_real.Gemma4VisionConfig, + text_hidden_size: int, + mesh: jax.sharding.Mesh | None = None, + dtype: jnp.dtype | None = None, +): + """Builds a `Gemma4VisionStack` and loads vision weights from safetensors.""" + return safetensors_loader.load_and_create_model( + file_dir=file_dir, + model_class=lambda cfg, rngs: vision_real.Gemma4VisionStack( + cfg, text_hidden_size, rngs=rngs + ), + config=config, + key_mapping=vision_key_mapping, + mesh=mesh, + dtype=dtype, + ) diff --git a/tunix/models/gemma4/vision_real.py b/tunix/models/gemma4/vision_real.py new file mode 100644 index 000000000..dc9fca212 --- /dev/null +++ b/tunix/models/gemma4/vision_real.py @@ -0,0 +1,533 @@ +# Copyright 2026 Google LLC +# +# 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. + +"""Faithful JAX/nnx port of the REAL Gemma 4 vision tower (`gemma4_vision`). + +This is a from-source port of HuggingFace +`transformers.models.gemma4.modeling_gemma4` (the classes `Gemma4Vision*` and +`Gemma4MultimodalEmbedder`), NOT the SigLIP encoder borrowed from Gemma 3. + +The real `gemma4_vision` tower (per `google/gemma-4-e2b-it`'s config.json) is a +rotary, gated-MLP, sandwich-norm ViT with: + * a patch embedder that flattens 3*patch^2 pixels and adds a factored 2D + learned position embedding (one-hot @ a [2, position_embedding_size, hidden] + table); + * multidimensional (2D) RoPE applied per spatial axis; + * per-head q/k RMSNorm (scaled) and v RMSNorm (unscaled); + * 4-norm sandwich encoder layers; + * a spatial average pooler that reduces patches to soft tokens. + +WIP STATUS (read before trusting this): + * Wiring + shapes are smoke-tested on random init (see + tests/models/gemma4/vision_real_test.py). + * NUMERICS ARE NOT YET VALIDATED against the real checkpoint. Bit-exact + parity requires loading the real safetensors weights into both this module + and HF `Gemma4VisionModel` and comparing activations (needs torch + the + checkpoint). See docs/gemma4_vision_port.md, "Stage 3: parity". + * Batched padding handling in the pooler (HF does a dynamic boolean + `hidden_states[pooler_mask]` gather) is intentionally left to the caller / + a later stage; here we return a padded `[B, output_length, hidden]` tensor + plus the validity mask, which is JAX-jit friendly. + +Reference dimensions (Gemma4VisionConfig defaults / e2b): + hidden_size=768, intermediate_size=3072, num_hidden_layers=16, + num_attention_heads=12, num_key_value_heads=12, head_dim=64, + patch_size=16, position_embedding_size=10240, pooling_kernel_size=3, + rope_theta=100.0, hidden_activation="gelu_pytorch_tanh", rms_norm_eps=1e-6. +""" + +from __future__ import annotations + +import dataclasses + +from flax import nnx +import jax +from jax import numpy as jnp +import jaxtyping +from tunix.utils import compat + + +@dataclasses.dataclass(frozen=True) +class Gemma4VisionConfig: + """Mirror of HF `Gemma4VisionConfig` (only the fields the tower needs).""" + + hidden_size: int = 768 + intermediate_size: int = 3072 + num_hidden_layers: int = 16 + num_attention_heads: int = 12 + num_key_value_heads: int = 12 + head_dim: int = 64 + rms_norm_eps: float = 1e-6 + patch_size: int = 16 + position_embedding_size: int = 10 * 1024 + pooling_kernel_size: int = 3 + rope_theta: float = 100.0 + # e2b ships with use_clipped_linears=False, so the ClippableLinear input/ + # output clamp buffers are +/-inf and are no-ops. We keep the `.linear.` + # nesting in the module tree so checkpoint keys map 1:1, but skip the clamp. + use_clipped_linears: bool = False + standardize: bool = False + param_dtype: jnp.dtype = jnp.float32 + dtype: jnp.dtype = jnp.float32 + + +def _gelu_tanh(x: jaxtyping.Array) -> jaxtyping.Array: + # HF "gelu_pytorch_tanh" == jax.nn.gelu(approximate=True). + return jax.nn.gelu(x, approximate=True) + + +class RMSNorm(nnx.Module): + """HF `Gemma4RMSNorm`: x_fp32 * rsqrt(mean(x^2)+eps) * weight. + + NOTE: unlike Gemma 1/2/3 (which scale by ``1 + weight``), Gemma 4 scales by + ``weight`` directly (weight is initialised to ones). ``with_scale=False`` + drops the parameter entirely (used by v_norm and the projector pre-norm). + """ + + def __init__( + self, + dim: int, + *, + eps: float, + with_scale: bool, + rngs: nnx.Rngs, + param_dtype: jnp.dtype, + ): + self.eps = eps + self.with_scale = with_scale + if with_scale: + self.scale = nnx.Param(jnp.ones((dim,), dtype=param_dtype)) + + def __call__(self, x: jaxtyping.Array) -> jaxtyping.Array: + in_dtype = x.dtype + x = x.astype(jnp.float32) + normed = x * jax.lax.rsqrt(jnp.mean(x * x, axis=-1, keepdims=True) + self.eps) + if self.with_scale: + normed = normed * self.scale.value.astype(jnp.float32) + return normed.astype(in_dtype) + + +class ClippableLinear(nnx.Module): + """HF `Gemma4ClippableLinear`: a bias-free Linear, optionally clamped. + + The module nests an inner ``linear`` so checkpoint keys + (``...{proj}.linear.weight``) map cleanly onto ``...{proj}.linear.kernel``. + For e2b (``use_clipped_linears=False``) the clamp is a no-op. + """ + + def __init__( + self, + in_features: int, + out_features: int, + *, + rngs: nnx.Rngs, + param_dtype: jnp.dtype, + dtype: jnp.dtype, + ): + self.linear = nnx.Linear( + in_features, + out_features, + use_bias=False, + rngs=rngs, + param_dtype=param_dtype, + dtype=dtype, + ) + + def __call__(self, x: jaxtyping.Array) -> jaxtyping.Array: + # use_clipped_linears is False for e2b; clamp buffers are +/-inf no-ops and + # are intentionally not ported here. Revisit if a clipped variant ships. + return self.linear(x) + + +# --------------------------------------------------------------------------- +# Rotary embedding (multidimensional / 2D). +# --------------------------------------------------------------------------- +def _rotate_half(x: jaxtyping.Array) -> jaxtyping.Array: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return jnp.concatenate((-x2, x1), axis=-1) + + +def _apply_rotary(x, cos, sin, unsqueeze_dim): + cos = jnp.expand_dims(cos, unsqueeze_dim) + sin = jnp.expand_dims(sin, unsqueeze_dim) + return (x * cos) + (_rotate_half(x) * sin) + + +def apply_multidimensional_rope(x, cos, sin, position_ids, unsqueeze_dim=2): + """Port of HF `apply_multidimensional_rope`. + + Splits the head_dim into ``ndim`` equal chunks (ndim == positions' last dim, + i.e. 2 for images) and applies 1D rotary independently to each chunk. + + x: [B, L, N, head_dim]; cos/sin: [B, L, head_dim]; position_ids: [B, L, ndim]. + """ + ndim = position_ids.shape[-1] + num_input_channels = x.shape[-1] + per_dim = 2 * (num_input_channels // (2 * ndim)) + if per_dim <= 0: + raise ValueError( + f"num_rotated_channels_per_dim must be > 0 (got channels=" + f"{num_input_channels}, ndim={ndim})" + ) + sizes = [per_dim] * ndim + bounds = [sum(sizes[:i]) for i in range(1, ndim)] + x_parts = jnp.split(x, bounds, axis=-1) + cos_parts = jnp.split(cos, bounds, axis=-1) + sin_parts = jnp.split(sin, bounds, axis=-1) + y = [ + _apply_rotary(x_parts[k], cos_parts[k], sin_parts[k], unsqueeze_dim) + for k in range(ndim) + ] + return jnp.concatenate(y, axis=-1) + + +class VisionRotaryEmbedding(nnx.Module): + """Port of HF `Gemma4VisionRotaryEmbedding` (default rope). + + inv_freq is computed on ``spatial_dim = head_dim // 2`` and reused for both + spatial axes; cos/sin for each axis are concatenated to width head_dim. + """ + + def __init__(self, config: Gemma4VisionConfig): + # Keep only Python scalars so nothing lands in the param tree: inv_freq is a + # derived constant (absent from the checkpoint) and is recomputed per call. + self.head_dim = config.head_dim + self.rope_theta = config.rope_theta + + def _inv_freq(self) -> jaxtyping.Array: + spatial_dim = self.head_dim // 2 + # 1.0 / base ** (arange(0, spatial_dim, 2)/spatial_dim) + exps = jnp.arange(0, spatial_dim, 2, dtype=jnp.float32) / spatial_dim + return jnp.asarray(1.0 / (self.rope_theta**exps)) # [spatial_dim/2] + + def __call__(self, position_ids: jaxtyping.Array): + # position_ids: [B, L, 2] -> cos/sin: [B, L, head_dim] + all_cos, all_sin = [], [] + inv = self._inv_freq().astype(jnp.float32) # [F] + for i in range(2): + pos = position_ids[:, :, i].astype(jnp.float32) # [B, L] + # outer product -> [B, L, F] + freqs = pos[:, :, None] * inv[None, None, :] + emb = jnp.concatenate((freqs, freqs), axis=-1) # [B, L, 2F] + all_cos.append(jnp.cos(emb)) + all_sin.append(jnp.sin(emb)) + cos = jnp.concatenate(all_cos, axis=-1) # [B, L, head_dim] + sin = jnp.concatenate(all_sin, axis=-1) + return cos, sin + + +class VisionAttention(nnx.Module): + """Port of HF `Gemma4VisionAttention` (bidirectional, scaling=1.0).""" + + def __init__(self, config: Gemma4VisionConfig, *, rngs: nnx.Rngs): + self.config = config + h = config.hidden_size + hd = config.head_dim + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = hd + kw = dict(rngs=rngs, param_dtype=config.param_dtype, dtype=config.dtype) + self.q_proj = ClippableLinear(h, self.num_heads * hd, **kw) + self.k_proj = ClippableLinear(h, self.num_kv_heads * hd, **kw) + self.v_proj = ClippableLinear(h, self.num_kv_heads * hd, **kw) + self.o_proj = ClippableLinear(self.num_heads * hd, h, **kw) + nkw = dict( + eps=config.rms_norm_eps, rngs=rngs, param_dtype=config.param_dtype + ) + self.q_norm = RMSNorm(hd, with_scale=True, **nkw) + self.k_norm = RMSNorm(hd, with_scale=True, **nkw) + self.v_norm = RMSNorm(hd, with_scale=False, **nkw) + + def __call__(self, x, cos, sin, position_ids, attn_bias): + b, l, _ = x.shape + q = self.q_norm(self.q_proj(x).reshape(b, l, self.num_heads, self.head_dim)) + q = apply_multidimensional_rope(q, cos, sin, position_ids) + k = self.k_norm(self.k_proj(x).reshape(b, l, self.num_kv_heads, self.head_dim)) + k = apply_multidimensional_rope(k, cos, sin, position_ids) + v = self.v_norm(self.v_proj(x).reshape(b, l, self.num_kv_heads, self.head_dim)) + + # [B, N, L, H] + q = jnp.transpose(q, (0, 2, 1, 3)) + k = jnp.transpose(k, (0, 2, 1, 3)) + v = jnp.transpose(v, (0, 2, 1, 3)) + n_rep = self.num_heads // self.num_kv_heads + if n_rep > 1: + k = jnp.repeat(k, n_rep, axis=1) + v = jnp.repeat(v, n_rep, axis=1) + + scaling = 1.0 # HF sets self.scaling = 1.0 for the vision tower. + scores = jnp.einsum("bnqh,bnkh->bnqk", q, k).astype(jnp.float32) * scaling + if attn_bias is not None: + scores = scores + attn_bias # [B, 1, 1, L] additive mask + weights = jax.nn.softmax(scores, axis=-1).astype(v.dtype) + out = jnp.einsum("bnqk,bnkh->bnqh", weights, v) + out = jnp.transpose(out, (0, 2, 1, 3)).reshape(b, l, -1) + return self.o_proj(out) + + +class VisionMLP(nnx.Module): + """Port of HF `Gemma4VisionMLP` (gated, gelu_tanh).""" + + def __init__(self, config: Gemma4VisionConfig, *, rngs: nnx.Rngs): + kw = dict(rngs=rngs, param_dtype=config.param_dtype, dtype=config.dtype) + self.gate_proj = ClippableLinear( + config.hidden_size, config.intermediate_size, **kw + ) + self.up_proj = ClippableLinear( + config.hidden_size, config.intermediate_size, **kw + ) + self.down_proj = ClippableLinear( + config.intermediate_size, config.hidden_size, **kw + ) + + def __call__(self, x): + return self.down_proj(_gelu_tanh(self.gate_proj(x)) * self.up_proj(x)) + + +class VisionEncoderLayer(nnx.Module): + """Port of HF `Gemma4VisionEncoderLayer` (4-norm sandwich).""" + + def __init__(self, config: Gemma4VisionConfig, *, rngs: nnx.Rngs): + nkw = dict( + eps=config.rms_norm_eps, + with_scale=True, + rngs=rngs, + param_dtype=config.param_dtype, + ) + self.input_layernorm = RMSNorm(config.hidden_size, **nkw) + self.post_attention_layernorm = RMSNorm(config.hidden_size, **nkw) + self.pre_feedforward_layernorm = RMSNorm(config.hidden_size, **nkw) + self.post_feedforward_layernorm = RMSNorm(config.hidden_size, **nkw) + self.self_attn = VisionAttention(config, rngs=rngs) + self.mlp = VisionMLP(config, rngs=rngs) + + def __call__(self, x, cos, sin, position_ids, attn_bias): + residual = x + x = self.input_layernorm(x) + x = self.self_attn(x, cos, sin, position_ids, attn_bias) + x = self.post_attention_layernorm(x) + x = residual + x + + residual = x + x = self.pre_feedforward_layernorm(x) + x = self.mlp(x) + x = self.post_feedforward_layernorm(x) + return residual + x + + +class VisionPatchEmbedder(nnx.Module): + """Port of HF `Gemma4VisionPatchEmbedder`. + + pixel_values are pre-flattened patches of shape [B, P, 3*patch^2]. + """ + + def __init__(self, config: Gemma4VisionConfig, *, rngs: nnx.Rngs): + self.config = config + in_dim = 3 * config.patch_size**2 + self.input_proj = nnx.Linear( + in_dim, + config.hidden_size, + use_bias=False, + rngs=rngs, + param_dtype=config.param_dtype, + dtype=config.dtype, + ) + self.position_embedding_table = nnx.Param( + jnp.ones( + (2, config.position_embedding_size, config.hidden_size), + dtype=config.param_dtype, + ) + ) + + def _position_embeddings(self, pixel_position_ids, padding_positions): + # pixel_position_ids: [B, P, 2] -> [B, P, hidden] + clamped = jnp.maximum(pixel_position_ids, 0) + one_hot = jax.nn.one_hot( + clamped, self.config.position_embedding_size, dtype=jnp.float32 + ) # [B, P, 2, pos_size] + one_hot = jnp.transpose(one_hot, (0, 2, 1, 3)) # [B, 2, P, pos_size] + table = self.position_embedding_table.value.astype(jnp.float32) # [2,pos,h] + pos = jnp.einsum("baps,ash->baph", one_hot, table) # [B, 2, P, h] + pos = jnp.sum(pos, axis=1) # sum over the 2 spatial axes -> [B, P, h] + pos = jnp.where(padding_positions[..., None], 0.0, pos) + return pos + + def __call__(self, pixel_values, pixel_position_ids, padding_positions): + pixel_values = 2.0 * (pixel_values - 0.5) + hidden = self.input_proj(pixel_values.astype(self.input_proj.kernel.value.dtype)) + pos = self._position_embeddings(pixel_position_ids, padding_positions) + return hidden + pos.astype(hidden.dtype) + + +class VisionEncoder(nnx.Module): + """Port of HF `Gemma4VisionEncoder`.""" + + def __init__(self, config: Gemma4VisionConfig, *, rngs: nnx.Rngs): + self.config = config + self.rotary_emb = VisionRotaryEmbedding(config) + self.layers = compat.ModuleList() + for _ in range(config.num_hidden_layers): + self.layers.append(VisionEncoderLayer(config, rngs=rngs)) + + def __call__(self, inputs_embeds, valid_mask, pixel_position_ids): + # valid_mask: [B, P] True for real patches. Build additive bidirectional + # bias [B, 1, 1, P]. + attn_bias = jnp.where( + valid_mask[:, None, None, :], 0.0, jnp.finfo(jnp.float32).min + ) + cos, sin = self.rotary_emb(pixel_position_ids) + x = inputs_embeds + for layer in self.layers: + x = layer(x, cos, sin, pixel_position_ids, attn_bias) + return x + + +class VisionPooler(nnx.Module): + """Port of HF `Gemma4VisionPooler` (2D average pool + sqrt(hidden) scale). + + Returns padded `[B, output_length, hidden]` plus a `[B, output_length]` + validity mask. (HF additionally gathers valid rows into a ragged tensor; we + defer that to the merge step to stay jit-friendly.) + """ + + def __init__(self, config: Gemma4VisionConfig): + self.hidden_size = config.hidden_size + self.root_hidden_size = config.hidden_size**0.5 + + def _avg_pool_by_positions(self, hidden_states, pixel_position_ids, length): + b, input_seq_len, h = hidden_states.shape + k = int((input_seq_len // length) ** 0.5) + k_squared = k**2 + if k_squared * length != input_seq_len: + raise ValueError( + f"Cannot pool seq_len={input_seq_len} to {length}: k={k}^2 * " + f"{length} must equal {input_seq_len}." + ) + clamped = jnp.maximum(pixel_position_ids, 0) + max_x = jnp.max(clamped[..., 0], axis=-1, keepdims=True) + 1 # [B,1] + kernel_idxs = clamped // k # [B, P, 2] + kernel_idxs = kernel_idxs[..., 0] + (max_x // k) * kernel_idxs[..., 1] # [B,P] + weights = jax.nn.one_hot(kernel_idxs, length, dtype=jnp.float32) / k_squared + output = jnp.einsum("bpl,bph->blh", weights, hidden_states.astype(jnp.float32)) + mask = jnp.logical_not(jnp.all(weights == 0, axis=1)) # [B, length] + return output.astype(hidden_states.dtype), mask + + def __call__( + self, hidden_states, pixel_position_ids, padding_positions, output_length + ): + if output_length > hidden_states.shape[1]: + raise ValueError( + f"Requested {output_length} soft tokens > " + f"{hidden_states.shape[1]} patches." + ) + hidden_states = jnp.where( + padding_positions[..., None], 0.0, hidden_states + ) + if hidden_states.shape[1] != output_length: + hidden_states, mask = self._avg_pool_by_positions( + hidden_states, pixel_position_ids, output_length + ) + else: + mask = jnp.logical_not(padding_positions) + hidden_states = hidden_states * self.root_hidden_size + return hidden_states, mask + + +class Gemma4VisionModel(nnx.Module): + """Port of HF `Gemma4VisionModel` (patch embed -> encoder -> pool).""" + + def __init__(self, config: Gemma4VisionConfig, *, rngs: nnx.Rngs): + self.config = config + self.patch_embedder = VisionPatchEmbedder(config, rngs=rngs) + self.encoder = VisionEncoder(config, rngs=rngs) + self.pooler = VisionPooler(config) + + def __call__(self, pixel_values, pixel_position_ids): + pk = self.config.pooling_kernel_size + output_length = pixel_values.shape[-2] // (pk * pk) + padding_positions = jnp.all(pixel_position_ids == -1, axis=-1) # [B, P] + x = self.patch_embedder(pixel_values, pixel_position_ids, padding_positions) + x = self.encoder(x, jnp.logical_not(padding_positions), pixel_position_ids) + hidden_states, mask = self.pooler( + x, pixel_position_ids, padding_positions, output_length + ) + return hidden_states, mask + + +class Gemma4MultimodalEmbedder(nnx.Module): + """Port of HF `Gemma4MultimodalEmbedder` (the `embed_vision` projector).""" + + def __init__( + self, + vision_hidden_size: int, + text_hidden_size: int, + *, + eps: float, + rngs: nnx.Rngs, + param_dtype: jnp.dtype = jnp.float32, + dtype: jnp.dtype = jnp.float32, + ): + self.embedding_pre_projection_norm = RMSNorm( + vision_hidden_size, + eps=eps, + with_scale=False, + rngs=rngs, + param_dtype=param_dtype, + ) + self.embedding_projection = nnx.Linear( + vision_hidden_size, + text_hidden_size, + use_bias=False, + rngs=rngs, + param_dtype=param_dtype, + dtype=dtype, + ) + + def __call__(self, soft_tokens): + return self.embedding_projection( + self.embedding_pre_projection_norm(soft_tokens) + ) + + +class Gemma4VisionStack(nnx.Module): + """Vision tower + ``embed_vision`` projector under the checkpoint's names. + + The submodule names (``vision_tower``, ``embed_vision``) match the real + checkpoint prefixes, so loading is a plain ``model.`` strip + per-tensor + transpose. This is also the shape Stage 4 plugs into ``Gemma4``. + """ + + def __init__( + self, + config: Gemma4VisionConfig, + text_hidden_size: int, + *, + rngs: nnx.Rngs, + ): + self.config = config + self.vision_tower = Gemma4VisionModel(config, rngs=rngs) + self.embed_vision = Gemma4MultimodalEmbedder( + config.hidden_size, + text_hidden_size, + eps=config.rms_norm_eps, + rngs=rngs, + param_dtype=config.param_dtype, + dtype=config.dtype, + ) + + def __call__(self, pixel_values, pixel_position_ids): + """Returns (soft_tokens_in_text_space [B, S, text_hidden], valid_mask).""" + soft, mask = self.vision_tower(pixel_values, pixel_position_ids) + return self.embed_vision(soft), mask diff --git a/tunix/rl/common.py b/tunix/rl/common.py index 8a0efac88..96b73fef9 100644 --- a/tunix/rl/common.py +++ b/tunix/rl/common.py @@ -14,6 +14,7 @@ """Common RL helper classes and functions.""" from functools import partial # pylint: disable=g-importing-member +import inspect from typing import Any, Iterable import flax @@ -28,6 +29,34 @@ build_positions_from_mask = utils.build_positions_from_mask +def _model_accepts_segment_ids(model: Any) -> bool: + """Returns True if ``model.__call__`` declares a ``segment_ids`` parameter. + + Some models (e.g. Qwen3) thread ``segment_ids`` into a pallas splash-attention + kernel to suppress attention across packed-document or padding boundaries; + others (Gemma 2/3/4, Llama) don't accept the keyword at all. Passing it to a + model that doesn't accept it raises a TypeError, so callers should gate the + pass-through on this helper. + + Args: + model: A live nnx.Module instance (typically from ``nnx.merge``). + + Returns: + True if the model's ``__call__`` accepts ``segment_ids`` as a named + parameter, or declares a ``**kwargs`` catch-all; False otherwise (including + if introspection raises for any reason). + """ + try: + sig = inspect.signature(model.__call__) + except (ValueError, TypeError): + return False + if 'segment_ids' in sig.parameters: + return True + return any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) + + class RepeatIterable(Iterable[Any]): """An iterable that processes a list of rollout batches. @@ -347,20 +376,11 @@ def compute_per_token_logps( "attention_mask": attn_mask, } # Pass through any segment ids so the model's attention kernel can respect - # them only if the model signature accepts it: caller-provided packing ids take - # precedence; otherwise we pass the per-position non-pad mask derived in - # ``process_ids`` so flash-attention variants that lack a separate - # padding-mask input still skip pad positions. - import inspect # pylint: disable=g-import-not-at-top - try: - sig = inspect.signature(model.__call__) - has_segment_ids = ("segment_ids" in sig.parameters) or any( - p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() - ) - except Exception: - has_segment_ids = False - - if has_segment_ids: + # them, but only if the model signature accepts the keyword. Caller-provided + # packing ids take precedence; otherwise we pass the per-position non-pad + # mask derived in ``process_ids`` so flash-attention variants that lack a + # separate padding-mask input still skip pad positions. + if _model_accepts_segment_ids(model): if segment_ids is not None: model_kwargs["segment_ids"] = segment_ids elif input_seg_ids is not None: @@ -434,10 +454,17 @@ def compute_score( model_kwargs = {"positions": calculated_positions, "cache": None} if segment_ids is not None: + # Caller explicitly asked for packed mode; pass segment_ids through. If the + # model doesn't accept it the resulting TypeError points at user code, not + # at a hidden auto-derived value. model_kwargs["segment_ids"] = segment_ids else: model_kwargs["attention_mask"] = attn_mask - if input_seg_ids is not None: + # Only forward the auto-derived per-position non-pad mask to models that + # actually accept ``segment_ids`` (e.g. Qwen3). Without this gate, Gemma + # 2/3/4 / Llama would crash with ``TypeError: __call__() got an unexpected + # keyword argument 'segment_ids'``. See google/tunix#1539. + if input_seg_ids is not None and _model_accepts_segment_ids(model): model_kwargs["segment_ids"] = input_seg_ids out = model(prompt_completion_ids, **model_kwargs)