[model] Support gemma4

src/mcore_bridge/bridge/gpt_bridge.py

@@ -40,9 +40,12 @@ class GPTBridge:
     hf_o_proj_key = 'o_proj'
     hf_attn_prefix = 'self_attn'
     hf_mlp_prefix = 'mlp'
+    hf_post_attention_layernorm = 'post_attention_layernorm'
     hf_gate_key = 'gate.weight'
     hf_shared_expert_key = None
     hf_expert_bias_key = 'gate.e_score_correction_bias'
+    additional_dim0_keys = set()
+    additional_dim1_keys = set()
 
     def __init__(self, config: ModelConfig):
         self.config = config
@@ -124,11 +127,11 @@ def _get_tp_split_dim(self, mg_key: Optional[str]) -> Optional[int]:
             'linear_kv_up_proj',
             # mtp
             'eh_proj',
-        }
+        } & self.additional_dim0_keys
         if self.config.task_type in {'causal_lm', 'generative_reranker'}:
             dim0_keys.add('output_layer')
         # RowLinear
-        dim1_keys = {'out_proj', 'linear_proj', 'linear_fc2'}
+        dim1_keys = {'out_proj', 'linear_proj', 'linear_fc2'} & self.additional_dim1_keys
         if 'lora_A' not in mg_key and 'lora_B' not in mg_key:
             key, suffix = mg_key.rsplit('.', 2)[-2:]
             if suffix == 'layer_norm_weight':
@@ -439,7 +442,8 @@ def _set_state_dict(self,
                         to_mcore: bool,
                         *,
                         offset: float = 0,
-                        is_expert: bool = False):
+                        is_expert: bool = False,
+                        _check_mg_param: bool = True):
         if '.' in mg_key:
             module_key, param_key = mg_key.rsplit('.', 1)
         else:
@@ -487,7 +491,11 @@ def _set_state_dict(self,
             else:
                 mg_param = deep_getattr(sub_module, param_key)
             if to_mcore:
-                assert mg_param is not None, f'mg_module: {mg_module}, mg_key: {mg_key}'
+                if mg_param is None:
+                    if _check_mg_param:
+                        raise ValueError(f'mg_module: {mg_module}, mg_key: {mg_key}')
+                    else:
+                        return
                 hf_weight = hf_state_dict[hf_key].load()
                 if module_key in {
                         'embedding.word_embeddings', 'output_layer'
@@ -1587,13 +1595,13 @@ def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool
             hf_state_dict.update(
                 self._set_moe_state(
                     mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore, is_mtp=is_mtp))
-            self._set_state_dict(mg_layer, 'pre_mlp_layernorm.weight', hf_state_dict, 'post_attention_layernorm.weight',
-                                 to_mcore)
+            self._set_state_dict(mg_layer, 'pre_mlp_layernorm.weight', hf_state_dict,
+                                 f'{self.hf_post_attention_layernorm}.weight', to_mcore)
         else:
             hf_state_dict.update(
                 self._set_mlp_state(mg_mlp, hf_state_dict, f'{self.hf_mlp_prefix}.', layer_idx, to_mcore))
             self._set_state_dict(mg_layer, 'mlp.linear_fc1.layer_norm_weight', hf_state_dict,
-                                 'post_attention_layernorm.weight', to_mcore)
+                                 f'{self.hf_post_attention_layernorm}.weight', to_mcore)
         return hf_state_dict
 
     def _set_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: int, to_mcore: bool):
