Gemma-4 vision: correct std_bias/std_scale classification per HF standardize (fidelity fix)

[lokic233]

Summary

Correct the classification and construction of the Gemma-4 vision standardization state (std_bias/std_scale) so it matches HuggingFace modeling_gemma4 per model variant, and make the contract explicit and fail-closed for every Gemma-4 config.

This is a standalone fidelity fix. It does not add or change the multimodal feature itself (that is #4790), and it does not fix the separate params-only checkpoint-loader issue (tracked separately, see Non-goals).

Exact model variants affected

Variant	HF vision_config.standardize	std tensors in checkpoint	Correct MaxText state
gemma4-26b	True	2 (std_bias, std_scale)	checkpoint-resident, non-trainable buffers
gemma4-31b	True	2	checkpoint-resident, non-trainable buffers
gemma4-e2b	False	0	none (standardize is an exact identity)
gemma4-e4b	False	0	none

(Measured from checkpoint_conversion/utils/hf_model_configs.py, which mirrors the HF configs: 26B/31B standardize=True, E2B/E4B standardize=False.)

The bug

The vision encoder unconditionally constructed std_bias/std_scale as trainable nnx.Param leaves and always applied the standardize op. For E2B/E4B (standardize=False) this fabricates two identity nnx.Param leaves that are absent from the HF model and its checkpoint, and that wrongly enter the nnx.Param gradient filter. For 26B/31B (standardize=True) the state existed but was modeled as trainable, whereas HF registers it via register_buffer (non-trainable).

Separately, the previous version of this field defaulted to False globally. Since a per-model-family semantic distinction must not inherit a global default, a 26B/31B config that failed to set the field would silently drop its std buffers.

The fix

1. Explicit, fail-closed config contract. standardize_for_vit: bool | None = None. A model_validator hard-fails any gemma4-* model whose standardize_for_vit is unset. All four Gemma-4 model YAMLs set it explicitly (E2B/E4B false, 26B/31B true).
2. Correct state construction.
   • standardize_for_vit=False: construct no std state; the standardize op is an exact identity.
   • standardize_for_vit=True: construct std_bias/std_scale as VisionStdVar — a plain nnx.Variable (not nnx.Param). It is checkpoint-resident and restored by Orbax, but non-trainable: excluded from the optimizer, weight decay, and gradient filter (which all key on nnx.Param), matching HF's register_buffer semantics. This reuses the existing MaxText idiom for non-trainable buffers (MoEBiasVar/Tid2EidVar in layers/moe.py).

26B/31B compatibility

• Both remain standardize_for_vit=true; the std state is still constructed, used in the forward standardize op, and saved/restored — only its Variable type changes from nnx.Param to a non-trainable nnx.Variable, which is the HF-faithful classification.
• Orbax save/reload preserves the values byte-exactly; a known bias/scale transforms the output exactly.

Tests

tests/unit/gemma4_vision_standardize_test.py (real model configs):

• E2B (standardize=false) → no std state anywhere (Param or Variable), no attribute.
• 26B (standardize=true) → exactly two VisionStdVar leaves, zero in nnx.Param.
• All four configs resolve standardize_for_vit to their HF truth from the real YAMLs.

Additional verified (state semantics, §A2): optimizer built with wrt=nnx.Param contains zero std leaves in its opt_state; Orbax round-trip preserves std_bias/std_scale exactly and keeps them as VisionStdVar; (x - std_bias) * std_scale matches expected within 1e-6.

Composite integration validation (backprop evidence provenance)

The end-to-end backward/optimizer-apply evidence was produced by the composite stack — not by checking out this PR alone:

PR #4790 head <797842b1>            (Gemma-4 E2B multimodal feature: vision clipped-linears, PLE, masking)
+ PR #4807 candidate <this branch>  (std-fidelity fix)
+ a full-state checkpoint remap/splice bridge   (params-only θ0 -> full train-state; see Non-goals)
+ stock train.py

• CPU: real E2B θ0 → forward → backward → optimizer apply → save θ1: loss 14.941 → 11.596 (finite, decreasing); θ1 fresh-reload green; 448 vision clip bounds byte-identical θ0→θ1 (value-hash 750460d3…); 751/751 non-clip trainable leaves changed.
• TPU (tpu7x, single host): same trainer path: loss 15.917 → 14.020 (finite, decreasing); θ1 saved with populated opt_state (proves gradients flowed + optimizer applied); trainer exit 0; no array deletion.

Clip-immutability provenance: CPU byte-proves the 448-bound immutability (pre/post value-hash identical). TPU independently proves backward, optimizer apply, and θ1 save. (A TPU-side pre/post 448-bound byte receipt is not asserted here.)

Scope / non-goals

• This PR does not enable or modify the multimodal feature (that is Gemma-4 E2B: enable image (multimodal) E2E parity — vision clipped-linears + padded-patch masking + PLE substitution (E4B vision-tower validated; E4B E2E pending) #4790) and must not be read as replacing or superseding Gemma-4 E2B: enable image (multimodal) E2E parity — vision clipped-linears + padded-patch masking + PLE substitution (E4B vision-tower validated; E4B E2E pending) #4790.
• This PR does not fix the separate load_parameters_path params-only checkpoint issue (a converted params-only checkpoint currently drops a 69-leaf vision tail on restore; the E2E evidence used a full-state remap/splice bridge). That is a distinct checkpoint-loader issue tracked separately; the remap/splice bridge is evidence-only, not a proposed upstream loader path.
• The standalone "deleted array under value_and_grad" symptom seen in earlier debugging was an artifact of extracting/donating state from a live committed Orbax-restored model inside the transform (eager value_and_grad/nnx.jit/nnx.split); the production train.py does not follow that pattern. It is not the claim of this PR; the claim is the std-state fidelity correction above.

[gemini-code-assist]

Code Review

This pull request introduces a new configuration parameter standardize_for_vit to control whether the Gemma-4 vision encoder standardizes exit embeddings with learned std_bias and std_scale parameters. When disabled (the default), these parameters are not created, preventing them from incorrectly entering the gradient filter as trainable parameters. Unit tests have been added to verify this behavior. I have no additional feedback to provide.

[lokic233]

Note on the checkpoint-conversion path (for reviewers):

checkpoint_conversion/utils/param_mapping.py already conditions the std synthesis hooks on the HF-side flag vcfg["standardize"] (it synthesizes std_scale=ones / std_bias=zeros only when standardize=False, under a block currently marked as dead for gemma4-small MM). This model-side change is the symmetric counterpart on the model side: standardize_for_vit mirrors the same HF vision_config.standardize flag, so the model and the converter stay consistent (both keyed off the HF config). For E2B/E4B (standardize=False) the converted checkpoints carry 0 std keys — verified against a real converted E2B checkpoint (0/2012 std keys) — so with this fix the model has no std Params and the restore has no extra/missing-key mismatch.

This PR is intentionally scoped to the model-side guard + test; it does not modify the (currently dead) conversion hooks.

[lokic233]

Additional validation — backward pass confirmed on real TPU (single-host tpu7x):

Beyond the CPU result, the same R2-fixed trainer was run on a real TPU (tpu7x-standard-1t, single host, batch 1) restoring the real Gemma-4 E2B checkpoint via the production full-state path:

restoring from this run's directory step 0
completed step: 0, loss: 15.917 (finite)
completed step: 1, loss: 14.020 (finite, decreasing)   # after backward + optimizer apply
Saved a checkpoint at step 1        # full train-state: params(1199) + opt_state(2400) + nnx_aux + step
trainer exit = 0                    # no "Array has been deleted"

Finite, decreasing loss and a saved θ1 (with populated opt_state, proving gradients flowed and the optimizer applied) on the actual target platform. Vision QK-clip bounds remained byte-identical across the step (clip-freeze). This confirms the fix unblocks E2B/E4B multimodal training with backprop on TPU, not just CPU.

[lokic233]

Hardening update (force-pushed with --force-with-lease; previous head d4c3dd18 backed up locally):

• Rebased onto current main (3f313d17e), squashed to a single commit, 7 files, +163/-9.
• Explicit fail-closed config contract: standardize_for_vit: bool | None = None + a model_validator that hard-fails any gemma4-* with it unset. All four Gemma-4 YAMLs now set it explicitly (E2B/E4B=false, 26B/31B=true), so no variant can silently inherit the wrong value.
• 26B/31B (standardize=true) state corrected: std_bias/std_scale are now VisionStdVar (a non-trainable nnx.Variable), not nnx.Param — checkpoint-resident + Orbax-restored but excluded from optimizer/weight-decay/gradient, matching HF register_buffer. This reuses MaxText's existing non-trainable-buffer idiom (MoEBiasVar/Tid2EidVar). Verified: optimizer(wrt=nnx.Param) has zero std leaves; Orbax round-trip preserves values; forward transform exact.
• Tests use the real model configs (not a 26B config edited to simulate E2B): 3/3 pass.

Correction on the earlier TPU comment: there is no TPU-side pre/post 448-bound byte receipt. CPU byte-proves the 448 clip-bound immutability (pre/post value-hash identical); TPU independently proves backward + optimizer-apply + θ1 save. The PR body reflects this.

Still draft; not marking ready-for-review pending human owner decision.