Gemma-4 E2B/E4B: enable image (multimodal) parity — vision clipped-linears + padded-patch masking + PLE fix

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 <proj>_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}
 

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"

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

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

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",

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):

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