[1tex]

Summary

transformers 5.15 (huggingface/transformers#47384) turned the Gemma 4 text config into a heterogeneous per-layer config:

• global_head_dim and num_global_key_value_heads are no longer attributes of Gemma4TextConfig; they are folded into per_layer_config for the full_attention layers (as head_dim / num_key_value_heads).
• Reading head_dim or num_key_value_heads on the global config raises AmbiguousGlobalPerLayerAttributeError (a RuntimeError, so hasattr does not catch it).

requirements.txt allows these versions (the bound went to <5.17.0 in #185 and was removed in #232), and with them loading any Gemma 4 checkpoint fails in hf_to_mcore_config:

from transformers import AutoConfig
from mcore_bridge.config import hf_to_mcore_config

hf_to_mcore_config(AutoConfig.from_pretrained('google/gemma-4-31B'))
# transformers 5.16:
# AmbiguousGlobalPerLayerAttributeError: 'num_key_value_heads' is a per-layer attribute and may vary across layers. ...

There is a second, silent issue behind it. Gemma4TextGPTModel._set_inv_freq passes head_dim_key='global_head_dim' for the proportional RoPE of the full-attention layers. In transformers>=5.15 the RoPE init function resolves per_layer_config['full_attention'] first, which has no global_head_dim, so it falls back to hidden_size // num_attention_heads (168 instead of 512 for the 31B model, i.e. 84 frequencies instead of 256).

Change

• config/parser.py
  • get_gemma4_attention_dims(text_config, layer_type) returns (head_dim, num_key_value_heads) for a layer type. With per-layer configs it reads text_config.per_layer_config[layer_type]; otherwise it keeps the existing global_head_dim / num_global_key_value_heads logic.
  • get_gemma4_rope_head_dim_key(text_config) returns 'head_dim' with per-layer configs and 'global_head_dim' otherwise.
  • _convert_config skips attributes listed in config.per_layer_attributes, so the generic mapping no longer touches values that have no single global meaning. For Gemma 4, kv_channels / num_query_groups are then set explicitly from the sliding layers, which is what the global values meant before.
• model/mm_gpts/gemma4.py: Gemma4SelfAttention, Gemma4Bridge._set_qkv, _set_inv_freq and unpack_pp_input use the helpers instead of reading global_head_dim / num_global_key_value_heads directly.

This does not set allow_global_per_layer_attribute_access; values are read from the layer they belong to. Behavior with transformers<5.15 is unchanged: the helpers take the old branch, and per_layer_attributes is empty there.

Verification

Tested on main at 8b567b3.

tests/test_gemma4_per_layer_config.py builds Gemma 4 31B and E2B text configs (real attention fields, 6 layers; 31B has attention_k_eq_v=True, E2B does not). Most cases load the parser functions from source like test_dsv4_mla_rope.py, so they run on CPU without Transformer Engine. They check:

• get_gemma4_attention_dims returns (256, 16) / (512, 4) for 31B sliding / full layers and (256, 1) / (512, 1) for E2B;
• _convert_config does not raise and does not report per-layer values;
• the full-attention RoPE inv_freq computed with get_gemma4_rope_head_dim_key equals HF's own Gemma4TextRotaryEmbedding(...).full_attention_inv_freq (256 entries for head_dim 512);
• hf_to_mcore_config returns the sliding-layer kv_channels / num_query_groups (skipped without Transformer Engine).

CPU, no Transformer Engine:
transformers 5.14.0 / 5.15.0 / 5.16.0: 4 passed, 2 skipped (each)

NGC PyTorch 26.05 container (Transformer Engine, megatron-core 0.18), transformers 5.16.1:
6 passed

On main the same configs give: 5.14 works; 5.15 and 5.16 raise AmbiguousGlobalPerLayerAttributeError in _convert_config, and the full-attention RoPE with head_dim_key='global_head_dim' has 84 frequencies instead of 256.

End to end on 8x H200 (Gemma 4 31B, ms-swift Megatron CPT, TP4/PP2, packing):

• transformers 5.16.1: the model loads and trains; on ms-swift 4.5.3 the first two losses are identical to an earlier run that used a local workaround (1.6207763, 1.61427021);
• transformers 5.14.0, same data: the first two losses are identical to the 5.16.1 run (1.69193351, 1.63562369).

(These runs also enabled the head_dim 512 attention kernel from a follow-up PR; the config path tested here is the same either way.)

flake8, isort and yapf (versions from .pre-commit-config.yaml) are clean on the changed files.