fix(gemma4): read per-layer attention dims for transformers>=5.15

src/mcore_bridge/config/parser.py

@@ -3,7 +3,7 @@
 from functools import partial
 from megatron.core.activations import squared_relu
 from transformers import PretrainedConfig
-from typing import Any, Dict
+from typing import Any, Dict, Tuple
 
 from mcore_bridge.utils import get_env_args
 
@@ -135,10 +135,40 @@
 }
 
 
+# transformers>=5.15 (huggingface/transformers#47384) keeps attributes that differ across layers in
+# `config.per_layer_config`; reading them on the global config raises AmbiguousGlobalPerLayerAttributeError.
+def has_per_layer_config(config) -> bool:
+    return getattr(config, 'per_layer_config', None) is not None
+
+
+def get_gemma4_attention_dims(text_config, layer_type: str) -> Tuple[int, int]:
+    """Return `(head_dim, num_key_value_heads)` used by Gemma 4 layers of `layer_type`."""
+    if has_per_layer_config(text_config):
+        layer_config = text_config.per_layer_config[layer_type]
+        return layer_config.head_dim, layer_config.num_key_value_heads
+    # transformers<5.15: full-attention layers read `global_head_dim` / `num_global_key_value_heads`.
+    is_sliding = layer_type == 'sliding_attention'
+    head_dim = text_config.global_head_dim if not is_sliding and text_config.global_head_dim else text_config.head_dim
+    if not is_sliding and text_config.attention_k_eq_v:
+        num_key_value_heads = text_config.num_global_key_value_heads
+    else:
+        num_key_value_heads = text_config.num_key_value_heads
+    return head_dim, num_key_value_heads
+
+
+def get_gemma4_rope_head_dim_key(text_config) -> str:
+    # transformers>=5.15 resolves `per_layer_config['full_attention']` in the RoPE init (global head_dim there).
+    return 'head_dim' if has_per_layer_config(text_config) else 'global_head_dim'
+
+
 def _convert_config(config, _internal_call=False) -> Dict[str, Any]:
     megatron_config = {}
+    # Per-layer attributes have no single global value; model-specific code reads them per layer.
+    per_layer_attributes = getattr(config, 'per_layer_attributes', None) or ()
     for k, hf_keys in config_mapping.items():
         for hf_k in hf_keys:
+            if hf_k in per_layer_attributes:
+                continue
             if hasattr(config, hf_k):
                 hf_v = getattr(config, hf_k)
                 if hf_v is None:
@@ -300,6 +330,9 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
     elif llm_model_type in {'ernie4_5', 'ernie4_5_moe', 'glm4'}:
         res['rotary_interleaved'] = True
     elif hf_model_type in {'gemma4', 'gemma4_unified'}:
+        # Global defaults follow the sliding layers; `Gemma4SelfAttention` overrides them for full-attention layers.
+        res['kv_channels'], res['num_query_groups'] = get_gemma4_attention_dims(hf_config.text_config,
+                                                                                'sliding_attention')
         res['qk_layernorm'] = True
         res['window_size'] = f'{window_size - 1},0'
         window_attn_skip_freq = ','.join(['1' if lt == 'sliding_attention' else '0' for lt in layer_types])

src/mcore_bridge/model/mm_gpts/gemma4.py

@@ -28,6 +28,7 @@
 
 from mcore_bridge.bridge import MultimodalGPTBridge
 from mcore_bridge.config import ModelConfig
+from mcore_bridge.config.parser import get_gemma4_attention_dims, get_gemma4_rope_head_dim_key
 from mcore_bridge.utils import get_logger
 
 from ..constant import ModelType
