diff --git a/src/mcore_bridge/config/parser.py b/src/mcore_bridge/config/parser.py index 23aae0a..2a56215 100644 --- a/src/mcore_bridge/config/parser.py +++ b/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]) diff --git a/src/mcore_bridge/model/mm_gpts/gemma4.py b/src/mcore_bridge/model/mm_gpts/gemma4.py index 21624ee..468af07 100644 --- a/src/mcore_bridge/model/mm_gpts/gemma4.py +++ b/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]) diff --git a/tests/test_gemma4_per_layer_config.py b/tests/test_gemma4_per_layer_config.py new file mode 100644 index 0000000..0edaa0e --- /dev/null +++ b/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]