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

[lokic233]

What

Enables Gemma-4 E2B/E4B image (multimodal) parity in MaxText. Previously E2B/E4B multimodal was gated off — the validator noted it was "pending clipped-linears in the vision encoder" — and produced diverging image logits. This PR lands the full set of pieces the HF reference contract needs, validated end-to-end.

To enable image inputs for these models you now set use_multimodal=true and use_clipped_linears_for_vit=true; the E2B/E4B model configs supply the rest of the image contract. Everything is off by default, so other models and text-only runs are unchanged.

Changes

1. Vision clipped-linears (models/gemma4_vision.py, checkpoint_conversion/utils/param_mapping.py, configs/*)
The Gemma-4 reference checkpoint ships per-projection activation clip bounds (7 vision projections × {input,output}_{min,max} × 16 blocks = 448 bounds). Added as opt-in, checkpoint-resident, non-trainable scalars gated on use_clipped_linears_for_vit (exact no-op when off); Gemma4Attention and a Gemma4ClippedMlpBlock apply clamp-in → projection → clamp-out. NaN-sentinel validation; path-based optimizer-freeze mask. The converter maps the 448 bounds.

2. Image-parity path (models/gemma4_vision.py, layers/encoders.py, models/models.py, layers/nnx_decoders.py)
Clip-bounds alone are necessary but not sufficient; this adds the rest of the reference contract:

• Padded-patch vision masking: Gemma4EncoderBlock threads decoder_segment_ids (valid=1/pad=2) so phantom padded patches are masked out of vision self-attention; Gemma4VisionEncoderLayer gains a padded-patch path (pre-patchified patches + real per-patch positions with -1 sentinels), pools by real positions, and returns the pooled-token validity mask.
• Position/mask threading: encoder_image_position_ids flows through the model forward → vision encoder; the returned validity mask is routed into MultimodalInput.image_masks so exactly the valid pooled tokens land in the image placeholders.
• Decoder PLE image-row substitution: for E2B/E4B the per-layer embeddings map image placeholder tokens → pad_token_id (matching HF modeling_gemma4), instead of feeding the placeholder id into the PLE path. Without this the per-layer embeddings at image positions diverge, corrupting the image-span/post-image logits.
• Causal image spans: E2B/E4B image attention is causal (not bidirectional).

3. Config (configs/types.py, configs/base.yml, configs/models/gemma4-e2b.yml, configs/models/gemma4-e4b.yml)
Adds use_clipped_linears_for_vit, use_bidirectional_image_attn, ple_pad_substitute_image_rows, ple_pad_mode, image_placeholder_token_id, ple_pad_token_id (defaults preserve behavior for other models); the E2B/E4B model configs set the image-contract defaults; the validator allows E2B/E4B multimodal when use_clipped_linears_for_vit is set.

Validation (E2E, honest)

CPU, teacher-forced 340-token forward on a converted E2B checkpoint vs the HF reference logits. Frozen gate: post-image max_KL ≤ 1.26e-3 and argmax ≥ 0.995.

segment	clip-ON (full fix)	clip-OFF (bounds neutralized)
pre_image	max_KL 2.7e-5, argmax 1.0	2.7e-5, 1.0
image_span	max_KL 9.0e-5, argmax 1.0	5.14, 0.60
post_image	max_KL 4.3e-5, argmax 1.0 ✅	1.195, 0.84 ❌

• clip-ON with the full fix passes the gate (post-image max_KL 4.3e-5, argmax 1.0), matching an independently-measured reference baseline (~4.2e-5) for the same fixture.
• clip-OFF diverges (post-image max_KL 1.195), confirming the clip-bounds are required.
• Verified via the clean runtime path (only use_multimodal=true + use_clipped_linears_for_vit=true, model config supplies the rest) — same result.

Additional unit/component checks: clip math bit-exact vs jnp.clip, dtype-preserving, exact no-op when disabled; NaN/Inf bounds hard-fail; the vision padded-patch path yields exactly the expected valid pooled-token count.

[gemini-code-assist]

Code Review

This pull request introduces support for Gemma-4 E2B/E4B multimodal models by implementing opt-in, checkpoint-resident, non-trainable activation clip bounds (clipped-linears) for the vision encoder. The changes include adding configuration options, parameter mapping, and custom projection overrides in Gemma4Attention and Gemma4ClippedMlpBlock to apply the clipping. The review feedback highlights critical integration gaps: the validation function validate_clip_bounds and the optimizer freeze mask clip_optimizer_freeze_mask are defined but never actually called or integrated into the model initialization or optimizer setup. Additionally, the reviewer recommends raising an error if fused_qkv is enabled, as it would currently bypass the attention projection clipping.

[gemini-code-assist]

The validate_clip_bounds function is defined here, and corresponding validate_clip_bounds methods are implemented in Gemma4Attention and Gemma4ClippedMlpBlock. However, these validation methods are never actually called anywhere in the codebase (e.g., during model initialization, warmup, or after checkpoint loading). This means the NaN sentinel check is completely bypassed, and any missing or NaN bounds in the checkpoint will silently propagate into jnp.clip, turning the entire activation tensor into NaNs.

Please ensure that validate_clip_bounds is called on the model layers after checkpoint loading or during model initialization (outside of JIT-compiled functions). Additionally, we can simplify the extraction of the scalar value and the finiteness check using float(v) and math.isfinite.

Suggested change

	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 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(v)
	import math
	if not math.isfinite(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.")

[gemini-code-assist]

The clip_optimizer_freeze_mask function is defined to exclude the immutable clip bounds from optimizer updates and weight decay. However, this function is never imported or called in the optimizer setup (e.g., in train.py or optimizers.py). As a result, if freeze_vision_encoder_params is set to False (e.g., during full fine-tuning or pre-training), these supposedly non-trainable clip bounds will be updated by the optimizer, which is incorrect.

Please integrate clip_optimizer_freeze_mask into the optimizer creation flow to ensure the clip bounds are correctly masked out and frozen.

[gemini-code-assist]

If fused_qkv is enabled in the configuration, the query, key, and value projections are fused into a single qkv_proj matrix multiplication. In this case, Gemma4Attention does not override qkv_projection to apply the per-projection clip bounds, meaning the clipping will be silently bypassed for Q, K, and V. Just like Gemma4ClippedMlpBlock raises an error when fused_mlp is enabled, Gemma4Attention should raise an error if fused_qkv is enabled when clipped linears are used.

Suggested change

	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 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
	if getattr(self.config, "fused_qkv", False):
	raise ValueError("Gemma4Attention with clipped linears requires fused_qkv=False.")
	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()

[lokic233]

Converted to draft and narrowed the scope after end-to-end validation.

What changed vs the initial version of this PR:

• Removed the validator change that would have unblocked E2B/E4B multimodal. Clip-bounds are necessary but not sufficient for image parity, so unblocking would have implied working image support that isn't there yet.
• Retitled to reflect that this lands the clip-bounds building block only.

E2E findings (CPU, teacher-forced, vs HF reference logits):

• The 448 clip bounds convert + load correctly (real finite values), and the clamp math is exact — verified at the helper, forward-component, and checkpoint-load levels.
• With clip-bounds loaded, the text path is exact (pre-image max KL ≈ 2.7e-5) but the image span still diverges (post-image max KL ≈ 10, argmax ≈ 0.55 — far from parity). Clip-ON vs clip-OFF barely differ, i.e. the clip effect is swamped by a larger vision-forward divergence.
• Root cause is not the clip math: the current Gemma-4 vision forward is missing other pieces of the reference contract (pad-patch attention masking, external image-position threading, and likely more). Those are out of scope for this change and being worked separately.

This clip-bounds change is off by default and safe to land on its own; keeping it as draft until the remaining vision-forward pieces land, at which point E2B/E4B image parity can be demonstrated end-to-end.

[lokic233]

Update — this PR now closes E2B/E4B image parity end-to-end.

It initially landed only the vision clipped-linears (marked draft, since clip-bounds alone are necessary but not sufficient). After root-causing the remaining image-span divergence, the missing pieces were:

1. Padded-patch vision self-attention masking + real per-patch position threading (phantom pad patches must be masked; pooled by real positions to the valid tokens).
2. Decoder per-layer-embedding (PLE) image-row substitution — E2B/E4B map image placeholder tokens → pad_token_id in the PLE path (HF modeling_gemma4); feeding the placeholder id instead was the dominant residual.
3. Causal image spans for E2B/E4B.

With all three plus the clip-bounds, the teacher-forced forward now matches the HF reference: post-image max_KL 4.3e-5, argmax 1.0 (gate ≤ 1.26e-3 / ≥ 0.995), vs 1.195 with clip-bounds disabled. Marking ready for review. All behavior is opt-in and defaults keep other models unchanged.