[Bug]: Gemma 4 E2B QAT checkpoints fail to load: k_norm required for KV-shared layers (num_kv_shared_layers ignored by loader)

[xbill9]

Summary

The Gemma 4 E2B QAT checkpoints fail to load on TPU in every currently serveable form, on both the JAX and torchax model paths. The root cause for the bf16 QAT checkpoint appears to be that the Gemma 4 model implementation instantiates self_attn.k_norm for all layers and requires them from the checkpoint, without accounting for num_kv_shared_layers — the KV-shared upper layers have no k_norm of their own, and the QAT export (correctly, per the config) omits them.

Environment

• Image: vllm/vllm-tpu:nightly (vLLM v0.23.1rc1.dev1076+g5c342876a), pulled 2026-07-21
• Hardware: v6e-1 (ct6e-standard-1t GCE flex-start VM, europe-west4-a)
• --tensor-parallel-size 1 --max-model-len 65536 --gpu-memory-utilization 0.9 --max_num_batched_tokens 4096 --disable_chunked_mm_input

Failure 1: google/gemma-4-E2B-it-qat-q4_0-unquantized, JAX path (default)

ValueError: Following weights were not initialized from checkpoint:
{'model.language_model.layers.15.self_attn.k_norm.weight', ...
 'model.language_model.layers.34.self_attn.k_norm.weight'}   # layers 15–34, i.e. exactly the 20 KV-shared layers

(With --load-format runai_streamer the same underlying problem surfaces later and more cryptically, as TypeError: Argument 'model.states[0][223]' of shape bfloat16[256] of type <class 'jax.ShapeDtypeStruct'> is not a valid JAX type during the first jitted call — the un-loaded param is still an abstract value.)

Failure 2: same checkpoint, MODEL_IMPL_TYPE=vllm (torchax path)

File ".../model_loader/default_loader.py", line 467, in track_weights_loading
ValueError: Following weights were not initialized from checkpoint:
{'language_model.model.layers.15.self_attn.k_norm.weight', ... 'language_model.model.layers.34.self_attn.k_norm.weight'}

Checkpoint forensics (why this is a loader bug, not a broken export)

Read directly from the safetensors headers of both HF repos:

checkpoint	k_norm tensors present	config
google/gemma-4-E2B-it	all 35 language-model layers	num_hidden_layers: 35, num_kv_shared_layers: 20
google/gemma-4-E2B-it-qat-q4_0-unquantized	only the 15 non-KV-shared layers (+ towers)	identical

The configs are byte-identical on these fields. Layers 15–34 are the 20 KV-shared layers; they reuse K/V from lower layers, so the QAT export ships no k_norm for them. The plain checkpoint happens to include (unused) k_norm tensors for those layers, which is the only reason it loads while the QAT one does not. The plain google/gemma-4-E2B-it serves fine in the same environment.

Expected behavior: the Gemma 4 implementation should not create/require k_norm (or K/V projections) for the num_kv_shared_layers upper layers, so that architecturally-minimal exports like the QAT one load.

Failure 3 (related): google/gemma-4-E2B-it-qat-w4a16-ct

The compressed-tensors int4 QAT variant fails earlier, in scheme resolution:

File ".../tpu_inference/layers/jax/quantization/compressed_tensors.py", line 146, in get_quant_method
NotImplementedError: compressed-tensors scheme for layer 'model.language_model.per_layer_model_projection' is not yet supported in the JAX path.

Happy to provide full logs or test candidate fixes on live v6e capacity.

Repro

docker run --privileged --net=host --shm-size 10gb -v /dev/shm:/dev/shm \
  -e HF_HOME=/dev/shm -e HF_TOKEN=... \
  vllm/vllm-tpu:nightly vllm serve google/gemma-4-E2B-it-qat-q4_0-unquantized \
  --tensor-parallel-size 1 --max-model-len 65536