Gemma-4 E2B: enable image (multimodal) E2E parity — vision clipped-linears + padded-patch masking + PLE substitution (E4B vision-tower validated; E4B E2E pending)

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

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)

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

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,

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