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..c4a4b26c26 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,10 @@ 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 diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml index b8da0d1d46..cd752ae22b 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,10 @@ 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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a548422..712c042c85 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2137,6 +2137,43 @@ 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' (token-id path only) " + "or 'both' (also substitute the pad embedding in the context path)." + ), + ) + 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 +3609,22 @@ 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. - 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.") + # 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." + ) 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/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..51b32ea112 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -2231,10 +2231,31 @@ def _apply_gemma4_small_layers( """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)) + if str(getattr(cfg, "ple_pad_mode", "identity")) == "both" and hasattr(self, "shared_embedding"): + _pad_vec = self.shared_embedding(jnp.full_like(decoder_input_tokens, _pad_id).astype(jnp.int32)) + 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..60ab571f85 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,108 @@ 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_in(x, cb): + """clamp(x, input_min, input_max) in x's dtype; no-op if ``cb`` is None.""" + if cb is None: + return x + xd = x.dtype + return jnp.clip(x, cb.input_min.value.astype(xd), cb.input_max.value.astype(xd)) + + +def _clip_out(y, cb): + """clamp(y, output_min, output_max) in y's dtype; no-op if ``cb`` is None.""" + if cb is None: + return y + yd = y.dtype + return jnp.clip(y, cb.output_min.value.astype(yd), cb.output_max.value.astype(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. Raises ValueError otherwise. No-op if ``cb`` is None.""" + if cb is None: + return + 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.") + + + + + def factorized_posemb(posemb: jax.Array, positions_xy: jax.Array, precision) -> jax.Array: """Computes factorized position embedding from (x, y) coordinates. @@ -409,7 +513,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 +530,116 @@ 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 + 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 +709,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 +733,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 +745,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 +821,106 @@ 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.""" + 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 ---- + 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) + + 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: - inputs = jnp.expand_dims(inputs, 1) - b, n, h, w, c = inputs.shape - inputs_flat = jnp.reshape(inputs, (b * n, h, w, c)) + 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)) - x, positions_xy = self.vision_entry(inputs_flat) + 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) 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,