Gemma-4 E2B: enable image (multimodal) E2E parity — vision clipped-linears + padded-patch masking + PLE substitution (E4B vision-tower validated; E4B E2E pending) - #4790
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
| 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.") |
There was a problem hiding this comment.
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.
| 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.") |
There was a problem hiding this comment.
Fixed in the pushed head (35ce712). validate_all_vision_clip_bounds(model) now runs in model_creation_utils.from_pretrained AFTER the Orbax restore and BEFORE the first JIT, gated on use_clipped_linears_for_vit and use_multimodal. It walks the model with nnx.iter_graph, validates every _ClipBounds module (finite, scalar, min<=max), and asserts EXACTLY 112 modules / 448 scalar bounds — hard-failing on missing/NaN/Inf/wrong-count (incl. zero) rather than letting a NaN propagate into jnp.clip. Real-TPU: this validation PASSED (112/448) on TPU7x. Tests: tests/unit/gemma4_clip_linears_test.py (rejects NaN/Inf/min>max; expected-count constants).
| 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) |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Fixed in the pushed head (525d38a). clip_optimizer_freeze_mask is now wired into optimizers.get_optimizer: when use_clipped_linears_for_vit is set, the 448 clip-bound leaves are composed into the optax.multi_transform partition as set_to_zero() (no updates, no weight decay, no momentum/variance slots) — and this holds even when freeze_vision_encoder_params=False. Defense in depth: lax.stop_gradient is also applied at every clamp use-site (_clip_in/_clip_out). Unit tests assert the freeze-mask polarity/count and that the bounds receive zero gradient (bounds_are_stop_gradient).
| 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() |
There was a problem hiding this comment.
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.
| 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() |
There was a problem hiding this comment.
Fixed in the pushed head (35ce712). Gemma4Attention.enable_vision_clipped_linears now raises if fused_qkv=True (mirrors the existing fused_mlp guard). Also added a static config-validation gate in configs/types.py that hard-fails E2B/E4B multimodal with fused_qkv=True or fused_mlp=True before model construction. Config tests cover both (tests/unit/gemma4_mm_contract_gate_test.py: test_fused_qkv_hard_fails, test_fused_mlp_hard_fails).
…erequisite)
The Gemma-4 E2B/E4B reference checkpoint ships per-projection activation clip
bounds for the vision tower: each of the 7 vision projections
(self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj) in every encoder block
carries a scalar {input,output}_{min,max} (16 blocks x 7 x 4 = 448 bounds), and
the reference forward clamps each projection's input and output by those bounds.
MaxText does not model them today, so the values are silently dropped on
conversion. They are one of the pieces required before E2B/E4B image inputs can
match the reference.
This adds the clip bounds as an opt-in, checkpoint-resident, non-trainable
feature gated on use_clipped_linears_for_vit (exact no-op when False):
- gemma4_vision.py: clip-bound helpers (_clip_in/_clip_out, _ClipBounds,
NaN-sentinel validate_clip_bounds, path-based clip_optimizer_freeze_mask).
Gemma4Attention overrides its q/k/v/o projection methods to clamp-in ->
DenseGeneral -> clamp-out; Gemma4ClippedMlpBlock does the same for
gate/up/down. The shared attentions.Attention / linears.MlpBlock are untouched,
and the underlying DenseGeneral weights and checkpoint key paths are unchanged
(the bounds are separate scalar leaves).
- param_mapping.py: maps the 448 clip-bound scalars from the HF checkpoint when
the flag is set.
- types.py / base.yml: adds the use_clipped_linears_for_vit flag (default False).
The bounds are plain nnx.Param leaves so they round-trip through the
nnx->linen->orbax checkpoint path and map to the canonical params collection;
clip_optimizer_freeze_mask keeps them out of optimizer updates and weight decay
via a leaf-path mask. A NaN sentinel + validate_clip_bounds hard-fails on a
missing/non-finite bound rather than silently degrading to an identity clamp.
Scope / status (deliberately conservative):
- This does NOT unblock E2B/E4B multimodal. The existing validator that gates
E2B/E4B image inputs is left in place, because clip-bounds are necessary but
NOT sufficient for image parity on their own.
- Verified: the 448 bounds convert and load with real finite values; the clamp
math is bit-exact vs jnp.clip, dtype-preserving, and an exact no-op when
disabled (checked at the helper level and on a real Gemma4EncoderBlock forward,
where wide bounds reproduce the disabled-path output exactly and tight bounds
bound the projected activations).
- End-to-end teacher-forced image parity is NOT yet achieved: with clip-bounds
loaded, the text path is exact but the image span still diverges from the HF
reference, because the current Gemma-4 vision forward is missing other pieces
of the reference contract (e.g. pad-patch attention masking and external image
position threading). Those are out of scope for this change and tracked
separately; this PR lands the clip-bounds building block on its own.
7428140 to
dc7c56b
Compare
lokic233
changed the title
|
Converted to draft and narrowed the scope after end-to-end validation. What changed vs the initial version of this PR:
E2E findings (CPU, teacher-forced, vs HF reference logits):
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. |
…-row substitution) Builds on the vision clipped-linears to make Gemma-4 E2B/E4B image inputs match the HF reference end to end. The clipped-linears alone are necessary but not sufficient; the reference contract also needs the vision padded-patch handling and a decoder-side per-layer embedding (PLE) fix. Vision (models/gemma4_vision.py): - Gemma4EncoderBlock threads decoder_segment_ids into attention (valid=1 / pad=2) so the phantom padded patches are masked out of vision self-attention. - Gemma4VisionEncoderLayer gains a padded-patch path (image_position_ids != None): consume pre-patchified patches + real per-patch positions (-1 = pad), build the segment ids, pool by the real positions, and return the pooled-token validity mask. image_position_ids is None -> byte-identical legacy path. Threading (layers/encoders.py, models/models.py): - VisionEncoder / the model forward thread encoder_image_position_ids to the Gemma-4 vision encoder and route the returned validity mask into MultimodalInput.image_masks, so exactly the valid pooled tokens land in the image placeholders (merge_mm_embeddings.token_masks). Decoder (layers/nnx_decoders.py): - ple_pad_substitute_image_rows: Gemma-4 E2B/E4B build the per-layer inputs from llm_input_ids with image placeholder tokens mapped to pad_token_id (HF modeling_gemma4), rather than feeding the placeholder id into the PLE path. Without this the per-layer embeddings at the image positions diverge, corrupting the image-span and post-image logits. - use_bidirectional_image_attn: E2B/E4B image spans are causal; suppress the bidirectional attention carve-out unless explicitly enabled (bidirectional-image models set it True). Config (configs/types.py, configs/base.yml): - Adds 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). - Allows E2B/E4B multimodal when use_clipped_linears_for_vit is set. Validation (CPU, teacher-forced 340-token forward vs HF reference logits, converted E2B checkpoint): with the full fix and clip-bounds enabled, pre_image max_KL 2.7e-5, argmax 1.0 image_span max_KL 9.0e-5, argmax 1.0 post_image max_KL 4.3e-5, argmax 1.0 (frozen gate: <= 1.26e-3 and argmax >= 0.995) matching the reference. With clip-bounds disabled the image span diverges, confirming the clip-bounds are required. Defaults keep all changes off for non-Gemma-4 models.
So enabling image inputs only requires use_multimodal=true + use_clipped_linears_for_vit=true; the E2B/E4B model configs supply the decoder image-contract flags (causal image spans, PLE image-row pad substitution, placeholder/pad token ids). Also drops the stale 'multimodal not yet supported' comment.
lokic233
changed the title
|
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:
With all three plus the clip-bounds, the teacher-forced forward now matches the HF reference: post-image |
…ut without image_position_ids The legacy all-valid path expects full images ([B,H,W,C]/[B,N,H,W,C]). Pre-patchified pixel_values require image_position_ids (the padded-patch path). Previously a 3D pre-patchified input with no positions would crash cryptically on shape-unpack; now it raises a clear contract error instead of silently mis-handling it.
…mizer freeze integration
3.1 _clip_in/_clip_out now read bounds through jax.lax.stop_gradient (defense in depth: clip-bound
gradients can never contaminate gradient statistics even if the freeze mask were misconfigured).
3.2 get_optimizer() now wires clip_optimizer_freeze_mask into optax.multi_transform when
use_clipped_linears_for_vit is set: the 448 clip bounds map to set_to_zero() so they receive no
updates, no weight decay, and no momentum/variance slots. Composes with trainable_parameters_mask
(a leaf is trainable only if trainable under the whitelist AND not a clip bound).
…closed) + fused_qkv guard
M4: validate_all_vision_clip_bounds(model) walks the model graph, validates every clip-state module
(finite, scalar, input_min<=input_max, output_min<=output_max) and asserts EXACTLY 112 modules /
448 scalar bounds. Wired into from_pretrained AFTER restore, BEFORE first JIT, gated on
use_clipped_linears_for_vit + use_multimodal. Exact-count check prevents a zero-module traversal
from silently passing. Also strengthened validate_clip_bounds with the min<=max ordering check.
M5: Gemma4Attention.enable_vision_clip_bounds now hard-fails on fused_qkv=True (distinct q/k/v clip
bounds require separate projections); fused_mlp=True guard already present in Gemma4ClippedMlpBlock.
…, all pass) CI-grade unit tests covering the merge-critical invariants: clip math bit-exact vs jnp.clip, dtype-preserving, exact no-op when disabled, stop_gradient => zero grad wrt bound, NaN/Inf/min>max hard-fail, freeze-mask polarity+count, expected 112/448 count constants, public symbols present. Runs on CPU in grr-maxtext:latest (13 passed).
M1: thread image_position_ids through the stock loss_fn (train.py) for both
Linen and NNX paths, guarded on presence in the batch (the stock fixed-grid
Gemma-4 processor does not emit it -> legacy all-valid path, correct for its
contract; native-resolution/pan-and-scan processors do -> Option-S path).
M2: make ple_pad_mode a validated enum (hard-fail unknown); 'identity' documented
as the HF-faithful default (matches Transformers 5.9.0 get_per_layer_inputs,
which ignores inputs_embeds when input_ids is given); wire 'both' as a real
(HF-divergent, ablation-only) mode by threading shared_embedding into the
Gemma-4-small PLE path -- removing the previously-dead hasattr(self,...) branch.
M6: Gemma-4 E2B/E4B multimodal static contract gate in config validation:
enum PLE mode, non-negative token ids (single source of truth), causal image
spans (use_bidirectional_image_attn=False), fused_qkv/fused_mlp=False.
Pure-Linen Gemma-4-small MM path now fail-closed (PLE substitution is NNX-only).
No behavior change for non-Gemma-4 models or text-only Gemma-4.
…n-finite pad from valid outputs)
…ard test Packed multimodal is unsupported by the stock data pipeline; packing image spans risks cross-doc image attention + PLE image-row substitution across segment boundaries. MaxText forbade this for multimodal SFT only; extend to every training mode. + in-tree packing-guard test.
lokic233
changed the title
|
Update — hardening pushed to the PR head (797842b); ready for re-review (E2B). The three review-bot findings are addressed with tests (replies in-thread): (1) clip validation now runs post-restore with exact 112/448 accounting; (2) clip bounds are immutable (stop_gradient at use-site + optimizer Scope has been made honest: E2B full E2E parity is validated (post_image max_KL 4.29e-5/argmax 1.0 structured; 4.71e-5/1.0 natural; full component ablation in the body). E4B is vision-tower-validated only (cosine 0.999999); full E4B decoder E2E parity is pending and is not claimed. Real-TPU: |
1 participant
What
Enables Gemma-4 E2B full end-to-end image (multimodal) parity in MaxText. E4B runs through the same
contract-gated code path and its vision tower is validated (cosine 0.999999 vs HF), but E4B full decoder
E2E parity is PENDING (no E4B E2E fixture/converted checkpoint this cycle) — E4B is not claimed as
E2E-validated.
Reference contract (pinned)
transformers/models/gemma4/modeling_gemma4.py), semantic reference:Gemma4ClippableLinear(clamp-in → linear → clamp-out) + the PLE block inGemma4ForConditionalGeneration.forward.The four load-bearing pieces (all required; ablation below):
modules / 448 scalar bounds).
pad_token_idin the token-identity PLE path (matchesHF; adjudicated against 5.9.0:
get_per_layer_inputsignoresinputs_embedswheninput_idsis given, sothe context path keeps the merged image features →
ple_pad_mode="identity"is HF-faithful;"both"isHF-divergent and ablation-only).
Stock trainer contract
The stock Gemma-4 processor uses a fixed-grid 672×960 → 2520 all-valid patches → 280 soft tokens (no
padding). For that contract the legacy all-valid vision path is correct and
image_position_idsis absent.The stock loss_fn now threads
image_position_idsto the vision encoder when a native-resolution/pan-and-scanprocessor emits it (Option-S); pre-patchified input without positions fails closed (refuses to silently
take the legacy path). Multimodal sequence packing fails closed (unsupported; would risk cross-document
image attention); the E2B/E4B ymls default
packing=false.Hardening invariants
validate_all_vision_clip_bounds(nnx.iter_graph) asserts exactly 112modules / 448 bounds, all finite, scalar,
min<=max; hard-fails on missing/NaN/Inf/wrong-count (incl. zero).Wired into
from_pretrainedbefore the first JIT.lax.stop_gradientat every clamp use-site +clip_optimizer_freeze_maskwired intoget_optimizer(optax.multi_transform→set_to_zero, no updates/decay/slots), independent offreeze_vision_encoder_params.[-1,y]/[x,-1]), placeholder==pooled count.A real correctness bug the garbage-pad test found
The garbage-pad invariant test exposed a real bug: NaN-valued padded-patch rows could contaminate valid
pooled outputs before attention masking (attention computes q/k/v on all patches before the segment mask). The
fix sanitizes padded rows to zero before the numerically-unsafe projections; finite-garbage (1e6) and NaN-pad
tests now leave the valid pooled outputs byte-identical. This is a load-bearing test success.
Validation
E2B, MaxText NNX vs pinned HF 5.9.0, teacher-forced (gate: post_image max_KL ≤ 1.26e-3 = 3× text-only FP32 floor
4.2e-4; argmax ≥ 0.995 — see provenance):
Ablation (E2B structured — every component load-bearing):
In-tree tests (CPU):
tests/unit/gemma4_clip_linears_test.py(18) +tests/unit/gemma4_mm_contract_gate_test.py(8).Real TPU/BF16:
from_pretrainedload +validate_all_vision_clip_boundsPASS (112/448) on TPU7x; a cleanone-step trainer+checkpoint-roundtrip demonstration is pending (a flax-version nnx grad buffer-lifetime issue
in the demonstration harness, not the model code; the clip-immutability mechanism is unit-proven).
Reproduction (no private GCS)
JAX_PLATFORMS=cpu python -m pytest tests/unit/gemma4_clip_linears_test.py tests/unit/gemma4_mm_contract_gate_test.py -vThe E2E parity uses a converted HF→MaxText E2B checkpoint + the frozen fixtures; see REPRO_COMMANDS.
Limitations