Problem

Fine-tuning Gemma-4 26B-A4B or 31B Vision with dtype = torch.float16 produces NaN training loss from step 1 on transformers 5.5.0 through 5.9.x (and the yanked 5.10.0). Text and audio are unaffected, as are E2B / E4B vision and all bf16 runs.

Root cause is an upstream transformers bug: Gemma4VisionPooler.forward scales pooled activations by sqrt(hidden_size) in the input dtype:

hidden_states *= self.root_hidden_size

The 26B / 31B vision tower (hidden 1152, sqrt = 33.94, no clipped linears) produces fp16 activations around 2300, which scale to roughly 79k, past the fp16 max of 65504, giving inf and then NaN through embed_vision. This engages under forced-float32 mode, which keeps the unquantized vision tower in fp16 with autocast off. Localized empirically on a 26B fp16 run: the pooler input was finite (amax 2330), its output inf, and every module before it finite. Upstream fixed this in transformers 5.10.1 via huggingface/transformers#46277.

Fix

patch_Gemma4VisionPoolerFP16 in temporary_patches/gemma4.py mirrors the upstream fix with two pass-through wrappers:

• Gemma4VisionPooler.forward: upcasts a float16 hidden_states input to fp32. The buggy pooler's ops (masked_fill, avg pool, scale) all preserve input dtype, so this matches the fixed pooler's fp32 output up to fp16 rounding of the pooled activations. Dtype promotion then carries fp32 through the caller's standardization, where the std_bias subtraction cancels the large values (buffer amax is 53760, inside fp16 range).
• Gemma4VisionModel.forward: casts last_hidden_state back to the encoder's working dtype, taken from the patch embedder's input_proj weight (the embedder casts pixels to it, so this equals upstream's inputs_embeds.dtype even when callers pass fp32 processor pixels to an fp16 tower). Only engages for fp16 towers and also handles return_dict = False tuple returns.

Design points:

• Source-pattern classification instead of version compares: known-buggy sources get patched, known-fixed sources (5.10.1 and later) are skipped entirely, unrecognized future sources are left untouched with an optional log line, and transformers without gemma4 (4.57.6) no-op via ImportError. This also handles the yanked 5.10.0 correctly.
• The wrappers call the original forwards, so the HF decorators on Gemma4VisionModel.forward (merge_with_config_defaults, capture_outputs) are preserved, as is signature introspection via functools.wraps.
• The vision wrapper installs before the pooler wrapper, making the pooler marker the commit point; a missing vision wrapper next to a marked pooler is re-installed rather than reported as already patched.
• bf16 and fp32 setups, uniform or mixed input dtypes, are bit-identical to upstream (the pooler gate and the cast-back both check for fp16).
• No compiler companion is needed: the compiler never traces Gemma4VisionPooler or Gemma4VisionModel (verified by grepping generated unsloth_compiled_cache modules from 52 notebook-matrix runs plus fresh validation runs).
• vLLM's gemma4 multimodal path calls the transformers pooler with keyword arguments and casts to its model dtype before embed_vision, so the fp32 pooler output composes there too and fixes the same overflow for fp16 vLLM multimodal on old transformers.

Testing

CPU (tests/test_gemma4_vision_pooler_fp16.py, 19 tests): overflow reproduction, fp32-reference match, standardize on and off, mixed fp32-pixels/fp16-tower, fp32-pixels/bf16-tower untouched, bf16/fp32 bit-identical, tuple returns, kwarg calls, idempotency, repair path, drift and no-source classification, gradient flow, and a real-source canary that fails loudly if upstream rewrites the pooler. Green on transformers 4.57.6 (no gemma4, no-op), 5.5.0 (patches), and 5.14.1 (skips), and the classification was checked against every 5.x release tag: 5.5.0-5.10.0 classify buggy, 5.10.1-5.14.1 classify fixed.

GPU (B200, transformers 5.5.0, 3-step LoRA notebook replicas with inference and merged save):

Also reproduced the mixed-dtype path on real 5.5.0 classes: fp32 pixels into an fp16 Gemma4VisionModel return fp16 finite features after the patch.

Note: the package currently pins transformers<=5.5.0, which is exactly the buggy release users get; if the pin is later lifted into 5.6-5.9 the patch covers those too, and from 5.10.1 it automatically stands down. The complementary notebook-side fix would be bumping the Gemma-4 vision notebooks off the ==5.5.0 pin once supported.