diff --git a/src/maxtext/configs/models/gemma4-26b.yml b/src/maxtext/configs/models/gemma4-26b.yml index b20397dee7..77cd678b9b 100644 --- a/src/maxtext/configs/models/gemma4-26b.yml +++ b/src/maxtext/configs/models/gemma4-26b.yml @@ -61,3 +61,4 @@ num_attention_heads_for_vit: 16 image_placeholder: "<|image|>" vision_output_length: 280 num_position_embeddings_for_vit: 10240 +standardize_for_vit: true # HF vision_config.standardize=true: std_bias/std_scale are checkpoint-resident buffers diff --git a/src/maxtext/configs/models/gemma4-31b.yml b/src/maxtext/configs/models/gemma4-31b.yml index 7e97f9dba5..59f69c0ec6 100644 --- a/src/maxtext/configs/models/gemma4-31b.yml +++ b/src/maxtext/configs/models/gemma4-31b.yml @@ -57,3 +57,4 @@ num_attention_heads_for_vit: 16 image_placeholder: "<|image|>" vision_output_length: 280 num_position_embeddings_for_vit: 10240 +standardize_for_vit: true # HF vision_config.standardize=true: std_bias/std_scale are checkpoint-resident buffers diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml index 81c8d4ea66..226a0d4362 100644 --- a/src/maxtext/configs/models/gemma4-e2b.yml +++ b/src/maxtext/configs/models/gemma4-e2b.yml @@ -58,3 +58,4 @@ num_attention_heads_for_vit: 12 image_placeholder: "<|image|>" vision_output_length: 280 num_position_embeddings_for_vit: 10240 +standardize_for_vit: false # HF vision_config.standardize=false: no std_bias/std_scale in checkpoint (identity) diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml index b8da0d1d46..9c84ef6475 100644 --- a/src/maxtext/configs/models/gemma4-e4b.yml +++ b/src/maxtext/configs/models/gemma4-e4b.yml @@ -59,3 +59,4 @@ num_attention_heads_for_vit: 12 image_placeholder: "<|image|>" vision_output_length: 280 num_position_embeddings_for_vit: 10240 +standardize_for_vit: false # HF vision_config.standardize=false: no std_bias/std_scale in checkpoint (identity) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index dc536f9799..cfe5df5ba9 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2143,6 +2143,20 @@ class MultimodalGeneral(BaseModel): description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').", ) freeze_vision_encoder_params: bool = Field(True, description="Freeze the parameters of the vision encoder.") + standardize_for_vit: bool | None = Field( + None, + description=( + "Whether the vision encoder standardizes exit embeddings with std_bias/std_scale. Mirrors HF" + " Gemma-4 vision_config.standardize and is a per-model-family semantic distinction, so it must be set" + " EXPLICITLY by every Gemma-4 multimodal model config (no global default is safe): HF registers" + " std_bias/std_scale as checkpoint-resident (non-trainable) buffers ONLY when standardize=True." + " Gemma-4 26B/31B ship standardize=True (2 std tensors); E2B/E4B ship standardize=False (0 std tensors)." + " When False the standardize op is an exact identity and MaxText must NOT construct std_bias/std_scale" + " (fabricating them wrongly enters checkpoint-restored committed arrays into the nnx.Param gradient" + " filter). None means unset: a Gemma-4 model with use_multimodal=True and this unset is a hard config" + " error (see validate_standardize_for_vit) — it must never silently inherit a value." + ), + ) freeze_audio_encoder_params: bool = Field(True, description="Freeze the parameters of the audio encoder.") use_audio: bool = Field(False, description="Enable audio encoder for multimodal models.") image_size_for_vit: int | list[int] | None = Field(896, description="Input image size for the Vision Transformer.") @@ -3581,6 +3595,15 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de raise ValueError( f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with nn.scan)." ) + # Gemma 4 vision standardization is a per-model-family semantic distinction (HF vision_config.standardize) + # and must be set EXPLICITLY — it must never silently inherit a global default, or a standardize=True variant + # (26B/31B) would drop its checkpoint-resident std buffers, or a standardize=False variant (E2B/E4B) would + # fabricate spurious trainable std leaves. Fail closed when unset for any Gemma-4 model. + if self.model_name.startswith("gemma4-") and self.standardize_for_vit is None: + raise ValueError( + f"{self.model_name}: standardize_for_vit must be set explicitly (True for 26B/31B, False for E2B/E4B) " + "to mirror HF vision_config.standardize. Set it in the model config YAML; it must not inherit a default." + ) if self.use_multimodal: # Gemma 4 small (E2B / E4B) only supports text for now; multimodal # support is pending clipped-linears in the vision encoder. diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py index 72cd841f81..ea2f19839c 100644 --- a/src/maxtext/models/gemma4_vision.py +++ b/src/maxtext/models/gemma4_vision.py @@ -85,6 +85,17 @@ def patchify(images: jax.Array, patch_size: int) -> tuple[jax.Array, jax.Array]: return patches, jnp.broadcast_to(positions_xy, tuple(b) + positions_xy.shape) +class VisionStdVar(nnx.Variable): + """Checkpoint-resident, NON-trainable vision standardization state (std_bias / std_scale). + + Mirrors HF `modeling_gemma4`, which registers `std_bias`/`std_scale` via `register_buffer` + (non-trainable) and only when `vision_config.standardize=True` (Gemma-4 26B/31B). As a plain + `nnx.Variable` (not `nnx.Param`), it is naturally excluded from the optimizer/weight-decay/grad + filter — which key on `nnx.Param` — while still being saved and restored by Orbax as part of the + model state. This is the same idiom MaxText uses for other non-trainable buffers (e.g. `MoEBiasVar`). + """ + + class VisionEntry(nnx.Module): """The vision entry layer.""" @@ -566,12 +577,22 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs): rngs=self.rngs, precision=config.matmul_precision, ) - self.std_bias = nnx.Param( - nnx.initializers.zeros(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,) - ) - self.std_scale = nnx.Param( - nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,) - ) + # Vision standardization (std_bias/std_scale) is CHECKPOINT-RESIDENT state that exists ONLY when the + # checkpoint ships std tensors, i.e. HF vision_config.standardize=True (Gemma-4 26B/31B). HF registers them + # via register_buffer (NON-trainable). Gemma-4 E2B/E4B have standardize=False and DO NOT store std_scale/ + # std_bias in their safetensors; when disabled the standardize op is an EXACT identity, so we construct NO std + # state at all (fabricating trainable identity nnx.Param leaves wrongly enters checkpoint-restored committed + # arrays into the nnx.Param gradient filter). When enabled, we construct them as VisionStdVar (a plain + # nnx.Variable): checkpoint-resident and restored by Orbax, but NON-trainable — excluded from the optimizer / + # weight-decay / gradient filter (which key on nnx.Param), matching HF's register_buffer semantics. + self._standardize_for_vit = bool(config.standardize_for_vit) + if self._standardize_for_vit: + self.std_bias = VisionStdVar( + nnx.initializers.zeros(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype) + ) + self.std_scale = VisionStdVar( + nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype) + ) def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array: """Applies the vision encoder layer.""" @@ -593,9 +614,10 @@ def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array: # We take the first result. (embeddings, _) = vision_exit_results[0] - embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype( - embeddings.dtype - ) + if self._standardize_for_vit: + embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype( + embeddings.dtype + ) # Unflatten batch and num_images final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2])) diff --git a/tests/unit/gemma4_vision_standardize_test.py b/tests/unit/gemma4_vision_standardize_test.py new file mode 100644 index 0000000000..941591a6b3 --- /dev/null +++ b/tests/unit/gemma4_vision_standardize_test.py @@ -0,0 +1,105 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Gemma-4 vision standardization (std_bias/std_scale) contract. + +HF `modeling_gemma4` registers vision `std_bias`/`std_scale` via `register_buffer` (non-trainable) and ONLY when +`vision_config.standardize=True`. Per-variant HF truth: Gemma-4 26B/31B ship standardize=True (2 std tensors); +E2B/E4B ship standardize=False (0 std tensors). This test locks the MaxText contract: + + * The `standardize_for_vit` config field is a per-model-family semantic distinction and must be set explicitly + by every Gemma-4 model config (no silent global default); an unset Gemma-4 config is a hard config error. + * standardize_for_vit=False -> NO std state at all (the standardize op is an exact identity). + * standardize_for_vit=True -> exactly two std leaves, constructed as VisionStdVar (a non-trainable + `nnx.Variable`, NOT `nnx.Param`), so they are checkpoint-resident but excluded from the optimizer/gradient. +""" + +import unittest + +from flax import nnx +import jax +import numpy as np +from jax.sharding import Mesh + +from maxtext.configs import pyconfig +from maxtext.utils.globals import MAXTEXT_REPO_ROOT +from maxtext.models.gemma4_vision import Gemma4VisionEncoderLayer, VisionStdVar + + +def _init(model_name): + base_config_path = f"{MAXTEXT_REPO_ROOT}/src/maxtext/configs/base.yml" + argv = [ + "", + base_config_path, + f"model_name={model_name}", + "attention=dot_product", + "matmul_precision=highest", + "dtype=float32", + "dtype_mm=float32", + "weight_dtype=float32", + "skip_jax_distributed_system=true", + "enable_checkpointing=false", + "run_name=std_test", + "base_output_directory=/tmp/std_test", + ] + if model_name in ("gemma4-e2b", "gemma4-e4b"): + argv.append("scan_layers=false") + return pyconfig.initialize(argv) + + +class TestGemma4VisionStandardizeContract(unittest.TestCase): + """standardize_for_vit gates whether/how std_bias/std_scale exist, per HF vision_config.standardize.""" + + def setUp(self): + self.mesh = Mesh(np.array(jax.devices()[:1]), axis_names=("data",)) + + def _std_leaves(self, model, filt): + names = set() + for path, _ in jax.tree_util.tree_flatten_with_path(nnx.state(model, filt))[0]: + ps = "/".join(str(getattr(p, "key", p)) for p in path) + if "std_bias" in ps or "std_scale" in ps: + names.add(ps) + return names + + def test_e2b_standardize_false_has_no_std_state(self): + """E2B (HF standardize=false): no std state at all; standardize is an exact identity.""" + cfg = _init("gemma4-e2b") + self.assertFalse(cfg.standardize_for_vit) + model = Gemma4VisionEncoderLayer(cfg, self.mesh, rngs=nnx.Rngs(0)) + self.assertFalse(hasattr(model, "std_bias")) + self.assertFalse(hasattr(model, "std_scale")) + self.assertEqual(self._std_leaves(model, nnx.Param), set()) + self.assertEqual(self._std_leaves(model, nnx.Variable), set()) + + def test_26b_standardize_true_has_nontrainable_std_buffers(self): + """26B (HF standardize=true): exactly two std leaves as VisionStdVar, NOT in nnx.Param.""" + cfg = _init("gemma4-26b") + self.assertTrue(cfg.standardize_for_vit) + model = Gemma4VisionEncoderLayer(cfg, self.mesh, rngs=nnx.Rngs(0)) + self.assertIsInstance(model.std_bias, VisionStdVar) + self.assertIsInstance(model.std_scale, VisionStdVar) + # Present as (non-Param) Variables, absent from the Param (trainable) collection. + self.assertEqual(self._std_leaves(model, nnx.Param), set()) + self.assertEqual(len(self._std_leaves(model, nnx.Variable)), 2) + + def test_e2b_e4b_false_and_26b_31b_true(self): + """All four Gemma-4 configs resolve standardize_for_vit to their HF truth from the real YAMLs.""" + self.assertFalse(_init("gemma4-e2b").standardize_for_vit) + self.assertFalse(_init("gemma4-e4b").standardize_for_vit) + self.assertTrue(_init("gemma4-26b").standardize_for_vit) + self.assertTrue(_init("gemma4-31b").standardize_for_vit) + + +if __name__ == "__main__": + unittest.main()