@@ -1610,13 +1618,16 @@ def _set_layer_state(self, mg_layer, hf_state_dict, hf_prefix: str, layer_idx: i
             hf_state_dict = self._add_prefix(hf_state_dict, hf_prefix)
         return hf_state_dict
 
+    def _set_word_embeddings(self, mg_model, hf_state_dict, to_mcore):
+        lm_model = getattr(mg_model, 'language_model') if self.is_multimodal else mg_model
+        self._set_state_dict(lm_model, 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key, to_mcore)
+
     def _convert_pre_process(self, mg_model, hf_state_dict, hf_prefix: str, to_mcore):
         if to_mcore:
             hf_state_dict = self._remove_prefix(hf_state_dict, hf_prefix)
         else:
             hf_state_dict = {}
-        lm_model = getattr(mg_model, 'language_model') if self.is_multimodal else mg_model
-        self._set_state_dict(lm_model, 'embedding.word_embeddings.weight', hf_state_dict, self.hf_embed_key, to_mcore)
+        self._set_word_embeddings(mg_model, hf_state_dict, to_mcore)
         if self.is_multimodal:
             for prefix, mg_prefix in self.module_mapping.items():
                 mg_module = deep_getattr(mg_model, f'visual.{mg_prefix}')

src/mcore_bridge/config/parser.py

@@ -149,6 +149,8 @@ def hf_to_mcore_config(hf_config: PretrainedConfig) -> Dict[str, Any]:
         n_shared_experts = res.pop('n_shared_experts')
     elif llm_model_type in {'ernie4_5', 'ernie4_5_moe', 'glm4'}:
         res['rotary_interleaved'] = True
+    elif hf_model_type in {'gemma4'}:
+        res['qk_layernorm'] = True
     elif llm_model_type == 'gpt_oss':
         res['add_bias_linear'] = True
         res['bias_dropout_fusion'] = False

src/mcore_bridge/model/constant.py

@@ -28,6 +28,7 @@ class MLLMModelType:
     glm4v_moe = 'glm4v_moe'
     kimi_vl = 'kimi_vl'
     llama4 = 'llama4'
+    gemma4 = 'gemma4'
 
     kimi_k25 = 'kimi_k25'
 

src/mcore_bridge/model/gpt_model.py

@@ -100,9 +100,7 @@ def __init__(
             for i in range(len(self.decoder.layers)):
                 if hasattr(self.decoder.layers[i].self_attention, 'rotary_pos_emb'):
                     del self.decoder.layers[i].self_attention.rotary_pos_emb
-        self.attention_scaling = 1.
-        new_inv_freq, self.attention_scaling = get_rope_inv_freq(config)
-        self.rotary_pos_emb.inv_freq = new_inv_freq.to(self.rotary_pos_emb.inv_freq.device)
+        self._set_inv_freq()
         if self.config.task_type == 'seq_cls' and self.post_process:
             self.output_layer = OutputLayerLinear(
                 config.hidden_size,
@@ -212,7 +210,36 @@ def _preprocess(
         if decoder_input is not None and self.training and torch.is_grad_enabled() and not decoder_input.requires_grad:
             # fix LoRA incompatibility with gradient checkpointing
             decoder_input = decoder_input.requires_grad_(True)
+        rotary_pos_emb, decoder_rotary_pos_emb, rotary_pos_cos, rotary_pos_sin = self._get_rotary_pos_emb(
+            decoder_input, position_ids, packed_seq_params=packed_seq_params)
+
+        if (in_inference_mode and ((self.config.enable_cuda_graph and self.config.cuda_graph_scope != 'full_iteration')
+                                   or self.config.flash_decode) and rotary_pos_cos is not None
+                and inference_context.is_static_batching()):
+            current_batch_size = input_ids.shape[0]
+            sequence_len_offset = torch.tensor(
+                [inference_context.sequence_len_offset] * current_batch_size,
+                dtype=torch.int32,
+                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
+            )
+        else:
+            sequence_len_offset = None
 
+        # Wrap decoder_input to allow the decoder (TransformerBlock) to delete the
+        # reference held by this caller function, enabling early garbage collection for
+        # inference. Skip wrapping if decoder_input is logged after decoder completion.
+        if in_inference_mode and not has_config_logger_enabled(self.config):
+            decoder_input = WrappedTensor(decoder_input)
+
+        return (decoder_input, rotary_pos_emb, decoder_rotary_pos_emb, rotary_pos_cos, rotary_pos_sin,
+                sequence_len_offset)
+
+    def _set_inv_freq(self):
+        self.attention_scaling = 1.
+        new_inv_freq, self.attention_scaling = get_rope_inv_freq(self.config)
+        self.rotary_pos_emb.inv_freq = new_inv_freq.to(self.rotary_pos_emb.inv_freq.device)
+
+    def _get_rotary_pos_emb(self, decoder_input, position_ids, packed_seq_params, inference_context=None):
         # Rotary positional embeddings (embedding is None for PP intermediate devices)
         rotary_pos_emb = None
         rotary_pos_cos = None
@@ -247,26 +274,13 @@ def _preprocess(
                         rotary_seq_len,
                         packed_seq=packed_seq,
                     )
+        decoder_rotary_pos_emb = rotary_pos_emb
+        packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'
+        if self.position_embedding_type == 'rope' and packed_seq and not self.config.apply_rope_fusion:
+            assert position_ids.shape[0] == 1, f'position_ids.shape: {position_ids.shape}'
+            decoder_rotary_pos_emb = rotary_pos_emb[position_ids[0]]
 
-        if (in_inference_mode and ((self.config.enable_cuda_graph and self.config.cuda_graph_scope != 'full_iteration')
-                                   or self.config.flash_decode) and rotary_pos_cos is not None
-                and inference_context.is_static_batching()):
-            current_batch_size = input_ids.shape[0]
-            sequence_len_offset = torch.tensor(
-                [inference_context.sequence_len_offset] * current_batch_size,
-                dtype=torch.int32,
-                device=rotary_pos_cos.device,  # Co-locate this with the rotary tensors
-            )
-        else:
-            sequence_len_offset = None
-
-        # Wrap decoder_input to allow the decoder (TransformerBlock) to delete the
-        # reference held by this caller function, enabling early garbage collection for
-        # inference. Skip wrapping if decoder_input is logged after decoder completion.
-        if in_inference_mode and not has_config_logger_enabled(self.config):
-            decoder_input = WrappedTensor(decoder_input)
-
-        return decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset
+        return rotary_pos_emb, decoder_rotary_pos_emb, rotary_pos_cos, rotary_pos_sin
 
     # Code borrowed from NVIDIA/Megatron-LM
     def forward(
@@ -296,20 +310,15 @@ def forward(
         """
 
         inference_context = deprecate_inference_params(inference_context, inference_params)
-
-        decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
+        # There is a difference in whether rotary_pos_emb can be fused between the decoder and MTP.
+        decoder_input, rotary_pos_emb, decoder_rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset = (
             self._preprocess(
                 input_ids=input_ids,
                 position_ids=position_ids,
                 decoder_input=decoder_input,
                 inference_context=inference_context,
                 packed_seq_params=packed_seq_params,
             ))
-        decoder_rotary_pos_emb = rotary_pos_emb
-        packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'
-        if self.position_embedding_type == 'rope' and packed_seq and not self.config.apply_rope_fusion:
-            assert position_ids.shape[0] == 1, f'position_ids.shape: {position_ids.shape}'
-            decoder_rotary_pos_emb = rotary_pos_emb[position_ids[0]]
 
         mtp_decoder_input = decoder_input
         if self.config.is_multimodal and self.config.mtp_num_layers and decoder_input is None:

src/mcore_bridge/model/gpts/minimax_m2.py

@@ -27,9 +27,11 @@ def __init__(
         k_layernorm = submodules.k_layernorm
         submodules.q_layernorm = IdentityOp
         submodules.k_layernorm = IdentityOp
-        super().__init__(config, submodules, *args, **kwargs)
-        submodules.q_layernorm = q_layernorm
-        submodules.k_layernorm = k_layernorm
+        try:
+            super().__init__(config, submodules, *args, **kwargs)
+        finally:
+            submodules.q_layernorm = q_layernorm
+            submodules.k_layernorm = k_layernorm
         self.q_norm = build_module(
             submodules.q_layernorm,
             hidden_size=self.hidden_size_per_attention_head * config.num_attention_heads,

src/mcore_bridge/model/mm_gpts/__init__.py

@@ -1,2 +1,2 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from . import glm, internvl, kimi_vl, llama4, llava, qwen, qwen3_5, qwen3_5_gdn, qwen3_omni, qwen3_vl
+from . import gemma4, glm, internvl, kimi_vl, llama4, llava, qwen, qwen3_5, qwen3_5_gdn, qwen3_omni, qwen3_vl