K_NORM MISSING

#2
by dhifafaz - opened

Is the k_norm really missing when you guys exported the model?

Thanks for flagging this. We audited the serialized checkpoint and performed a full Transformers load check. No required k_norm tensors are missing.

This is Gemma 4 E2B with 35 language layers and num_kv_shared_layers: 20. In Gemma 4, the final 20 language layers reuse key/value states from the non-shared stack. They therefore do not instantiate or serialize separate k_proj, v_proj, or k_norm tensors. The checkpoint correctly contains:

  • language k_norm for layers 0-14: 15 tensors;
  • vision k_norm for layers 0-15: 16 tensors;
  • no language k_norm for shared-KV layers 15-34.

With transformers==5.12.1, AutoModelForImageTextToText.from_pretrained(..., output_loading_info=True) reports zero missing keys, zero unexpected keys, and zero mismatched keys. We also added this compatibility note to the model card and now gate future exports by checking the serialized layer sets against config.json.

If your loader expects a separate k_norm on all 35 language layers, it is not applying Gemma 4's shared-KV configuration. Please update to transformers>=5.12.1; do not synthesize or copy the absent shared-layer tensors. The exact Transformers 5.12.1 Gemma 4 implementation is here: https://github.com/huggingface/transformers/blob/v5.12.1/src/transformers/models/gemma4/modeling_gemma4.py

If the warning persists, please post the loader name, version, and exact warning so we can reproduce that integration specifically.

I think you need to upgrade your version of Transformers. Transformers 5.12.1 loads all 1,951 serialized tensors with zero missing, unexpected, or mismatched keys. The apparent 60-key difference is the intended shared-KV representation, not lost weights, which is intended. E2B shares key/value states across 20 language layers, so those layers should not serialize their own k_proj, v_proj, or k_norm.

MirilAI changed discussion status to closed

Sign up or log in to comment