diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 26359cdddc..22f80d9e09 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -3107,7 +3107,9 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers } ) - # TODO: gemma4-small multimodal not yet supported — vision-encoder mappings below are dead. + # Gemma-4 E2B/E4B vision-encoder param mapping. Active when use_multimodal is set; + # the clipped-linears activation clip bounds are additionally mapped when + # use_clipped_linears_for_vit is enabled (required for image parity on E2B/E4B). if maxtext_config.use_multimodal and vcfg: nvis = vcfg.get("num_hidden_layers", 0) mapping.update( @@ -3163,6 +3165,23 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers f"{prefix}-mlp-wo-kernel": f"{hf_prefix}.mlp.down_proj.linear.weight", } ) + # Gemma-4 vision clipped-linears: per-projection activation clip bounds + # (scalar {input,output}_{min,max}) carried in the reference checkpoint. + # Only mapped when the clipped-linears path is enabled; the nnx leaves live + # at _clip.{input,output}_{min,max} under attention/mlp. + if getattr(maxtext_config, "use_clipped_linears_for_vit", False): + _clip_proj = { + "attention-q_clip": f"{hf_prefix}.self_attn.q_proj", + "attention-k_clip": f"{hf_prefix}.self_attn.k_proj", + "attention-v_clip": f"{hf_prefix}.self_attn.v_proj", + "attention-o_clip": f"{hf_prefix}.self_attn.o_proj", + "mlp-gate_clip": f"{hf_prefix}.mlp.gate_proj", + "mlp-up_clip": f"{hf_prefix}.mlp.up_proj", + "mlp-down_clip": f"{hf_prefix}.mlp.down_proj", + } + for mt_sub, hf_proj in _clip_proj.items(): + for bound in ("input_min", "input_max", "output_min", "output_max"): + mapping[f"{prefix}-{mt_sub}-{bound}"] = f"{hf_proj}.{bound}" return {k: v for k, v in mapping.items() if v is not None} diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 7a725fc4ca..3b61b967f4 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1188,6 +1188,18 @@ freeze_vision_encoder_params: true freeze_audio_encoder_params: true dtype_mm: "float32" # Data type for multimodal model's vision encoder remat_policy_for_vit: "minimal" # Remat policy for multimodal model's vision encoder. Check `remat_policy` for options. +# Gemma-4 vision only: apply the per-projection activation clip bounds carried in the +# reference checkpoint (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj each have +# scalar {input,output}_{min,max}). A prerequisite for Gemma-4 E2B/E4B image parity +# (necessary but not on its own sufficient); no-op for other vision encoders. Bounds are +# checkpoint-resident, non-trainable scalars. +use_clipped_linears_for_vit: false +# Gemma-4 E2B/E4B decoder image-handling (defaults preserve behavior for other models): +use_bidirectional_image_attn: false # E2B/E4B image spans are causal +ple_pad_substitute_image_rows: false # substitute pad id for image rows in the per-layer-embedding path (HF gemma4) +ple_pad_mode: "identity" # 'identity' (token-id path) or 'both' (also context/embedding path) +image_placeholder_token_id: 258880 # GEMMA4_TOKEN_PLACEHOLDER +ple_pad_token_id: 0 # E2B text_config.pad_token_id image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg" video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4" diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml index 81c8d4ea66..502e84737b 100644 --- a/src/maxtext/configs/models/gemma4-e2b.yml +++ b/src/maxtext/configs/models/gemma4-e2b.yml @@ -44,7 +44,8 @@ global_rope_proportion: 0.25 local_rope_proportion: 1.0 final_logits_soft_cap: 30.0 -# Vision encoder flags — multimodal not yet supported for E2B / E4B. +# Vision encoder flags. Image (multimodal) parity requires use_clipped_linears_for_vit=true +# (set at runtime alongside use_multimodal=true); the flags below configure the E2B image contract. vision_encoder_block: "gemma4" rope_theta_for_vit: 100 image_size_for_vit: [672, 960] @@ -58,3 +59,11 @@ num_attention_heads_for_vit: 12 image_placeholder: "<|image|>" vision_output_length: 280 num_position_embeddings_for_vit: 10240 +# E2B image contract (decoder-side): image spans are causal and the per-layer-embedding path +# substitutes the pad token for image placeholder rows (matches HF modeling_gemma4). +use_bidirectional_image_attn: False +ple_pad_substitute_image_rows: True +ple_pad_mode: "identity" +image_placeholder_token_id: 258880 +ple_pad_token_id: 0 +packing: false # Gemma-4 E2B/E4B multimodal does not support packed image spans (see configs/types.py gate) diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml index b8da0d1d46..d2cf5bdd7e 100644 --- a/src/maxtext/configs/models/gemma4-e4b.yml +++ b/src/maxtext/configs/models/gemma4-e4b.yml @@ -45,7 +45,8 @@ global_rope_proportion: 0.25 local_rope_proportion: 1.0 final_logits_soft_cap: 30.0 -# Vision encoder flags — multimodal not yet supported for E2B / E4B. +# Vision encoder flags. Image (multimodal) parity requires use_clipped_linears_for_vit=true +# (set at runtime alongside use_multimodal=true); the flags below configure the E4B image contract. vision_encoder_block: "gemma4" rope_theta_for_vit: 100 image_size_for_vit: [672, 960] @@ -59,3 +60,11 @@ num_attention_heads_for_vit: 12 image_placeholder: "<|image|>" vision_output_length: 280 num_position_embeddings_for_vit: 10240 +# E4B image contract (decoder-side): image spans are causal and the per-layer-embedding path +# substitutes the pad token for image placeholder rows (matches HF modeling_gemma4). +use_bidirectional_image_attn: False +ple_pad_substitute_image_rows: True +ple_pad_mode: "identity" +image_placeholder_token_id: 258880 +ple_pad_token_id: 0 +packing: false # Gemma-4 E2B/E4B multimodal does not support packed image spans (see configs/types.py gate) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a548422..69a3d29270 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2137,6 +2137,47 @@ class MultimodalGeneral(BaseModel): use_multimodal: bool = Field(False, description="Enable multimodal capabilities.") attention_for_vit: str = Field("dot_product", description="The attention algorithm to use for vision encoder.") + use_clipped_linears_for_vit: bool = Field( + False, + description=( + "Gemma-4 vision only: apply the per-projection activation clip bounds carried in the reference " + "checkpoint (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj each have scalar " + "{input,output}_{min,max}). A prerequisite for Gemma-4 E2B/E4B image parity (necessary but not " + "on its own sufficient); no-op for other encoders." + ), + ) + use_bidirectional_image_attn: bool = Field( + False, + description=( + "Whether image placeholder tokens attend bidirectionally in the text decoder. Gemma-4 E2B/E4B " + "use causal image spans (False); bidirectional-image models (Gemma-3, gemma4-26b/31b) use True." + ), + ) + ple_pad_substitute_image_rows: bool = Field( + False, + description=( + "Gemma-4 E2B/E4B per-layer-embedding (PLE) path: substitute ple_pad_token_id for image placeholder " + "rows before the per-layer embedder, matching HF modeling_gemma4 (llm_input_ids pad substitution). " + "Default False preserves the native PLE for other models." + ), + ) + ple_pad_mode: str = Field( + "identity", + description=( + "PLE pad-substitution scope when ple_pad_substitute_image_rows=True. 'identity' (DEFAULT, " + "HF-faithful): map image placeholder rows -> pad in the token-identity PLE path only; the context/" + "projection path keeps the merged image features (matches HF Transformers 5.9.0 " + "Gemma4ForConditionalGeneration, where get_per_layer_inputs ignores inputs_embeds when input_ids is " + "provided). 'both': additionally overwrite the context path with the pad embedding — this is " + "HF-DIVERGENT and provided only for ablation. Validated as an enum; unknown values hard-fail." + ), + ) + image_placeholder_token_id: int = Field( + 258880, description="Gemma-4 image placeholder token id (GEMMA4_TOKEN_PLACEHOLDER)." + ) + ple_pad_token_id: int = Field( + 0, description="Pad token id used for PLE image-row substitution (Gemma-4 E2B text_config.pad_token_id=0)." + ) vision_encoder_block: VisionEncoderBlockType = Field( VisionEncoderBlockType.NONE, description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').", @@ -3572,16 +3613,84 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with nn.scan)." ) if self.use_multimodal: - # Gemma 4 small (E2B / E4B) only supports text for now; multimodal - # support is pending clipped-linears in the vision encoder. + # Gemma 4 small (E2B / E4B) multimodal requires the vision-encoder clipped-linears AND the + # padded-patch masking / position-threading path; gate on the clipped-linears flag. + if self.model_name in ("gemma4-e2b", "gemma4-e4b") and not self.use_clipped_linears_for_vit: + raise ValueError( + f"Multimodal for {self.model_name} requires use_clipped_linears_for_vit=True " + "(the vision encoder ships per-projection activation clip bounds; without them the " + "image span diverges). Set use_clipped_linears_for_vit=True to enable image inputs." + ) + # ---- Gemma-4 E2B/E4B multimodal STATIC contract gate (fail-closed) ---- + # These invariants encode the semantics validated against pinned HF Transformers 5.9.0 + # (Gemma4ForConditionalGeneration). Any deviation silently corrupts image/post-image logits, so we + # refuse to build the model rather than degrade to a wrong-but-runnable path. if self.model_name in ("gemma4-e2b", "gemma4-e4b"): - raise ValueError(f"Multimodal is not yet supported for {self.model_name}; only text inputs are supported.") + # (a) PLE pad-substitution mode must be a known value. HF maps image placeholder tokens -> pad in the + # token-identity PLE path only (the context path keeps the merged image features); that is + # "identity". "both" additionally overwrites the context path, which is HF-DIVERGENT. We keep the + # knob for experimentation but hard-fail unknown/typo values instead of silently defaulting. + _valid_ple_modes = ("identity", "both") + if str(self.ple_pad_mode) not in _valid_ple_modes: + raise ValueError( + f"ple_pad_mode='{self.ple_pad_mode}' is not one of {_valid_ple_modes}. " + f"For Gemma-4 E2B/E4B the HF-faithful contract is 'identity' (token-identity PLE path maps " + f"image rows -> pad; context path keeps merged image features). 'both' is HF-divergent and " + f"provided only for ablation. Refusing to run with an unrecognized PLE mode." + ) + # (b) The image placeholder id and PLE pad id are semantic constants tied to the tokenizer/model. We + # require them to be set explicitly (single source of truth: the model yml derived from HF config) + # so a silent hidden default cannot mask a tokenizer mismatch. + if self.ple_pad_substitute_image_rows: + if int(self.image_placeholder_token_id) < 0: + raise ValueError( + "image_placeholder_token_id must be a valid non-negative token id when " + "ple_pad_substitute_image_rows=True (derive it from the model/tokenizer config)." + ) + if int(self.ple_pad_token_id) < 0: + raise ValueError( + "ple_pad_token_id must be a valid non-negative token id when " + "ple_pad_substitute_image_rows=True (Gemma-4 E2B text_config.pad_token_id=0)." + ) + # (c) Gemma-4 E2B/E4B image spans are CAUSAL. Bidirectional image attention is a Gemma-3 / 26B / 31B + # feature and would change the attention pattern for E2B/E4B. + if bool(getattr(self, "use_bidirectional_image_attn", False)): + raise ValueError( + f"{self.model_name} uses CAUSAL image spans; use_bidirectional_image_attn must be False. " + "Bidirectional image attention is for Gemma-3 / gemma4-26b / gemma4-31b." + ) + # (d) The clipped path clamps q/k/v and gate/up/down per-projection; a fused QKV or fused MLP would + # apply a single clamp and silently bypass the per-projection clip semantics. Fail closed here + # (defense in depth alongside the runtime guards in gemma4_vision). + if bool(getattr(self, "fused_qkv", False)): + raise ValueError( + f"{self.model_name} multimodal clipped-linears require fused_qkv=False " + "(distinct q/k/v activation clip bounds must be applied per-projection)." + ) + if bool(getattr(self, "fused_mlp", False)): + raise ValueError( + f"{self.model_name} multimodal clipped-linears require fused_mlp=False " + "(distinct gate/up/down activation clip bounds must be applied per-projection)." + ) + # (e) Sequence packing is not supported for Gemma-4 E2B/E4B multimodal in ANY training mode. The stock + # data pipeline does not pack image spans, and packing image-bearing examples together would risk + # cross-document image attention and PLE image-row substitution bleeding across segment boundaries. + # MaxText already forbids packing for multimodal SFT; extend that to every mode for these models + # (fail closed rather than silently produce cross-doc image attention). + if bool(getattr(self, "packing", False)): + raise ValueError( + f"{self.model_name} multimodal does not support sequence packing (packing=True). The stock data " + "pipeline does not pack image spans; packing image-bearing examples risks cross-document image " + "attention and PLE image-row substitution across segment boundaries. Set packing=False." + ) valid_mm_models = ( "gemma3-4b", "gemma3-12b", "gemma3-27b", "gemma4-26b", "gemma4-31b", + "gemma4-e2b", + "gemma4-e4b", "llama4-17b-16e", "llama4-17b-128e", "qwen3-omni-30b-a3b", diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 42753eb752..fad5eefa0b 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -1679,6 +1679,18 @@ def _apply_gemma4_small_layers( per_layer_inputs = None if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0: + # Fail-closed: the pure-Linen Gemma-4-small decoder does NOT implement the PLE image-row pad + # substitution (the HF-faithful mapping of image placeholder rows -> pad in the token-identity PLE + # path). That logic lives in the NNX decoder (nnx_decoders._apply_gemma4_small_layers) and is exercised + # by the supported enable_nnx=True path. Running E2B/E4B multimodal through this pure-Linen path would + # silently diverge on the image-span / post-image logits, so we refuse rather than degrade quietly. + if bool(getattr(cfg, "ple_pad_substitute_image_rows", False)) and multimodal_input is not None: + raise NotImplementedError( + "Gemma-4 E2B/E4B multimodal PLE pad-substitution is only implemented on the NNX decoder path " + "(enable_nnx=True, the default). The pure-Linen decoder does not perform the image-row PLE " + "substitution and would silently diverge from the HF reference. Set enable_nnx=True (default) to " + "run Gemma-4 E2B/E4B multimodal." + ) per_layer_inputs = gemma4_small.PLEToLinen( config=cfg, mesh=mesh, diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py index e5e32d044a..a5a0c38ff8 100644 --- a/src/maxtext/layers/encoders.py +++ b/src/maxtext/layers/encoders.py @@ -106,13 +106,21 @@ def _setup_vision_encoder_layers(self): return encoder_name, projector_name - def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False): + def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False, + image_position_ids=None): # vision encoder output, frozen params in many cases encoder = getattr(self, self.encoder_name) + vision_image_masks = None if self.vision_encoder_block.value.startswith("qwen3") and input_masks is not None: encoder_output = encoder( input_images, video_mask=input_masks, video_grid_thw=video_grid_thw, deterministic=deterministic ) + elif self.vision_encoder_block == VisionEncoderBlockType.GEMMA4 and image_position_ids is not None: + # Gemma-4 padded-patch path: pre-patchified patches + per-patch positions (-1 = pad). The + # encoder returns (embeddings, image_masks); the mask marks the valid pooled tokens. + encoder_output = encoder(input_images, deterministic=deterministic, image_position_ids=image_position_ids) + embeddings, vision_image_masks = encoder_output + encoder_output = embeddings else: encoder_output = encoder(input_images, deterministic=deterministic) deep_feats = None @@ -131,7 +139,7 @@ def __call__(self, input_images, input_masks=None, video_grid_thw=None, determin projector = getattr(self, self.projector_name) embeddings = projector(embeddings) - return embeddings, deep_feats + return embeddings, deep_feats, vision_image_masks class MultimodalMLPProjector(nnx.Module): diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 895ea27c14..68412bb6fc 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -1754,6 +1754,7 @@ def __call__( attention_metadata=attention_metadata, previous_chunk=previous_chunk, slot=slot, + shared_embedding=shared_embedding, ) elif cfg.scan_layers: if self.is_deepseek: @@ -2227,14 +2228,47 @@ def _apply_gemma4_small_layers( attention_metadata=None, previous_chunk=None, slot=None, + shared_embedding=None, ): """Apply Gemma 4 small (E2B/E4B) decoder layers (pure-NNX).""" cfg = self.config bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None + # Gemma-4 E2B/E4B image spans are causal (text_config.use_bidirectional_attention is unset), unlike + # the bidirectional-image Gemma-3 / 26B / 31B models. Suppress the bidirectional attention carve-out + # unless explicitly enabled via config (gate on a flag, not the model name). + if not bool(getattr(cfg, "use_bidirectional_image_attn", False)): + bidirectional_mask_value = None per_layer_inputs = None if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0: - per_layer_inputs = self.per_layer_embedder(decoder_input_tokens, y) + ple_tokens = decoder_input_tokens + ple_context = y + # Gemma-4 E2B/E4B build the per-layer inputs from llm_input_ids with the image placeholder + # tokens mapped to pad_token_id (HF modeling_gemma4.py), rather than feeding the image + # placeholder id / merged image features into the PLE path. Without this substitution the + # per-layer embeddings at the image placeholder positions diverge from the reference, which + # corrupts the image-span and post-image logits. Gated on ple_pad_substitute_image_rows + # (default False preserves the native PLE for other models). + if bool(getattr(cfg, "ple_pad_substitute_image_rows", False)) and multimodal_input is not None: + _img_id = int(getattr(cfg, "image_placeholder_token_id", 258880)) + _pad_id = int(getattr(cfg, "ple_pad_token_id", 0)) + _img_row = decoder_input_tokens.astype(jnp.int32) == _img_id + ple_tokens = jnp.where(_img_row, _pad_id, decoder_input_tokens.astype(jnp.int32)) + # ple_pad_mode: 'identity' (HF-faithful default) substitutes ONLY the token-identity path (above), + # leaving the context/projection path reading the merged image features — this matches HF + # Transformers 5.9.0 (get_per_layer_inputs ignores inputs_embeds when input_ids is given). The + # 'both' mode additionally overwrites the context path with the pad embedding; this is HF-DIVERGENT + # and provided only for ablation (see M9 ablation matrix). Enum-validated in configs/types.py. + if str(getattr(cfg, "ple_pad_mode", "identity")) == "both": + if shared_embedding is None: + raise ValueError( + "ple_pad_mode='both' requires the shared token embedding to compute the pad context " + "vector, but shared_embedding was not threaded into the Gemma-4-small PLE path. This is " + "an HF-divergent ablation mode; use the default 'identity' for HF-faithful behavior." + ) + _pad_vec = shared_embedding(jnp.full_like(decoder_input_tokens, _pad_id).astype(jnp.int32), model_mode=model_mode) + ple_context = jnp.where(_img_row[..., None], _pad_vec, y) + per_layer_inputs = self.per_layer_embedder(ple_tokens, ple_context) layer_types = gemma4_small.build_layer_types(cfg.num_decoder_layers, cfg.model_name) num_kv_shared = cfg.num_kv_shared_layers diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py index 72cd841f81..b65836b39c 100644 --- a/src/maxtext/models/gemma4_vision.py +++ b/src/maxtext/models/gemma4_vision.py @@ -16,6 +16,8 @@ """Vision transformer implementation for Gemma4.""" from typing import cast +import functools +import operator import jax import jax.numpy as jnp from flax import linen as nn @@ -30,6 +32,230 @@ from maxtext.layers import normalizations +# ============================================================================= +# Gemma-4 vision clipped-linears (Navi upstream contribution) +# ----------------------------------------------------------------------------- +# The Gemma-4 E2B/E4B vision tower ships per-projection activation clip bounds in +# the reference (HF) checkpoint: for each of the 7 vision projections +# (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj) in each of the 16 encoder +# blocks, a scalar {input_min,input_max,output_min,output_max} = 16*7*4 = 448 +# checkpoint tensors. The reference forward clamps each projection's input by +# [input_min,input_max] and its output by [output_min,output_max]. Omitting the +# clamps produces large activation drift in the image span (empirically KL 4-17 +# on a 340-token teacher-forced parity harness), because a handful of vision +# activations blow up without the trained saturation. Upstream MaxText marks +# E2B/E4B multimodal "not yet supported" and does not model these bounds. +# +# This module adds them as OPT-IN, checkpoint-resident, NON-TRAINABLE scalars, +# gated on ``config.use_clipped_linears_for_vit`` (exact no-op when False). +# Design: plain ``nnx.Param`` bounds (so they map to the canonical ``params`` +# collection and round-trip through the nnx->linen->orbax checkpoint path), plus +# a leaf-PATH optimizer-freeze mask so the 448 bounds are excluded from optimizer +# updates and weight decay without a custom nnx.Variable subclass (subclasses are +# renamed by the linen bridge and silently dropped from the saved checkpoint). +# A NaN sentinel + ``validate_clip_bounds`` hard-fails on a missing/non-finite +# bound rather than silently degrading to an identity clamp. + +# Leaf-name tokens that identify a clip-bound scalar in a flattened params tree. +_CLIP_LEAF_TOKENS = ("q_clip", "k_clip", "v_clip", "o_clip", "gate_clip", "up_clip", "down_clip") +_CLIP_BOUND_NAMES = ("input_min", "input_max", "output_min", "output_max") + + +def _mk_clip_bound(init_val=jnp.nan): + """A checkpoint-resident scalar clip bound as a plain ``nnx.Param`` (maps to the + canonical ``params`` collection; NaN sentinel marks an unloaded bound).""" + return nnx.Param(jnp.asarray(init_val, dtype=jnp.float32)) + + +def _is_clip_bound_path(path) -> bool: + """True iff a flattened-params key path addresses a clip-bound scalar.""" + s = "/".join(str(getattr(p, "key", p)) for p in path) if not isinstance(path, str) else path + return any(tok in s for tok in _CLIP_LEAF_TOKENS) and any(b in s for b in _CLIP_BOUND_NAMES) + + +def clip_optimizer_freeze_mask(params_tree): + """Bool pytree (same structure as ``params_tree``): True for TRAINABLE leaves, + False for the immutable clip bounds. Feed to ``optax.masked``/``multi_transform`` + so the bounds get ``set_to_zero()`` updates. Path-based, so it survives the + nnx->linen->orbax round-trip regardless of leaf type erasure.""" + flat = jax.tree_util.tree_flatten_with_path(params_tree)[0] + leaves_mask = [not _is_clip_bound_path(path) for path, _ in flat] + treedef = jax.tree_util.tree_structure(params_tree) + return jax.tree_util.tree_unflatten(treedef, leaves_mask) + + +def _clip_bound_value(bound, dtype): + """Read a clip-bound scalar as a constant: stop_gradient defends the bounds from receiving + gradient at the clamp use-site (defense in depth alongside the optimizer freeze mask), so clip-bound + gradients can never contaminate gradient statistics even if the freeze mask were misconfigured.""" + return jax.lax.stop_gradient(bound.value).astype(dtype) + + +def _clip_in(x, cb): + """clamp(x, input_min, input_max) in x's dtype; no-op if ``cb`` is None. Bounds are stop_gradient'd.""" + if cb is None: + return x + xd = x.dtype + return jnp.clip(x, _clip_bound_value(cb.input_min, xd), _clip_bound_value(cb.input_max, xd)) + + +def _clip_out(y, cb): + """clamp(y, output_min, output_max) in y's dtype; no-op if ``cb`` is None. Bounds are stop_gradient'd.""" + if cb is None: + return y + yd = y.dtype + return jnp.clip(y, _clip_bound_value(cb.output_min, yd), _clip_bound_value(cb.output_max, yd)) + + +class _ClipBounds(nnx.Module): + """Holds the four checkpoint-resident scalar clip bounds as ``nnx.Param`` leaves.""" + + def __init__(self): + self.input_min = _mk_clip_bound() + self.input_max = _mk_clip_bound() + self.output_min = _mk_clip_bound() + self.output_max = _mk_clip_bound() + + +def _make_clip_state(): + """Four NaN-sentinel scalar bounds (checkpoint-resident, non-trainable).""" + return _ClipBounds() + + +def validate_clip_bounds(cb, where=""): + """Hard-fail: every bound finite + scalar, and input_min<=input_max, output_min<=output_max. + Raises ValueError otherwise. No-op if ``cb`` is None.""" + if cb is None: + return + vals = {} + for nm in ("input_min", "input_max", "output_min", "output_max"): + v = getattr(cb, nm).value + if getattr(v, "shape", ()) not in ((), (1,)): + raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} has non-scalar " + f"shape {v.shape}; expected scalar.") + fv = float(jnp.reshape(v, (-1,))[0]) + if not bool(jnp.isfinite(jnp.asarray(fv))): + raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} = {fv} is non-finite " + f"(missing/NaN/Inf). use_clipped_linears_for_vit=True declares a FINITE clipped model; " + f"refusing to fall back to an identity clamp.") + vals[nm] = fv + if vals["input_min"] > vals["input_max"]: + raise ValueError(f"Gemma4 vision clip bound{(' @ '+where) if where else ''}: input_min={vals['input_min']} " + f"> input_max={vals['input_max']} (a clamp with min>max would empty the interval).") + if vals["output_min"] > vals["output_max"]: + raise ValueError(f"Gemma4 vision clip bound{(' @ '+where) if where else ''}: output_min={vals['output_min']} " + f"> output_max={vals['output_max']} (a clamp with min>max would empty the interval).") + + +# Expected clipped-module accounting for a Gemma-4 vision tower: 16 encoder blocks x 7 projections +# (attention q/k/v/o = 4 modules holding q/k/v/o clip states + mlp gate/up/down = 3) -> but the clip +# STATE lives on 2 modules per block (the Gemma4Attention holding q/k/v/o_clip, and the +# Gemma4ClippedMlpBlock holding gate/up/down_clip). The scalar bound count is the invariant that matters: +# 16 blocks x 7 projections x 4 bounds = 448. We validate the projection-level clip states (16x7 = 112). +EXPECTED_CLIP_PROJECTIONS = 112 # 16 blocks x 7 projections (q,k,v,o,gate,up,down) +EXPECTED_CLIP_BOUNDS = 448 # 112 projections x 4 scalar bounds + + +def validate_all_vision_clip_bounds(model, *, expected_projections=EXPECTED_CLIP_PROJECTIONS, + expected_bounds=EXPECTED_CLIP_BOUNDS): + """Post-checkpoint-load, pre-first-JIT validation of ALL Gemma-4 vision clip bounds. Fail-closed. + + Walks the model graph, collects every ``_ClipBounds`` module (each carries the 4 scalars for one + projection), validates each (finite, scalar, min<=max), and asserts EXACT counts: + ``expected_projections`` clip-state modules and ``expected_bounds`` scalar leaves. Raises ValueError on + any deficiency. The exact-count check is what prevents a traversal that accidentally finds zero modules + from silently "passing". + """ + from flax import nnx # pylint: disable=import-outside-toplevel + n_proj = 0 + n_bounds = 0 + for path, mod in nnx.iter_graph(model): + # A clip-state module has exactly the four bound leaves. + if (hasattr(mod, "input_min") and hasattr(mod, "input_max") + and hasattr(mod, "output_min") and hasattr(mod, "output_max") + and not isinstance(mod, (int, float))): + where = "/".join(str(getattr(p, "key", p)) for p in path) if isinstance(path, (list, tuple)) else str(path) + validate_clip_bounds(mod, where) + n_proj += 1 + n_bounds += 4 + if n_proj != expected_projections: + raise ValueError( + f"Gemma-4 vision clip validation found {n_proj} clip-state modules, expected exactly " + f"{expected_projections} (16 blocks x 7 projections). A mismatch means the clip bounds were not " + f"loaded/mapped correctly; refusing to run with a partially-clipped vision tower." + ) + if n_bounds != expected_bounds: + raise ValueError( + f"Gemma-4 vision clip validation found {n_bounds} clip-bound scalars, expected exactly {expected_bounds}." + ) + return n_proj, n_bounds + + +# Sentinel value marking a padded (phantom) patch position in the Option-S padded-patch contract. +POSITIONS_PAD_VALUE = -1 + + +def option_s_position_sentinel_ok(positions_xy): + """Return (all_rows_valid_or_full_sentinel, n_bad_rows) for an Option-S position tensor. + + The padded-patch contract requires every per-patch position row to be EXACTLY one of: + * [x, y] with both coordinates >= 0 (a valid patch), or + * [POSITIONS_PAD_VALUE, POSITIONS_PAD_VALUE] (a fully-padded sentinel). + A MIXED row such as [-1, y] or [x, -1] is malformed: it would be classified as "valid" by the + ``(pos == -1).all(axis=-1)`` pad test (silently attended-to as a real patch) yet carry a negative + coordinate into pooling. This helper flags such rows so callers can fail closed rather than + silently mis-handle them. Pure (numpy/jax-eager friendly); does not allocate on-device state. + + Args: + positions_xy: int array [..., 2] of per-patch (x, y) positions. + Returns: + (ok: bool, n_bad: int) where ok is True iff no malformed (mixed-sentinel) row exists. + """ + import numpy as _np # local import: this is called eagerly (host) for validation, not in the JIT forward + p = _np.asarray(positions_xy) + if p.shape[-1] != 2: + raise ValueError(f"option_s positions must have last dim 2 (x,y); got shape {p.shape}") + is_neg = p == POSITIONS_PAD_VALUE + any_neg = is_neg.any(axis=-1) + all_neg = is_neg.all(axis=-1) + # A row is malformed iff it has SOME sentinel coord but is not FULLY sentinel. + bad = _np.logical_and(any_neg, _np.logical_not(all_neg)) + # Also reject any negative coordinate other than the exact sentinel (e.g. -2), which is never valid. + other_neg = _np.logical_and(p < 0, p != POSITIONS_PAD_VALUE).any(axis=-1) + bad = _np.logical_or(bad, other_neg) + n_bad = int(bad.sum()) + return (n_bad == 0), n_bad + + +def validate_option_s_positions(positions_xy, where=""): + """Hard-fail if any Option-S position row is a malformed mixed sentinel. No-op if positions is None.""" + if positions_xy is None: + return + ok, n_bad = option_s_position_sentinel_ok(positions_xy) + if not ok: + raise ValueError( + f"Gemma-4 Option-S position sentinel integrity{(' @ '+where) if where else ''}: found {n_bad} malformed " + f"row(s). Every patch position must be [x,y] with both coords >= 0, or [-1,-1] (full sentinel). Mixed " + f"rows like [-1, y] / [x, -1] (or other negatives) would be silently treated as valid patches and " + f"corrupt pooling; refusing to run." + ) + + +def option_s_pooled_matches_placeholder(image_mask, n_placeholder): + """Return (ok, n_valid_pooled) checking the count of valid pooled image tokens == n_placeholder. + + ``image_mask`` is the [..., K] validity mask returned by the vision pooler (True where a pooled slot + received at least one real patch). ``n_placeholder`` is the number of image placeholder positions in the + decoder sequence for the SAME sample. The Option-S contract requires these to be equal so that the pooled + image tokens scatter 1:1 onto the placeholder rows; a mismatch means the merge would drop or duplicate + image tokens regardless of padded-embedding ordering. + """ + import numpy as _np + m = _np.asarray(image_mask) + n_valid = int(_np.asarray(m).astype(bool).sum()) + return (n_valid == int(n_placeholder)), n_valid + + def factorized_posemb(posemb: jax.Array, positions_xy: jax.Array, precision) -> jax.Array: """Computes factorized position embedding from (x, y) coordinates. @@ -409,7 +635,15 @@ def __call__(self, inputs: jax.Array, positions: jax.Array) -> jax.Array: class Gemma4Attention(attentions.Attention): - """Gemma 4 specific Attention module.""" + """Gemma 4 specific Attention module. + + When ``use_clipped_linears`` is enabled, the q/k/v/o projections apply the + per-projection activation clip bounds carried in the Gemma-4 vision checkpoint + (input clamp before the matmul, output clamp after). The clamps are wired by + overriding the base ``Attention`` projection methods, so the underlying + ``DenseGeneral`` weights and their checkpoint key paths are unchanged. When the + flag is off, every override is an exact delegate to the base implementation. + """ def init_rotary_embedding(self) -> Gemma4VisionRotaryEmbedding: """Initializes the rotary position embedding module for Gemma 4 vision.""" @@ -418,6 +652,124 @@ def init_rotary_embedding(self) -> Gemma4VisionRotaryEmbedding: rotary_fraction=None, # Or assume it from config if available ) + def enable_vision_clip_bounds(self): + """Attach the four checkpoint-resident clip-bound scalars for each of q/k/v/o. + + Called once by ``Gemma4EncoderBlock`` after construction when + ``config.use_clipped_linears_for_vit`` is set. Idempotent. + """ + if getattr(self, "_use_clipped_linears", False): + return + # Fail-closed: the clipped path clamps q/k/v independently, which requires SEPARATE q/k/v projections. + # A fused QKV projection would apply a single clamp and silently bypass the per-projection clip semantics. + if bool(getattr(self.config, "fused_qkv", False)): + raise ValueError( + "Gemma-4 vision clipped-linears require fused_qkv=False: the checkpoint carries distinct " + "q/k/v activation clip bounds that must be applied per-projection. Refusing to silently clip a " + "fused QKV projection." + ) + self._use_clipped_linears = True + self.q_clip = _make_clip_state() + self.k_clip = _make_clip_state() + self.v_clip = _make_clip_state() + self.o_clip = _make_clip_state() + + def validate_clip_bounds(self): + if not getattr(self, "_use_clipped_linears", False): + return + validate_clip_bounds(self.q_clip, "q_proj") + validate_clip_bounds(self.k_clip, "k_proj") + validate_clip_bounds(self.v_clip, "v_proj") + validate_clip_bounds(self.o_clip, "o_proj") + + # --- projection overrides: clamp(input) -> DenseGeneral -> clamp(output) --- + def query_projection(self, inputs_q, out_sharding=None): + if not getattr(self, "_use_clipped_linears", False): + return super().query_projection(inputs_q, out_sharding=out_sharding) + x = _clip_in(inputs_q, self.q_clip) + y = self.query(x, out_sharding=out_sharding) + return _clip_out(y, self.q_clip) + + def kv_projection(self, inputs_kv, proj_name, out_sharding=None): + if not getattr(self, "_use_clipped_linears", False): + return super().kv_projection(inputs_kv, proj_name=proj_name, out_sharding=out_sharding) + if proj_name == "key": + cb, module = self.k_clip, self.key + elif proj_name == "value": + cb, module = self.v_clip, self.value + else: + raise ValueError(f"proj_name must be 'key' or 'value', but got {proj_name}") + x = _clip_in(inputs_kv, cb) + y = module(x, out_sharding=out_sharding) + return _clip_out(y, cb) + + def out_projection(self, out, out_sharding=None): + if not getattr(self, "_use_clipped_linears", False): + return super().out_projection(out, out_sharding=out_sharding) + x = _clip_in(out, self.o_clip) + y = self.out(x, out_sharding=out_sharding) + return _clip_out(y, self.o_clip) + + +class Gemma4ClippedMlpBlock(linears.MlpBlock): + """MlpBlock that applies the Gemma-4 vision per-projection activation clip bounds + to the gate (wi_0), up (wi_1) and down (wo) projections. + + Only the non-fused activation path is supported (E2B/E4B use + ``activations=("gelu", "linear")`` with ``fused_mlp=False``), because gate and up + carry distinct clip bounds. When ``use_clipped_linears`` is off, this delegates to + the base ``MlpBlock`` unchanged. + """ + + def __init__(self, *args, use_clipped_linears=False, **kwargs): + super().__init__(*args, **kwargs) + self._use_clipped_linears = bool(use_clipped_linears) + if self._use_clipped_linears: + self.gate_clip = _make_clip_state() # wi_0 + self.up_clip = _make_clip_state() # wi_1 + self.down_clip = _make_clip_state() # wo + + def validate_clip_bounds(self): + if not self._use_clipped_linears: + return + validate_clip_bounds(self.gate_clip, "gate_proj") + validate_clip_bounds(self.up_clip, "up_proj") + validate_clip_bounds(self.down_clip, "down_proj") + + def __call__(self, inputs, decode=False, deterministic=False, + intermediate_sharding=None, out_sharding=None): + if not self._use_clipped_linears: + return super().__call__(inputs, decode=decode, deterministic=deterministic, + intermediate_sharding=intermediate_sharding, out_sharding=out_sharding) + cfg = self.config + if getattr(cfg, "fused_mlp", False): + # Clipped vision MLP requires the unfused path so gate/up get their own output clamps. + raise ValueError("Gemma4ClippedMlpBlock requires fused_mlp=False (per-projection clip bounds).") + if self.mlp_layer_norm is not None: + inputs = self.mlp_layer_norm(inputs) + clips = [self.gate_clip, self.up_clip] # order matches activations ("gelu", "linear") == (gate, up) + activations = [] + for idx, act_fn in enumerate(self.activations): + dense_name = "wi" if len(self.activations) == 1 else f"wi_{idx}" + module = getattr(self, dense_name) + x = _clip_in(inputs, clips[idx]) + x = module(x, out_sharding=intermediate_sharding) + x = _clip_out(x, clips[idx]) + x = linears.checkpoint_name(x, "mlp" + dense_name) + if cfg.activations_in_float32: + x = x.astype(jnp.float32) + x = linears._convert_to_activation_function(act_fn)(x) + activations.append(x) + x = functools.reduce(operator.mul, activations).astype(self.dtype) + x = self.dropout(x, deterministic=deterministic) + x = self._maybe_shard_with_logical(x, self.intermediate_logical) + x = _clip_in(x, self.down_clip) + output = self.wo(x, out_sharding=out_sharding) + output = _clip_out(output, self.down_clip) + output = linears.checkpoint_name(output, "mlpwo") + return output + + class Gemma4EncoderBlock(nnx.Module): """Single transformer encoder block (MHSA + MLP).""" @@ -487,6 +839,11 @@ def __init__( is_vision=True, rngs=self.rngs, ) + # Opt-in Gemma-4 vision clipped-linears: attach the q/k/v/o checkpoint-resident + # clip bounds and route the projections through the clamp-in/clamp-out overrides. + self._use_clipped_linears = bool(getattr(config, "use_clipped_linears_for_vit", False)) + if self._use_clipped_linears: + self.attention.enable_vision_clip_bounds() self.pre_ffw_norm = normalizations.RMSNorm( num_features=config.hidden_size_for_vit, @@ -506,7 +863,9 @@ def __init__( rngs=self.rngs, ) - self.mlp = linears.MlpBlock( + mlp_cls = Gemma4ClippedMlpBlock if self._use_clipped_linears else linears.MlpBlock + mlp_kwargs = {"use_clipped_linears": True} if self._use_clipped_linears else {} + self.mlp = mlp_cls( config=config, mesh=mesh, in_features=config.hidden_size_for_vit, @@ -516,13 +875,32 @@ def __init__( weight_dtype=config.weight_dtype, intermediate_dropout_rate=config.dropout_rate, rngs=self.rngs, + **mlp_kwargs, ) - def __call__(self, x: jax.Array, positions: jax.Array | None = None, deterministic: bool = False) -> jax.Array: - """Applies the encoder block (MHSA + MLP) to the inputs.""" + def __call__( + self, + x: jax.Array, + positions: jax.Array | None = None, + deterministic: bool = False, + decoder_segment_ids: jax.Array | None = None, + ) -> jax.Array: + """Applies the encoder block (MHSA + MLP) to the inputs. + + When ``decoder_segment_ids`` is provided, patches carrying distinct segment + ids cannot attend to each other. This is used by the padded-patch path + (valid patches = segment 1, padded/sentinel patches = segment 2) so that the + phantom pad patches are masked out of vision self-attention. + """ x_normed = self.pre_attention_norm(x) - # Pass positions to attention for RoPE - x_attn, _ = self.attention(x_normed, x_normed, inputs_positions=positions, deterministic=deterministic) + # Pass positions to attention for RoPE (+ optional segment mask for padded patches). + x_attn, _ = self.attention( + x_normed, + x_normed, + inputs_positions=positions, + decoder_segment_ids=decoder_segment_ids, + deterministic=deterministic, + ) x_attn = self.post_attention_norm(x_attn) x_after_attn = x_attn + x @@ -573,34 +951,127 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs): nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,) ) - def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array: - """Applies the vision encoder layer.""" - if inputs.ndim == 4: - inputs = jnp.expand_dims(inputs, 1) - b, n, h, w, c = inputs.shape - inputs_flat = jnp.reshape(inputs, (b * n, h, w, c)) + def __call__( + self, + inputs: jax.Array, + deterministic: bool = False, + image_position_ids: jax.Array | None = None, + ): + """Applies the vision encoder layer. + + Two contracts: + + (A) Legacy all-valid (``image_position_ids is None``): ``inputs`` are raw images + [B, N, H, W, C] (or [B, H, W, C]); patchify -> full unmasked attention -> pool by the + derived positions -> return embeddings only (4D array [B, N, K, D]). + + (B) Padded-patch dynamic-N (``image_position_ids is not None``): ``inputs`` are ALREADY + patchified pixel_values with shape [B, L, P*P*C] (or [B, N, L, P*P*C]) and + ``image_position_ids`` is [B, L, 2] (or [B, N, L, 2]) with -1 sentinel rows marking padded + patches. The pre-patchified patches + REAL positions are fed to VisionEntry, per-patch + ``decoder_segment_ids`` (valid=1, pad=2) mask the phantom pad patches out of self-attention, + pooling uses the real positions (``avg_pool_by_positions`` maps a -1 patch to a zero-weight + bucket), and the VisionExit validity mask is returned. Returns a 2-tuple + ``(embeddings[B, N, K, D], image_masks[B*N, K])`` where ``image_masks.sum()`` is the number + of valid pooled tokens, threaded to ``merge_mm_embeddings.token_masks`` so exactly the valid + pooled tokens land in the image placeholders. + """ + if image_position_ids is None: + # ---- Legacy path: raw images -> patchify -> full (unmasked) attention ---- + # Fail-closed contract guard: the legacy path requires FULL images ([B,N,H,W,C] or [B,H,W,C]). + # Pre-patchified inputs ([B,L,F] or [B,N,L,F]) only make sense with per-patch image_position_ids + # (the padded-patch path). Refuse to silently mis-handle pre-patchified input as a full image, which + # would either crash cryptically or produce wrong pooling. + if inputs.ndim not in (4, 5): + raise ValueError( + "Gemma4 vision legacy path expects full images with shape [B, H, W, C] or [B, N, H, W, C]; " + f"got inputs.ndim={inputs.ndim} (shape {tuple(inputs.shape)}). Pre-patchified pixel_values must be " + "accompanied by image_position_ids (the padded-patch path). Refusing to silently take the legacy " + "all-valid vision path." + ) + if inputs.ndim == 4: + inputs = jnp.expand_dims(inputs, 1) + b, n, h, w, c = inputs.shape + inputs_flat = jnp.reshape(inputs, (b * n, h, w, c)) - x, positions_xy = self.vision_entry(inputs_flat) + x, positions_xy = self.vision_entry(inputs_flat) + + for i in range(self.config.num_hidden_layers_for_vit): + layer = getattr(self, f"layer_{i}") + x = layer(x, positions=positions_xy, deterministic=deterministic) + + vision_exit_results = self.vision_exit(x, positions_xy=positions_xy) + (embeddings, _) = vision_exit_results[0] + + embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype( + embeddings.dtype + ) + + # Unflatten batch and num_images + final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2])) + return final_x + + # ---- Padded-patch dynamic-N path: pre-patchified patches + sentinel positions ---- + # inputs: [B, L, P*P*C] pre-patchified pixel_values. Support an optional per-image N dim + # [B, N, L, F] -> flatten to [B*N, L, F] to match positions. + if inputs.ndim == 4: + b, n, l, f = inputs.shape + patches = jnp.reshape(inputs, (b * n, l, f)) + pos = jnp.reshape(image_position_ids, (b * n, l, 2)) + else: + assert inputs.ndim == 3, f"padded-patch path expects pre-patchified [B, L, F] patches, got {inputs.shape}" + b, l, f = inputs.shape + n = 1 + patches = inputs + pos = image_position_ids + if pos.ndim == 2: + pos = jnp.broadcast_to(pos, (b, l, 2)) + + pos = pos.astype(jnp.int32) + + # VisionEntry consumes pre-patchified patches + REAL positions (incl -1 sentinels). + x, positions_xy = self.vision_entry(patches, positions_xy=pos) + + # Segment ids: valid patch (any coord != -1) -> 1, padded sentinel patch -> 2. Distinct segments + # cannot attend to each other, masking the phantom pad patches out of self-attention. + is_pad = (positions_xy == -1).all(axis=-1) # [B*N, L] + decoder_segment_ids = jnp.where(is_pad, 2, 1).astype(jnp.int32) + + # Sanitize padded-patch rows to zero BEFORE attention. Segment masking already prevents pad patches + # from influencing valid patches' attention outputs, and pooling excludes pad positions — so for FINITE + # pad content the valid outputs are provably invariant (verified byte-identical under 1e6 garbage pad). + # However, attention computes q/k/v on ALL rows before masking, so a NON-FINITE (NaN/Inf) pad patch would + # produce NaN q/k/v and leak through the softmax normalization into valid rows (NaN*0=NaN). Zeroing pad + # rows here makes the isolation hold for arbitrary (incl. non-finite) pad content without changing any + # valid output. This is defense-in-depth for malformed processor output; the frozen processor emits finite + # pad, which was already invariant. + x = jnp.where(is_pad[..., None], jnp.zeros_like(x), x) for i in range(self.config.num_hidden_layers_for_vit): layer = getattr(self, f"layer_{i}") - x = layer(x, positions=positions_xy, deterministic=deterministic) + x = layer( + x, + positions=positions_xy, + deterministic=deterministic, + decoder_segment_ids=decoder_segment_ids, + ) + # Pool with REAL positions; avg_pool_by_positions returns (embeddings, validity_mask). vision_exit_results = self.vision_exit(x, positions_xy=positions_xy) - - # Return embeddings from VisionExit tuple - # vision_exit_results is a tuple of (embeddings, mask) tuples, one for each output length. - # We take the first result. - (embeddings, _) = vision_exit_results[0] + (embeddings, image_masks) = vision_exit_results[0] # embeddings [B*N, K, D], mask [B*N, K] embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype( embeddings.dtype ) - # Unflatten batch and num_images final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2])) + if image_masks is None: + image_masks = jnp.ones((b * n, embeddings.shape[1]), dtype=jnp.int32) + else: + # merge_mm_embeddings does argsort(-token_mask); use int32 so negation/sort is well-defined. + image_masks = image_masks.astype(jnp.int32) - return final_x + return final_x, image_masks class Gemma4VisionProjector(nnx.Module): diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index a70ad780a3..aab8f1bdab 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -129,6 +129,7 @@ def __call__( decoder_segment_ids=None, encoder_images: None | jnp.ndarray = None, encoder_image_masks: None | jnp.ndarray = None, + encoder_image_position_ids: None | jnp.ndarray = None, encoder_videos: None | jnp.ndarray = None, encoder_video_masks: None | jnp.ndarray = None, encoder_video_grid_thw: None | jnp.ndarray = None, @@ -164,17 +165,19 @@ def __call__( video_embeddings = None audio_embeddings = None deepstack_visual_embeds = None + vision_image_masks = None if getattr(self.config, "use_multimodal", False) and encoder_images is not None: - image_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] - input_images=encoder_images, deterministic=not enable_dropout + image_embeddings, deepstack_visual_embeds, vision_image_masks = self.vision_encoder( # pyrefly: ignore[not-callable] + input_images=encoder_images, deterministic=not enable_dropout, + image_position_ids=encoder_image_position_ids, ) bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision( self.config, decoder_input_tokens, is_video=False ) if getattr(self.config, "use_multimodal", False) and encoder_videos is not None: - video_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] + video_embeddings, deepstack_visual_embeds, _ = self.vision_encoder( # pyrefly: ignore[not-callable] input_images=encoder_videos, input_masks=encoder_video_masks, video_grid_thw=encoder_video_grid_thw, @@ -197,7 +200,7 @@ def __call__( if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None: multimodal_input = MultimodalInput( image_embeddings=image_embeddings, - image_masks=encoder_image_masks, + image_masks=vision_image_masks if vision_image_masks is not None else encoder_image_masks, video_embeddings=video_embeddings, video_masks=encoder_video_masks, audio_embeddings=audio_embeddings, @@ -449,6 +452,7 @@ def __call__( cache=None, encoder_images: jax.Array | None = None, encoder_image_masks: jax.Array | None = None, + encoder_image_position_ids: jax.Array | None = None, encoder_videos: jax.Array | None = None, encoder_video_masks: jax.Array | None = None, encoder_video_grid_thw: jax.Array | None = None, @@ -499,16 +503,18 @@ def __call__( video_embeddings = None audio_embeddings = None deepstack_visual_embeds = None + vision_image_masks = None if getattr(self.config, "use_multimodal", False) and encoder_images is not None: - image_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] - input_images=encoder_images, deterministic=not enable_dropout + image_embeddings, deepstack_visual_embeds, vision_image_masks = self.vision_encoder( # pyrefly: ignore[not-callable] + input_images=encoder_images, deterministic=not enable_dropout, + image_position_ids=encoder_image_position_ids, ) bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision( self.config, decoder_input_tokens, is_video=False ) if getattr(self.config, "use_multimodal", False) and encoder_videos is not None: - video_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] + video_embeddings, deepstack_visual_embeds, _ = self.vision_encoder( # pyrefly: ignore[not-callable] input_images=encoder_videos, input_masks=encoder_video_masks, video_grid_thw=encoder_video_grid_thw, @@ -531,7 +537,7 @@ def __call__( if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None: multimodal_input = MultimodalInput( image_embeddings=image_embeddings, - image_masks=encoder_image_masks, + image_masks=vision_image_masks if vision_image_masks is not None else encoder_image_masks, video_embeddings=video_embeddings, video_masks=encoder_video_masks, audio_embeddings=audio_embeddings, diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py index 67e1f589ca..64c2369d17 100644 --- a/src/maxtext/optimizers/optimizers.py +++ b/src/maxtext/optimizers/optimizers.py @@ -239,14 +239,32 @@ def get_optimizer(config, learning_rate_schedule, model=None): # When trainable_parameters_mask is empty, freeze_mask_fn is None and all parameters are trained. trainable_patterns = getattr(config, "trainable_parameters_mask", None) freeze_mask_fn = _get_path_mask_fn(trainable_patterns, match_returns_true=False) - if freeze_mask_fn is not None: - # Use optax.multi_transform to explicitly map frozen parameters to a stateless set_to_zero() optimizer. - # If we simply wrapped base_opt in optax.masked() or chained it, Optax would still allocate - # massive states (momentum, variance) for the entire model before zeroing the updates. - # By using multi_transform, only the trainable parameters get states allocated. + + # Gemma-4 vision clipped-linears: the 448 clip bounds are checkpoint-resident, immutable state and must + # NEVER receive optimizer updates, weight decay, or momentum/variance slots. Freeze them via a leaf-PATH + # mask (robust to the nnx->linen bridge's type erasure). This composes with any trainable_parameters_mask: + # a parameter is "trainable" only if it is trainable under the whitelist AND is not a clip bound. + freeze_clip_bounds = bool(getattr(config, "use_clipped_linears_for_vit", False)) + + if freeze_mask_fn is not None or freeze_clip_bounds: + def _partition(params): + # trainable_mask: True where trainable under the whitelist (all-True if no whitelist). + if freeze_mask_fn is not None: + trainable_mask = freeze_mask_fn(params) + else: + trainable_mask = jax.tree_util.tree_map(lambda _: True, params) + if freeze_clip_bounds: + # clip_optimizer_freeze_mask returns True for trainable leaves, False for clip bounds. + from maxtext.models import gemma4_vision # pylint: disable=import-outside-toplevel + clip_trainable = gemma4_vision.clip_optimizer_freeze_mask(params) + trainable_mask = jax.tree_util.tree_map(lambda a, b: bool(a) and bool(b), trainable_mask, clip_trainable) + return jax.tree_util.tree_map(lambda x: "trainable" if x else "frozen", trainable_mask) + + # Use optax.multi_transform so frozen (and clip-bound) params map to a stateless set_to_zero() optimizer + # and never allocate momentum/variance slots. return optax.multi_transform( {"trainable": base_opt, "frozen": optax.set_to_zero()}, - lambda params: jax.tree_util.tree_map(lambda x: "frozen" if x else "trainable", freeze_mask_fn(params)), + _partition, ) return base_opt diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 0906c3e786..d191426072 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -148,6 +148,14 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr decoder_segment_ids=data["inputs_segmentation"], encoder_images=data["images"] if config.use_multimodal else None, encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + # Gemma-4 padded-patch (Option-S) contract: when the data pipeline produces per-patch image + # position ids (native-resolution / pan-and-scan processors), thread them to the vision encoder + # so pad patches are masked and pooled by real coordinates. The stock fixed-grid Gemma-4 processor + # produces all-valid patches and does NOT emit this key, so the legacy all-valid path is used + # (correct for its contract). See gemma4_vision.Gemma4VisionEncoderLayer for the fail-closed guard. + encoder_image_position_ids=( + data["image_position_ids"] if config.use_multimodal and "image_position_ids" in data else None + ), enable_dropout=config.enable_dropout if is_train else False, rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] mutable=mutable_collections, @@ -197,6 +205,11 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr decoder_segment_ids=data["inputs_segmentation"], encoder_images=data["images"] if config.use_multimodal else None, encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + # See the Linen branch above: thread per-patch image position ids when the data pipeline emits + # them (Gemma-4 Option-S padded-patch contract); None for the stock fixed-grid all-valid processor. + encoder_image_position_ids=( + data["image_position_ids"] if config.use_multimodal and "image_position_ids" in data else None + ), enable_dropout=config.enable_dropout if is_train else False, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index a2a1403125..9e4bb54dce 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -1211,6 +1211,14 @@ def _walk_align(ckpt, model_arr, axes): "Please ensure the checkpoint format matches the scan_layers setting." ) + # Gemma-4 vision clipped-linears: fail-closed post-load validation BEFORE the first JIT forward/train. + # Asserts exactly 112 clip-state modules / 448 finite, correctly-ordered scalar bounds. A missing/NaN/ + # unordered bound or wrong module count raises here rather than silently degrading to an identity clamp. + if bool(getattr(config, "use_clipped_linears_for_vit", False)) and bool(getattr(config, "use_multimodal", False)): + from maxtext.models import gemma4_vision # pylint: disable=import-outside-toplevel + n_proj, n_bounds = gemma4_vision.validate_all_vision_clip_bounds(model) + max_logging.log(f"Gemma-4 vision clip validation OK: {n_proj} clip-state modules, {n_bounds} scalar bounds.") + if wrap_with_tunix_adapter: with mesh: use_no_op_mappings = "maxtext_config" in config.vllm_additional_config diff --git a/tests/unit/gemma4_clip_linears_test.py b/tests/unit/gemma4_clip_linears_test.py new file mode 100644 index 0000000000..fd2c08f766 --- /dev/null +++ b/tests/unit/gemma4_clip_linears_test.py @@ -0,0 +1,182 @@ +# 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. + +"""In-tree tests for the Gemma-4 vision clipped-linears (PR #4790). + +Pure mechanism + validation tests: no checkpoint, no TPU. Cover the merge-critical invariants: +clip math vs jnp.clip, dtype/no-op, NaN/Inf/ordering hard-fail, exact 112/448 counts, optimizer +freeze mask polarity, and fail-closed contract guards. +""" + +import unittest + +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.models import gemma4_vision as gv + +jax.config.update("jax_platform_name", "cpu") + + +def _finite_cb(in_lo=-2.5, in_hi=2.5, out_lo=-4.0, out_hi=4.0): + cb = gv._make_clip_state() + cb.input_min.value = jnp.asarray(in_lo, jnp.float32) + cb.input_max.value = jnp.asarray(in_hi, jnp.float32) + cb.output_min.value = jnp.asarray(out_lo, jnp.float32) + cb.output_max.value = jnp.asarray(out_hi, jnp.float32) + return cb + + +class ClipMathTest(unittest.TestCase): + + def test_clip_in_matches_jnp_clip(self): + cb = _finite_cb() + x = (jax.random.normal(jax.random.PRNGKey(0), (5, 17)) * 6.0).astype(jnp.float32) + self.assertTrue(bool(jnp.array_equal(gv._clip_in(x, cb), jnp.clip(x, -2.5, 2.5)))) + + def test_clip_out_matches_jnp_clip(self): + cb = _finite_cb() + x = (jax.random.normal(jax.random.PRNGKey(1), (5, 17)) * 6.0).astype(jnp.float32) + self.assertTrue(bool(jnp.array_equal(gv._clip_out(x, cb), jnp.clip(x, -4.0, 4.0)))) + + def test_none_is_exact_noop(self): + x = jax.random.normal(jax.random.PRNGKey(2), (3, 8)).astype(jnp.float32) + self.assertTrue(bool(jnp.array_equal(gv._clip_in(x, None), x))) + self.assertTrue(bool(jnp.array_equal(gv._clip_out(x, None), x))) + + def test_dtype_preserved(self): + cb = _finite_cb() + x = jax.random.normal(jax.random.PRNGKey(3), (4, 4)).astype(jnp.bfloat16) + self.assertEqual(gv._clip_in(x, cb).dtype, jnp.bfloat16) + + def test_bounds_are_stop_gradient(self): + # The clamp bound is read through jax.lax.stop_gradient at the use-site, so gradient wrt the bound + # value is exactly zero even when the clamp is active. Probe _clip_bound_value directly with a + # lightweight bound holder to avoid mutating an nnx Param inside a traced function. + class _B: + def __init__(self, v): + self.value = v + + def loss(bound_val): + x = jnp.ones((4,), jnp.float32) * 10.0 # above the bound -> clamp active + hi = gv._clip_bound_value(_B(bound_val), jnp.float32) + return jnp.sum(jnp.clip(x, -100.0, hi)) + + g = jax.grad(loss)(jnp.asarray(2.5, jnp.float32)) + self.assertEqual(float(g), 0.0) + + +class ValidateBoundsTest(unittest.TestCase): + + def test_accepts_finite_ordered(self): + gv.validate_clip_bounds(_finite_cb(), "ok") # no raise + + def test_rejects_nan_sentinel(self): + with self.assertRaises(ValueError): + gv.validate_clip_bounds(gv._make_clip_state(), "nan") + + def test_rejects_inf(self): + cb = _finite_cb() + cb.input_max.value = jnp.asarray(np.inf, jnp.float32) + with self.assertRaises(ValueError): + gv.validate_clip_bounds(cb, "inf") + + def test_rejects_min_gt_max(self): + cb = _finite_cb(in_lo=3.0, in_hi=-3.0) # input_min > input_max + with self.assertRaises(ValueError): + gv.validate_clip_bounds(cb, "order") + + def test_none_is_noop(self): + self.assertIsNone(gv.validate_clip_bounds(None)) + + +class FreezeMaskTest(unittest.TestCase): + + def test_freeze_mask_polarity_and_count(self): + tree = { + "vision_encoder": { + "layer_0": { + "attention": { + "q_clip": {"input_min": 1.0, "input_max": 1.0, "output_min": 1.0, "output_max": 1.0}, + "query": {"kernel": 1.0}, + }, + "mlp": {"gate_clip": {"input_min": 1.0, "output_max": 1.0}, "wi_0": {"kernel": 1.0}}, + } + }, + "decoder": {"layer_0": {"norm": {"scale": 1.0}}}, + } + mask = gv.clip_optimizer_freeze_mask(tree) + flat = jax.tree_util.tree_flatten_with_path(mask)[0] + n_clip = 0 + for path, v in flat: + is_clip = gv._is_clip_bound_path(path) + # clip leaves -> False (frozen); everything else -> True (trainable) + self.assertEqual(bool(v), (not is_clip)) + if is_clip: + n_clip += 1 + self.assertEqual(n_clip, 6) # q_clip(4) + gate_clip(2) + + +class SymbolsTest(unittest.TestCase): + + def test_expected_counts_constants(self): + self.assertEqual(gv.EXPECTED_CLIP_PROJECTIONS, 112) + self.assertEqual(gv.EXPECTED_CLIP_BOUNDS, 448) + + def test_public_symbols_present(self): + for sym in ("validate_all_vision_clip_bounds", "clip_optimizer_freeze_mask", + "Gemma4ClippedMlpBlock", "Gemma4Attention"): + self.assertTrue(hasattr(gv, sym), sym) + + +if __name__ == "__main__": + unittest.main() + + +class OptionSInvariantsTest(unittest.TestCase): + """M7: Option-S sentinel integrity + placeholder==pooled count invariants (pure validators).""" + + def test_valid_positions_pass(self): + pos = np.array([[[0, 0], [1, 0], [-1, -1], [-1, -1]]], dtype=np.int32) # 2 valid + 2 full-sentinel + ok, n_bad = gv.option_s_position_sentinel_ok(pos) + self.assertTrue(ok) + self.assertEqual(n_bad, 0) + gv.validate_option_s_positions(pos) # must not raise + + def test_mixed_sentinel_row_fails(self): + for bad_row in ([-1, 5], [5, -1]): + pos = np.array([[[0, 0], bad_row, [-1, -1]]], dtype=np.int32) + ok, n_bad = gv.option_s_position_sentinel_ok(pos) + self.assertFalse(ok, f"mixed row {bad_row} should be flagged") + self.assertEqual(n_bad, 1) + with self.assertRaises(ValueError): + gv.validate_option_s_positions(pos, where="test") + + def test_other_negative_coord_fails(self): + pos = np.array([[[0, 0], [-2, -2], [-1, -1]]], dtype=np.int32) # -2 is not the sentinel + ok, n_bad = gv.option_s_position_sentinel_ok(pos) + self.assertFalse(ok) + + def test_pooled_matches_placeholder(self): + mask = np.array([True, True, True, False, False]) # 3 valid pooled tokens + ok, n_valid = gv.option_s_pooled_matches_placeholder(mask, 3) + self.assertTrue(ok) + self.assertEqual(n_valid, 3) + + def test_pooled_placeholder_mismatch_detected(self): + mask = np.array([True, True, True, False, False]) # 3 valid + ok, n_valid = gv.option_s_pooled_matches_placeholder(mask, 4) # placeholder count 4 != 3 + self.assertFalse(ok) + self.assertEqual(n_valid, 3) diff --git a/tests/unit/gemma4_mm_contract_gate_test.py b/tests/unit/gemma4_mm_contract_gate_test.py new file mode 100644 index 0000000000..a60529079b --- /dev/null +++ b/tests/unit/gemma4_mm_contract_gate_test.py @@ -0,0 +1,88 @@ +# 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. +"""M2/M5/M6: Gemma-4 E2B/E4B multimodal STATIC contract gate (configs/types.py). + +Asserts the config validator ACCEPTS a valid E2B multimodal config and HARD-FAILS on each +contract violation: unknown PLE mode (M2 enum), bidirectional image attention (causal M6), +fused_qkv / fused_mlp (M5), and clipped-linears disabled. These encode the semantics verified +against pinned HF Transformers 5.9.0 and must fail closed rather than silently degrade. +""" +import os +import unittest + +from maxtext.configs import pyconfig +from maxtext.utils.globals import MAXTEXT_PKG_DIR + + +def _init(overrides): + base_yml = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml") + argv = [ + "test", + base_yml, + "model_name=gemma4-e2b", + "use_multimodal=true", + "override_model_config=true", + "tokenizer_path=/tmp/unused", + "per_device_batch_size=1", + "scan_layers=false", + "enable_checkpointing=false", + "run_name=gate_test", + "steps=1", + "skip_jax_distributed_system=true", + ] + overrides + return pyconfig.initialize(argv) + + +class Gemma4MmContractGateTest(unittest.TestCase): + + def test_valid_e2b_mm_config_builds(self): + # Clipped linears on, identity PLE, causal image spans, fused off -> the HF-faithful contract. + cfg = _init(["use_clipped_linears_for_vit=true"]) + self.assertEqual(cfg.model_name, "gemma4-e2b") + self.assertTrue(cfg.use_multimodal) + self.assertEqual(str(cfg.ple_pad_mode), "identity") + + def test_unknown_ple_pad_mode_hard_fails(self): + with self.assertRaises(Exception): + _init(["use_clipped_linears_for_vit=true", "ple_pad_mode=bogus"]) + + def test_both_ple_pad_mode_is_valid_ablation_enum(self): + cfg = _init(["use_clipped_linears_for_vit=true", "ple_pad_mode=both"]) + self.assertEqual(str(cfg.ple_pad_mode), "both") + + def test_bidirectional_image_attn_hard_fails(self): + # Gemma-4 E2B/E4B image spans are causal; bidirectional is a Gemma-3 / 26B / 31B feature. + with self.assertRaises(Exception): + _init(["use_clipped_linears_for_vit=true", "use_bidirectional_image_attn=true"]) + + def test_fused_qkv_hard_fails(self): + with self.assertRaises(Exception): + _init(["use_clipped_linears_for_vit=true", "fused_qkv=true"]) + + def test_fused_mlp_hard_fails(self): + with self.assertRaises(Exception): + _init(["use_clipped_linears_for_vit=true", "fused_mlp=true"]) + + def test_clipped_linears_off_under_mm_hard_fails(self): + with self.assertRaises(Exception): + _init(["use_clipped_linears_for_vit=false"]) + + def test_packing_hard_fails(self): + # Packed multimodal is unsupported for E2B/E4B in any mode (cross-doc image attention risk). + with self.assertRaises(Exception): + _init(["use_clipped_linears_for_vit=true", "packing=true"]) + + +if __name__ == "__main__": + unittest.main()