@@ -153,15 +154,10 @@ def __init__(
         self.layer_type = text_config.layer_types[layer_idx]
         self.is_sliding = self.layer_type == 'sliding_attention'
 
-        # Head dim: global layers may use a different head dim than sliding ones
-        self.head_dim = (
-            text_config.global_head_dim
-            if not self.is_sliding and text_config.global_head_dim else text_config.head_dim)
+        # Head dim / KV heads: global layers may differ from sliding ones
+        self.head_dim, self.num_key_value_heads = get_gemma4_attention_dims(text_config, self.layer_type)
 
         self.use_alternative_attention = (text_config.attention_k_eq_v and not self.is_sliding)
-        self.num_key_value_heads = (
-            text_config.num_global_key_value_heads
-            if self.use_alternative_attention else text_config.num_key_value_heads)
         # Shared KV across the trailing layers
         self.num_kv_shared_layers = getattr(text_config, 'num_kv_shared_layers', 0)
         if self.num_kv_shared_layers:
@@ -444,9 +440,9 @@ def _get_is_kv_shared_layer(self, layer_idx):
     def _set_qkv(self, mg_attn, hf_state_dict, to_mcore: bool, **kwargs):
         text_config = self.text_config
         layer_idx = kwargs['layer_idx']
-        is_sliding = text_config.layer_types[layer_idx] == 'sliding_attention'
-        head_dim = (
-            text_config.global_head_dim if not is_sliding and text_config.global_head_dim else text_config.head_dim)
+        layer_type = text_config.layer_types[layer_idx]
+        is_sliding = layer_type == 'sliding_attention'
+        head_dim, num_key_value_heads = get_gemma4_attention_dims(text_config, layer_type)
         is_kv_shared_layer = self._get_is_kv_shared_layer(layer_idx)
         if is_kv_shared_layer:
             self._set_state_dict(mg_attn, 'linear_qkv.weight', hf_state_dict, 'q_proj.weight', to_mcore)
@@ -455,9 +451,7 @@ def _set_qkv(self, mg_attn, hf_state_dict, to_mcore: bool, **kwargs):
             kwargs['kv_channels'] = head_dim
             use_alternative_attention = text_config.attention_k_eq_v and not is_sliding
             kwargs['attention_k_eq_v'] = use_alternative_attention
-            kwargs['num_query_groups'] = (
-                text_config.num_global_key_value_heads
-                if use_alternative_attention else text_config.num_key_value_heads)
+            kwargs['num_query_groups'] = num_key_value_heads
             return super()._set_qkv(mg_attn, hf_state_dict, to_mcore, **kwargs)
 
     def _set_router(self, mg_mlp, hf_state_dict, to_mcore, **kwargs):
@@ -584,7 +578,7 @@ def _set_inv_freq(self):
         self.config.rope_scaling = rope_scaling['full_attention']
         kwargs = {'layer_type': 'full_attention'}
         if self.config.rope_scaling['rope_type'] == 'proportional':
-            kwargs['head_dim_key'] = 'global_head_dim'
+            kwargs['head_dim_key'] = get_gemma4_rope_head_dim_key(self.config.hf_config.text_config)
         new_inv_freq, attention_scaling = get_rope_inv_freq(
             self.config, text_config=self.config.hf_config.text_config, **kwargs)
         self.full_rotary_pos_emb.inv_freq = new_inv_freq
@@ -680,7 +674,7 @@ def unpack_pp_input(self):
         per_layer_inputs_shape = [
             sequence_len, input_tensor.shape[1], self.config.num_layers, self.hidden_size_per_layer_input // tp_size
         ]
-        full_head_dim = self.text_config.global_head_dim
+        full_head_dim, _ = get_gemma4_attention_dims(self.text_config, 'full_attention')
         full_states_shape = per_layer_inputs_shape[:2] + [self.num_query_groups_per_partition, full_head_dim * 2]
         sliding_states_shape = full_states_shape[:3] + [self.config.kv_channels * 2]
         per_layer_inputs_dim = math.prod(per_layer_inputs_shape) // math.prod(input_tensor.shape[:2])

tests/test_gemma4_per_layer_config.py

@@ -0,0 +1,96 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Gemma 4 attention dims must be read the same way before and after transformers 5.15 (per-layer configs).
+The parser functions are loaded from source so most cases run on CPU without Transformer Engine."""
+import ast
+import pytest
+from pathlib import Path
+from typing import Any, Dict, Tuple
+
+gemma4_config = pytest.importorskip('transformers.models.gemma4.configuration_gemma4')
+
+_NAMES = {'config_mapping', 'has_per_layer_config', 'get_gemma4_attention_dims', 'get_gemma4_rope_head_dim_key'}
+
+
+def _load_parser():
+    path = Path(__file__).resolve().parents[1] / 'src/mcore_bridge/config/parser.py'
+    body = [
+        node for node in ast.parse(path.read_text()).body
+        if (isinstance(node, ast.FunctionDef) and node.name in _NAMES | {'_convert_config'}) or (
+            isinstance(node, ast.Assign) and getattr(node.targets[0], 'id', None) in _NAMES)
+    ]
+    namespace = {'Any': Any, 'Dict': Dict, 'Tuple': Tuple}
+    exec(compile(ast.fix_missing_locations(ast.Module(body=body, type_ignores=[])), str(path), 'exec'), namespace)
+    return namespace
+
+
+# Attention fields of the public checkpoints (fewer layers); expected (head_dim, kv_heads) for sliding / full.
+_CASES = {
+    'gemma-4-31B': (dict(
+        hidden_size=5376,
+        num_attention_heads=32,
+        num_key_value_heads=16,
+        num_global_key_value_heads=4,
+        attention_k_eq_v=True), (256, 16), (512, 4)),
+    'gemma-4-E2B': (dict(
+        hidden_size=1536,
+        num_attention_heads=8,
+        num_key_value_heads=1,
+        num_global_key_value_heads=None,
+        attention_k_eq_v=False), (256, 1), (512, 1)),
+}
+
+
+def _make_config(case):
+    rope = {
+        'full_attention': {
+            'partial_rotary_factor': 0.25,
+            'rope_theta': 1000000.0,
+            'rope_type': 'proportional'
+        },
+        'sliding_attention': {
+            'rope_theta': 10000.0,
+            'rope_type': 'default'
+        },
+    }
+    text_config = dict(
+        num_hidden_layers=6,
+        layer_types=['sliding_attention'] * 5 + ['full_attention'],
+        head_dim=256,
+        global_head_dim=512,
+        hidden_size_per_layer_input=0,
+        rope_parameters=rope,
+        **_CASES[case][0])
+    return gemma4_config.Gemma4Config(text_config=text_config)
+
+
+@pytest.mark.parametrize('case', list(_CASES))
+def test_attention_dims(case):
+    parser = _load_parser()
+    config = _make_config(case)
+    _, sliding, full = _CASES[case]
+    assert parser['get_gemma4_attention_dims'](config.text_config, 'sliding_attention') == sliding
+    assert parser['get_gemma4_attention_dims'](config.text_config, 'full_attention') == full
+    # The generic mapping must not read per-layer attributes (it raised on transformers>=5.15).
+    res = parser['_convert_config'](config)
+    assert res.get('kv_channels') in (None, sliding[0]) and res.get('num_query_groups') in (None, sliding[1])
+
+
+@pytest.mark.parametrize('case', list(_CASES))
+def test_full_attention_rope_matches_hf(case):
+    torch = pytest.importorskip('torch')
+    modeling = pytest.importorskip('transformers.models.gemma4.modeling_gemma4')
+    from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+    text_config = _make_config(case).text_config
+    head_dim_key = _load_parser()['get_gemma4_rope_head_dim_key'](text_config)
+    inv_freq, _ = ROPE_INIT_FUNCTIONS['proportional'](
+        text_config, 'cpu', layer_type='full_attention', head_dim_key=head_dim_key)
+    assert inv_freq.numel() == 512 // 2
+    torch.testing.assert_close(inv_freq, modeling.Gemma4TextRotaryEmbedding(text_config).full_attention_inv_freq)
+
+
+@pytest.mark.parametrize('case', list(_CASES))
+def test_hf_to_mcore_config(case):
+    pytest.importorskip('transformer_engine')
+    from mcore_bridge.config import hf_to_mcore_config
+    res = hf_to_mcore_config(_make_config(case))
+    assert (res['kv_channels'], res['num_query_groups']) == _CASES[case][1]