From 81df7b44247a38d87800413ad508575621866758 Mon Sep 17 00:00:00 2001 From: xingmingyyj Date: Fri, 26 Jun 2026 22:11:21 +0800 Subject: [PATCH 1/3] add gemma4 model --- paddleformers/datasets/template/template.py | 2 + .../transformers/auto/configuration.py | 9 + paddleformers/transformers/auto/modeling.py | 1 + .../transformers/gemma4_moe/__init__.py | 35 ++ .../transformers/gemma4_moe/configuration.py | 136 +++++ .../transformers/gemma4_moe/modeling.py | 504 ++++++++++++++++++ 6 files changed, 687 insertions(+) create mode 100644 paddleformers/transformers/gemma4_moe/__init__.py create mode 100644 paddleformers/transformers/gemma4_moe/configuration.py create mode 100644 paddleformers/transformers/gemma4_moe/modeling.py diff --git a/paddleformers/datasets/template/template.py b/paddleformers/datasets/template/template.py index 6d0d42cd39c..7eeb70ea7eb 100644 --- a/paddleformers/datasets/template/template.py +++ b/paddleformers/datasets/template/template.py @@ -995,3 +995,5 @@ def _get_gpt_oss_prefix(): chat_sep="<|assistant|>\n", mm_plugin=get_mm_plugin(name="glm_ocr", image_token="<|image|>"), ) + +# TODO(xingmingyyj) add template for Gemma4 diff --git a/paddleformers/transformers/auto/configuration.py b/paddleformers/transformers/auto/configuration.py index c04e1f34a5a..08c6ae54811 100644 --- a/paddleformers/transformers/auto/configuration.py +++ b/paddleformers/transformers/auto/configuration.py @@ -63,6 +63,9 @@ ("glm_ocr", "GlmOcrConfig"), ("qwen3_5", "Qwen3_5Config"), ("qwen3_5_moe", "Qwen3_5MoEConfig"), + # TODO(VL): When Gemma4 VL is implemented, "gemma4" should point to Gemma4Config (VL wrapper) + ("gemma4_text", "Gemma4MoeConfig"), + ("gemma4", "Gemma4MoeConfig"), # Temporary: no standalone text ckpt, extract text_config in from_dict ] ) @@ -91,6 +94,9 @@ ("glm_ocr", "GlmOcrForConditionalGeneration"), ("qwen3_5_moe", "Qwen3_5MoEForConditionalGeneration"), ("qwen3_5", "Qwen3_5ForConditionalGeneration"), + ("gemma4_moe", "Gemma4MoeForCausalLM"), + ("gemma4_text", "Gemma4MoeForCausalLM"), + ("gemma4", "Gemma4MoeForCausalLM"), ] ) @@ -104,6 +110,9 @@ ("qwen2_5_vl_text", "qwen2_5_vl"), ("qwen3_vl_text", "qwen3_vl"), ("qwen3_vl_moe_text", "qwen3_vl_moe"), + # TODO(VL): Remove these when Gemma4 VL module (gemma4/) is created + ("gemma4_text", "gemma4_moe"), + ("gemma4", "gemma4_moe"), ] ) diff --git a/paddleformers/transformers/auto/modeling.py b/paddleformers/transformers/auto/modeling.py index 11321baba1f..2e1f6deaa98 100644 --- a/paddleformers/transformers/auto/modeling.py +++ b/paddleformers/transformers/auto/modeling.py @@ -78,6 +78,7 @@ ("GptOss", "gpt_oss"), ("Phi3", "phi3"), ("Gemma3", "gemma3_text"), + ("Gemma4Moe", "gemma4_moe"), ("Glm4vMoe", "glm4v_moe"), ("GlmOcr", "glm_ocr"), ] diff --git a/paddleformers/transformers/gemma4_moe/__init__.py b/paddleformers/transformers/gemma4_moe/__init__.py new file mode 100644 index 00000000000..883173c3ec1 --- /dev/null +++ b/paddleformers/transformers/gemma4_moe/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import sys +from typing import TYPE_CHECKING + +from ...utils.lazy_import import _LazyModule + +import_structure = { + "configuration": ["Gemma4MoeConfig"], + "modeling": [ + "Gemma4MoeForCausalLM", + ], +} + +if TYPE_CHECKING: + from .configuration import * + from .modeling import * +else: + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + import_structure, + module_spec=__spec__, + ) diff --git a/paddleformers/transformers/gemma4_moe/configuration.py b/paddleformers/transformers/gemma4_moe/configuration.py new file mode 100644 index 00000000000..1912c11338b --- /dev/null +++ b/paddleformers/transformers/gemma4_moe/configuration.py @@ -0,0 +1,136 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from ..configuration_utils import PretrainedConfig + + +class Gemma4MoeConfig(PretrainedConfig): + """Gemma4 26B-A4B text-only MoE model configuration. + + NOTE: Gemma4 currently only has VL checkpoints (Gemma4ForConditionalGeneration), + no standalone text-only checkpoint exists. The config.json has top-level + model_type="gemma4" (VL wrapper) with text params nested in text_config + (model_type="gemma4_text"). + + Both "gemma4" and "gemma4_text" are registered in CONFIG_MAPPING to this class, + and from_dict() automatically extracts text_config from the VL wrapper. + + TODO(VL): When implementing full multimodal Gemma4: + 1. Create Gemma4Config (VL wrapper with text_config + vision_config) + 2. Change CONFIG_MAPPING "gemma4" to point to Gemma4Config + 3. Keep only "gemma4_text" -> Gemma4MoeConfig + 4. Remove the text_config extraction logic in from_dict() below + """ + + model_type = "gemma4_text" + keys_to_ignore_at_inference = ["past_key_values"] + + @classmethod + def from_dict(cls, config_dict, **kwargs): + """Load text config from VL wrapper config. + + Gemma4 only ships VL checkpoints, so config.json outer model_type="gemma4" + with text params in text_config. This extracts it automatically. + + TODO(VL): Move this extraction to Gemma4Config when VL model is implemented. + """ + if config_dict.get("model_type") == "gemma4" and "text_config" in config_dict: + config_dict = config_dict["text_config"] + return super().from_dict(config_dict, **kwargs) + + def __init__( + self, + vocab_size=262144, + hidden_size=2816, + intermediate_size=2112, + moe_intermediate_size=704, + num_hidden_layers=30, + num_attention_heads=16, + num_key_value_heads=8, + num_global_key_value_heads=2, + head_dim=256, + global_head_dim=512, + hidden_act="gelu_pytorch_tanh", + hidden_activation=None, + max_position_embeddings=262144, + rms_norm_eps=1e-6, + # RoPE + rope_parameters=None, + sliding_window_rope_base=10000.0, + full_attention_rope_base=1000000.0, + full_attention_rope_partial_factor=0.25, + # Attention pattern + layer_types=None, + sliding_window=1024, + interleaved_attn_pattern=(5, 1), + # MoE + num_experts=128, + top_k_experts=8, + scoring_func="sigmoid", + enable_moe_block=True, + # Gemma4-specific + attention_k_eq_v=True, + final_logit_softcapping=30.0, + scale_embeddings_by_hidden_size=True, + # Pipeline + pp_seg_method="layer:Gemma4TransformerLayer", + tie_word_embeddings=True, + **kwargs, + ): + super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.num_global_key_value_heads = num_global_key_value_heads + self.head_dim = head_dim + self.global_head_dim = global_head_dim + self.hidden_act = hidden_activation or hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.enable_moe_block = enable_moe_block + self.sliding_window = sliding_window + self.interleaved_attn_pattern = interleaved_attn_pattern + self.num_experts = num_experts + self.top_k_experts = top_k_experts + self.scoring_func = scoring_func + self.attention_k_eq_v = attention_k_eq_v + self.final_logit_softcapping = final_logit_softcapping + self.scale_embeddings_by_hidden_size = scale_embeddings_by_hidden_size + self.pp_seg_method = pp_seg_method + + # Parse rope_parameters from HF config format + if rope_parameters is not None: + sliding_rope = rope_parameters.get("sliding_attention", {}) + full_rope = rope_parameters.get("full_attention", {}) + self.sliding_window_rope_base = sliding_rope.get("rope_theta", sliding_window_rope_base) + self.full_attention_rope_base = full_rope.get("rope_theta", full_attention_rope_base) + self.full_attention_rope_partial_factor = full_rope.get( + "partial_rotary_factor", full_attention_rope_partial_factor + ) + else: + self.sliding_window_rope_base = sliding_window_rope_base + self.full_attention_rope_base = full_attention_rope_base + self.full_attention_rope_partial_factor = full_attention_rope_partial_factor + + # Build layer_types from interleaved pattern if not provided + if layer_types is None: + sliding_count, global_count = self.interleaved_attn_pattern + pattern = ["sliding_attention"] * sliding_count + ["full_attention"] * global_count + self.layer_types = (pattern * ((num_hidden_layers // len(pattern)) + 1))[:num_hidden_layers] + else: + self.layer_types = layer_types diff --git a/paddleformers/transformers/gemma4_moe/modeling.py b/paddleformers/transformers/gemma4_moe/modeling.py new file mode 100644 index 00000000000..766cf20d5d3 --- /dev/null +++ b/paddleformers/transformers/gemma4_moe/modeling.py @@ -0,0 +1,504 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Gemma4 MoE model provider and ForCausalLM entry. + +""" +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Callable + +import paddle + +from paddleformers.transformers.gpt_provider import GPTModelProvider +from paddleformers.transformers.model_utils import PretrainedModel + +logger = logging.getLogger(__name__) + +from paddlefleet.models.gpt.gemma4_layer_specs import Gemma4DualRotaryEmbedding +from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec +from paddlefleet.transformer.transformer_layer import Gemma4TransformerLayer + +from .configuration import Gemma4MoeConfig + + +def _patch_embedding_scale(embedding_layer, embed_scale): + """Monkey-patch LanguageModelEmbedding.forward to multiply by embed_scale. + + This is needed because PipelineLayer stores GPTEmbedding under numeric + keys, making it impossible to replace the embedding via attribute assignment. + Instead, we patch the forward method of the inner LanguageModelEmbedding. + """ + import types + + orig_forward = embedding_layer.forward.__func__ + + def _scaled_forward(self_inner, input_ids, position_ids, tokentype_ids=None): + result = orig_forward(self_inner, input_ids, position_ids, tokentype_ids) + return result * embed_scale + + embedding_layer.forward = types.MethodType(_scaled_forward, embedding_layer) + + +class Gemma4MoePreTrainedModel(PretrainedModel): + config_class = Gemma4MoeConfig + base_model_prefix = "gemma4_moe" + + +@dataclass +class Gemma4MoeModelProvider(GPTModelProvider): + """Provider for Gemma4 MoE model. Aligns with Megatron Gemma4ModelProvider.""" + + # Override defaults for Gemma4 26B-A4B + num_layers: int = 30 + hidden_size: int = 2816 + ffn_hidden_size: int = 2112 + num_attention_heads: int = 16 + num_query_groups: int = 8 + kv_channels: int = 256 + + # Gemma4-specific + global_head_dim: int = 512 + num_global_key_value_heads: int = 2 + layer_types: list = None + + # MoE + n_routed_experts: int = 128 + num_moe_experts: int = 128 + moe_router_topk: int = 8 + moe_ffn_hidden_size: int = 704 + moe_shared_expert_intermediate_size: int = 2112 + moe_token_dispatcher_type: str = "alltoall" + moe_grouped_gemm: bool = True + moe_layer_freq: int = 1 + scoring_func: str = "sigmoid" + + # RoPE + rotary_base: float = 10000 + sliding_window_rope_base: float = 10000.0 + full_attention_rope_base: float = 1000000.0 + global_rotary_percent: float = 0.25 + rotary_percent: float = 1.0 + rope_scaling: object = None + + # Model structure + share_embeddings_and_output_weights: bool = True + normalization: str = "RMSNorm" + layernorm_epsilon: float = 1e-6 + rms_norm_eps: float = 1e-6 + gated_linear_unit: bool = True + activation_func: str = "gelu_pytorch_tanh" + attention_k_eq_v: bool = True + final_logit_softcapping: float = 30.0 + scale_embeddings_by_hidden_size: bool = True + add_swa_attention_sink_bias: bool = False + add_full_attention_sink_bias: bool = False + + # Layer spec + transformer_layer_spec: Callable = None + + transform_rules = { + **GPTModelProvider.transform_rules, + "dtype": "params_dtype", + "num_experts": "n_routed_experts", + "top_k_experts": "num_experts_per_tok", + "num_hidden_layers": "num_layers", + "num_key_value_heads": "num_query_groups", + "head_dim": "kv_channels", + "moe_ffn_hidden_size": "moe_intermediate_size", + } + + def __post_init__(self): + # Set head_dim and num_key_value_heads to sliding-layer base values + # BEFORE super().__post_init__(), so TransformerConfig doesn't default them + # to hidden_size//num_heads and num_attention_heads. + # Global layers override these per-layer in Gemma4SelfAttention.__init__. + self.head_dim = self.kv_channels # 256 (sliding) + self.num_key_value_heads = self.num_query_groups # 8 (sliding) + + super().__post_init__() + if self.transformer_layer_spec is None: + self.transformer_layer_spec = self._get_decoder_layers_spec + if not hasattr(self, "num_experts_per_tok") or self.num_experts_per_tok == 2: + self.num_experts_per_tok = self.moe_router_topk + # Gemma4 controls shared expert via moe_shared_expert_intermediate_size directly. + # MoELayer needs n_shared_experts > 0 to create shared_experts. + if not getattr(self, "n_shared_experts", None): + self.n_shared_experts = 1 + # Sync num_hidden_layers from num_layers (transform_rules maps HF num_hidden_layers→num_layers) + if self.num_hidden_layers == 1 and self.num_layers > 1: + self.num_hidden_layers = self.num_layers + # Convert sliding_window from int (HF config) to tuple (PaddleFleet expects tuple[int, int]) + sw = getattr(self, "sliding_window", None) + if isinstance(sw, int): + self.sliding_window = (sw, 0) + + def _get_decoder_layers_spec(self, config): + """Generate layer specs for all Gemma4 layers via standard GPT path.""" + config.specific_layer = Gemma4TransformerLayer + num_layers = getattr(config, "num_layers", 30) + return [ + get_gpt_layer_local_spec( + config=config, + num_experts=None, + use_qk_norm=True, + normalization=getattr(config, "normalization", "RMSNorm"), + layer_number=i + 1, + attention_layer_type="gemma4", + ) + for i in range(num_layers) + ] + + def provide(self, pre_process=None, post_process=None, vp_stage=None, loss_fn=None): + """Build Gemma4 model using standard GPT spec path with gemma4 attention type.""" + from paddle.distributed.fleet.meta_parallel import LayerSpec, build_spec_layer + from paddlefleet.models.common.empty_layer import EmptyLayer + from paddlefleet.models.common.language_loss.language_loss import LanguageLoss + from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_spec + + # Build layers via standard get_gpt_layer_local_spec path + transformer_layers_spec = self._get_decoder_layers_spec(self) + + head_empty_layers_spec = [ + LayerSpec(layer=EmptyLayer, extra_kwargs={"config": self}) + for _ in range(self.num_empty_layers_add_in_head) + ] + tail_empty_layers_spec = [ + LayerSpec(layer=EmptyLayer, extra_kwargs={"config": self}) + for _ in range(self.num_empty_layers_add_in_tail) + ] + + gpt_spec = get_gpt_spec( + config=self, + transformer_layers_spec=transformer_layers_spec, + mtp_layers_spec=None, + vocab_size=self.vocab_size, + head_empty_layers_spec=head_empty_layers_spec, + tail_empty_layers_spec=tail_empty_layers_spec, + max_sequence_length=self.max_sequence_length, + position_embedding_type=self.position_embedding_type, + rotary_percent=self.rotary_percent, + rotary_base=self.rotary_base, + swa_rotary_base=getattr(self, "swa_rope_theta", None), + rope_scaling=self.rope_scaling, + parallel_output=self.parallel_output, + tie_word_embeddings=self.tie_word_embeddings, + ) + + pp_size = self.pipeline_model_parallel_size + fleet_model = build_spec_layer( + gpt_spec, + loss_fn=loss_fn if loss_fn else LanguageLoss(self), + num_stages=pp_size, + seg_method="layer:Gemma4TransformerLayer|EmptyLayer", + ) + + # Convert FleetLayer GPTModel → PaddleFormers GPTModel (PretrainedModel) + from paddleformers.transformers.gpt_provider import GPTModel + + model = GPTModel.__new__(GPTModel) + for attr_name in dir(fleet_model): + if not attr_name.startswith("__"): + try: + setattr(model, attr_name, getattr(fleet_model, attr_name)) + except: + pass + + # TODO(xingmingyyj) Support context parallel + # Replace RoPE with Dual RoPE. + # rotary_pos_emb lives inside GPTEmbedding (model.embedding.rotary_pos_emb), + # NOT as a top-level model attribute. + has_emb = hasattr(model, "embedding") + has_rpe_top = hasattr(model, "rotary_pos_emb") + has_rpe_emb = has_emb and hasattr(model.embedding, "rotary_pos_emb") + logger.info( + f"[Gemma4] RoPE replacement check: has_embedding={has_emb}, " + f"has_model.rotary_pos_emb={has_rpe_top}, " + f"has_model.embedding.rotary_pos_emb={has_rpe_emb}" + ) + if has_rpe_emb: + old_rpe = model.embedding.rotary_pos_emb + model.embedding.rotary_pos_emb = Gemma4DualRotaryEmbedding(self) + logger.info( + f"[Gemma4] Replaced model.embedding.rotary_pos_emb: " + f"{type(old_rpe).__name__} -> {type(model.embedding.rotary_pos_emb).__name__}" + ) + elif has_rpe_top: + old_rpe = model.rotary_pos_emb + model.rotary_pos_emb = Gemma4DualRotaryEmbedding(self) + logger.info( + f"[Gemma4] Replaced model.rotary_pos_emb: " + f"{type(old_rpe).__name__} -> {type(model.rotary_pos_emb).__name__}" + ) + else: + # Fallback: search sublayers for GPTEmbedding with rotary_pos_emb + logger.warning( + "[Gemma4] Could not find rotary_pos_emb via top-level or model.embedding. " "Searching sublayers..." + ) + for name, sublayer in model.named_sublayers(): + if hasattr(sublayer, "rotary_pos_emb") and sublayer.rotary_pos_emb is not None: + old_rpe = sublayer.rotary_pos_emb + sublayer.rotary_pos_emb = Gemma4DualRotaryEmbedding(self) + logger.info( + f"[Gemma4] Replaced {name}.rotary_pos_emb: " + f"{type(old_rpe).__name__} -> {type(sublayer.rotary_pos_emb).__name__}" + ) + break + else: + logger.error("[Gemma4] FAILED to find any rotary_pos_emb to replace!") + + # Logit Softcapping: patch GPTLMHead._forward to apply tanh softcapping. + # Cannot use Gemma4OutputLayer wrapper because GPTModel (PipelineLayer) + # stores the LM head in its internal pipeline registry, not as a direct + # `output_layer` attribute. Instead, find the GPTLMHead sublayer and + # monkey-patch its _forward method. + if self.final_logit_softcapping > 0: + import types + + from paddlefleet.models.gpt.lm_head import GPTLMHead + + softcap = self.final_logit_softcapping + + def _make_softcapped_forward(orig_fwd, cap): + def _forward_with_softcap(self_inner, hidden_states): + result = orig_fwd(self_inner, hidden_states) + if isinstance(result, tuple): + self_inner.config.fused_linear_ce_loss_chunk = 0 + result = orig_fwd(self_inner, hidden_states) + return paddle.tanh(result / cap) * cap + + return _forward_with_softcap + + for name, sublayer in model.named_sublayers(): + if isinstance(sublayer, GPTLMHead): + orig_fwd = sublayer._forward.__func__ + sublayer._forward = types.MethodType(_make_softcapped_forward(orig_fwd, softcap), sublayer) + break + + # Embedding scale: Gemma4 multiplies embeddings by sqrt(hidden_size). + # GPTModel (PipelineLayer) stores GPTEmbedding under numeric keys in _sub_layers, + # NOT as model.embedding. Must search via named_sublayers(). + if self.scale_embeddings_by_hidden_size: + embed_scale = self.hidden_size**0.5 + found_emb = False + + # Try direct access first (works if PipelineLayer exposes it) + if hasattr(model, "embedding"): + gpt_emb = model.embedding + if hasattr(gpt_emb, "embedding"): + _patch_embedding_scale(gpt_emb.embedding, embed_scale) + found_emb = True + logger.info(f"[Gemma4] Applied embedding scale √{self.hidden_size} via model.embedding") + + # Fallback: search sublayers for GPTEmbedding with .embedding + if not found_emb: + from paddlefleet.models.gpt.gpt_embedding import GPTEmbedding + + for name, sublayer in model.named_sublayers(): + if isinstance(sublayer, GPTEmbedding) and hasattr(sublayer, "embedding"): + _patch_embedding_scale(sublayer.embedding, embed_scale) + found_emb = True + logger.info(f"[Gemma4] Applied embedding scale √{self.hidden_size} via sublayer {name}") + break + + if not found_emb: + logger.error("[Gemma4] FAILED to find embedding layer for √hidden_size scaling!") + + return model + + +class Gemma4MoeForCausalLM(Gemma4MoePreTrainedModel): + """Gemma4 MoE ForCausalLM entry using PaddleFleet spec mode.""" + + @classmethod + def _gen_aoa_config(cls, config): + model_prefix = "model." + num_hidden_layers = config.num_hidden_layers + num_head_empty_layers = ( + config.num_empty_layers_add_in_head + if hasattr(config, "num_empty_layers_add_in_head") and config.num_empty_layers_add_in_head + else 0 + ) + aoa_config = { + "aoa_statements": [ + f"model.language_model.norm.weight -> {model_prefix}norm.weight", + f"model.language_model.embed_tokens.weight -> {model_prefix}embedding.embed_tokens.weight", + ] + } + if config.tie_word_embeddings: + aoa_config["aoa_statements"].append( + f"model.language_model.embed_tokens.weight -> {model_prefix}lm_head.weight" + ) + for layer_idx in range(num_hidden_layers): + lo = layer_idx + num_head_empty_layers + hf = f"model.language_model.layers.{layer_idx}" + pf = f"{model_prefix}layers.{lo}" + # Heterogeneous attention: global layers have different kv_heads + layer_types = getattr(config, "layer_types", None) + is_global = layer_types is not None and layer_types[layer_idx] == "full_attention" + kv_heads = ( + getattr(config, "num_global_key_value_heads", config.num_key_value_heads) + if is_global + else config.num_key_value_heads + ) + aoa_config["aoa_statements"] += [ + f"{hf}.input_layernorm.weight -> {pf}.input_layernorm.weight", + f"{hf}.post_attention_layernorm.weight -> {pf}.post_self_attn_layernorm.weight", + f"{hf}.pre_feedforward_layernorm.weight -> {pf}.pre_mlp_layernorm.weight", + f"{hf}.post_feedforward_layernorm.weight -> {pf}.post_mlp_layernorm.weight", + f"{hf}.layer_scalar -> {pf}.layer_scalar, dtype='float32'", + ] + # Attention: fused QKV + # Global layers have K=V tying: HF checkpoint has no v_proj, use k_proj as V + if is_global: + aoa_config["aoa_statements"].append( + f"{hf}.self_attn.q_proj.weight^T, {hf}.self_attn.k_proj.weight^T, {hf}.self_attn.k_proj.weight^T -> {pf}.self_attn.qkv_proj.weight, fused_qkv, num_heads={config.num_attention_heads}, num_key_value_groups={kv_heads}" + ) + else: + aoa_config["aoa_statements"].append( + f"{hf}.self_attn.q_proj.weight^T, {hf}.self_attn.k_proj.weight^T, {hf}.self_attn.v_proj.weight^T -> {pf}.self_attn.qkv_proj.weight, fused_qkv, num_heads={config.num_attention_heads}, num_key_value_groups={kv_heads}" + ) + aoa_config["aoa_statements"] += [ + f"{hf}.self_attn.o_proj.weight^T -> {pf}.self_attn.o_proj.weight", + f"{hf}.self_attn.q_norm.weight -> {pf}.self_attn.q_norm.weight", + f"{hf}.self_attn.k_norm.weight -> {pf}.self_attn.k_norm.weight", + f"{hf}.mlp.gate_proj.weight^T, {hf}.mlp.up_proj.weight^T -> {pf}.mlp.shared_experts.up_gate_proj.weight, fused_ffn", + f"{hf}.mlp.down_proj.weight^T -> {pf}.mlp.shared_experts.down_proj.weight", + f"{hf}.post_feedforward_layernorm_1.weight -> {pf}.mlp.post_shared_expert_layernorm.weight", + f"{hf}.pre_feedforward_layernorm_2.weight -> {pf}.mlp.pre_feedforward_layernorm_2.weight", + f"{hf}.post_feedforward_layernorm_2.weight -> {pf}.mlp.post_moe_layernorm.weight", + f"{hf}.router.proj.weight -> {pf}.mlp.gate.weight, dtype='float32'", + f"{hf}.router.per_expert_scale -> {pf}.mlp.gate.per_expert_scale, dtype='float32'", + f"{hf}.router.scale -> {pf}.mlp.gate.router_input_scale, dtype='float32'", + ] + # Routed experts + # HF: experts.gate_up_proj [E, I*2, H] -> PF: grouped_gemm_experts.weight1 [E, H, I*2] + # HF: experts.down_proj [E, H, I] -> PF: grouped_gemm_experts.weight2 [E, I, H] + # Both need permute="[0,2,1]" (swap last two dims within each expert slice) + # TODO(xingmingyyj) add assert + aoa_config["aoa_statements"] += [ + f"{hf}.experts.gate_up_proj -> {pf}.mlp.grouped_gemm_experts.weight1, permute='[0,2,1]'", + f"{hf}.experts.down_proj -> {pf}.mlp.grouped_gemm_experts.weight2, permute='[0,2,1]'", + ] + return aoa_config + + @classmethod + def _gen_inv_aoa_config(cls, config): + """PF -> HF weight mapping for saving checkpoints.""" + model_prefix = "model." + num_hidden_layers = config.num_hidden_layers + num_head_empty_layers = ( + config.num_empty_layers_add_in_head + if hasattr(config, "num_empty_layers_add_in_head") and config.num_empty_layers_add_in_head + else 0 + ) + aoa_statements = [ + f"{model_prefix}norm.weight -> model.language_model.norm.weight", + f"{model_prefix}embedding.embed_tokens.weight -> model.language_model.embed_tokens.weight", + ] + if config.tie_word_embeddings: + aoa_statements.append(f"{model_prefix}lm_head.weight -> _") + + for layer_idx in range(num_hidden_layers): + lo = layer_idx + num_head_empty_layers + hf = f"model.language_model.layers.{layer_idx}" + pf = f"{model_prefix}layers.{lo}" + layer_types = getattr(config, "layer_types", None) + is_global = layer_types is not None and layer_types[layer_idx] == "full_attention" + kv_heads = ( + getattr(config, "num_global_key_value_heads", config.num_key_value_heads) + if is_global + else config.num_key_value_heads + ) + + # Norms + aoa_statements += [ + f"{pf}.input_layernorm.weight -> {hf}.input_layernorm.weight", + f"{pf}.post_self_attn_layernorm.weight -> {hf}.post_attention_layernorm.weight", + f"{pf}.pre_mlp_layernorm.weight -> {hf}.pre_feedforward_layernorm.weight", + f"{pf}.post_mlp_layernorm.weight -> {hf}.post_feedforward_layernorm.weight", + ] + + # layer_scalar + aoa_statements.append(f"{pf}.layer_scalar -> {hf}.layer_scalar") + + # Attention: qkv_proj -> split q/k/v + transpose + # Global layers (K=V tying): HF has no v_proj, skip v output + aoa_statements.append( + f"{pf}.self_attn.qkv_proj.weight -> {pf}.self_attn.q_proj.weight, {pf}.self_attn.k_proj.weight, {pf}.self_attn.v_proj.weight, fused_qkv, num_heads={config.num_attention_heads}, num_key_value_groups={kv_heads}" + ) + aoa_statements += [ + f"{pf}.self_attn.q_proj.weight^T -> {hf}.self_attn.q_proj.weight", + f"{pf}.self_attn.k_proj.weight^T -> {hf}.self_attn.k_proj.weight", + ] + if not is_global: + aoa_statements.append(f"{pf}.self_attn.v_proj.weight^T -> {hf}.self_attn.v_proj.weight") + aoa_statements += [ + f"{pf}.self_attn.o_proj.weight^T -> {hf}.self_attn.o_proj.weight", + f"{pf}.self_attn.q_norm.weight -> {hf}.self_attn.q_norm.weight", + f"{pf}.self_attn.k_norm.weight -> {hf}.self_attn.k_norm.weight", + ] + + # Shared expert: up_gate_proj -> gate + up + transpose + aoa_statements += [ + f"{pf}.mlp.shared_experts.up_gate_proj.weight -> {pf}.mlp.shared_experts.gate_proj.weight, {pf}.mlp.shared_experts.up_proj.weight, fused_ffn", + f"{pf}.mlp.shared_experts.gate_proj.weight^T -> {hf}.mlp.gate_proj.weight", + f"{pf}.mlp.shared_experts.up_proj.weight^T -> {hf}.mlp.up_proj.weight", + f"{pf}.mlp.shared_experts.down_proj.weight^T -> {hf}.mlp.down_proj.weight", + ] + + # MoE norms + aoa_statements += [ + f"{pf}.mlp.post_shared_expert_layernorm.weight -> {hf}.post_feedforward_layernorm_1.weight", + f"{pf}.mlp.pre_feedforward_layernorm_2.weight -> {hf}.pre_feedforward_layernorm_2.weight", + f"{pf}.mlp.post_moe_layernorm.weight -> {hf}.post_feedforward_layernorm_2.weight", + ] + + # Router + aoa_statements += [ + f"{pf}.mlp.gate.weight -> {hf}.router.proj.weight, dtype='bfloat16'", + f"{pf}.mlp.gate.per_expert_scale -> {hf}.router.per_expert_scale", + f"{pf}.mlp.gate.router_input_scale -> {hf}.router.scale", + ] + + # Routed experts (inverse) + # PF: grouped_gemm_experts.weight1 [E, H, I*2] -> HF: experts.gate_up_proj [E, I*2, H] + # PF: grouped_gemm_experts.weight2 [E, I, H] -> HF: experts.down_proj [E, H, I] + # TODO(xingmingyyj) add assert + aoa_statements += [ + f"{pf}.mlp.grouped_gemm_experts.weight1 -> {hf}.experts.gate_up_proj, permute='[0,2,1]'", + f"{pf}.mlp.grouped_gemm_experts.weight2 -> {hf}.experts.down_proj, permute='[0,2,1]'", + ] + + return {"aoa_statements": aoa_statements} + + is_fleet = True + + def __new__(cls, config): + config.tensor_model_parallel_size = max(getattr(config, "tensor_model_parallel_size", 1), 1) + config.pipeline_model_parallel_size = max(getattr(config, "pipeline_model_parallel_size", 1), 1) + config.expert_model_parallel_size = max(getattr(config, "expert_model_parallel_size", 1), 1) + config.context_parallel_size = max(getattr(config, "context_parallel_size", 1), 1) + config.virtual_pipeline_model_parallel_size = max( + getattr(config, "virtual_pipeline_model_parallel_size", 1), 1 + ) + + model_provider = Gemma4MoeModelProvider.from_config(config) + gpt_model = model_provider.provide() + gpt_model._gen_aoa_config = cls._gen_aoa_config + gpt_model._gen_inv_aoa_config = cls._gen_inv_aoa_config + gpt_model.config_to_save = config + gpt_model.is_fleet = cls.is_fleet + return gpt_model From 987685290b22aaad9dc2e4790f0778e98299c806 Mon Sep 17 00:00:00 2001 From: xingmingyyj Date: Thu, 9 Jul 2026 17:42:03 +0800 Subject: [PATCH 2/3] fix --- paddleformers/datasets/template/template.py | 15 ++++++++++++++- paddleformers/transformers/gemma4_moe/modeling.py | 10 +++++----- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/paddleformers/datasets/template/template.py b/paddleformers/datasets/template/template.py index 7eeb70ea7eb..ad82221d7d4 100644 --- a/paddleformers/datasets/template/template.py +++ b/paddleformers/datasets/template/template.py @@ -996,4 +996,17 @@ def _get_gpt_oss_prefix(): mm_plugin=get_mm_plugin(name="glm_ocr", image_token="<|image|>"), ) -# TODO(xingmingyyj) add template for Gemma4 + +register_template( + name="gemma4", + format_user=StringFormatter(slots=["<|turn>user\n{{content}}\n<|turn>model\n"]), + format_assistant=StringFormatter(slots=["{{content}}"]), + format_system=StringFormatter(slots=["<|turn>system\n{{content}}\n"]), + format_observation=StringFormatter(slots=["<|turn>tool\n{{content}}\n<|turn>model\n"]), + format_prefix=EmptyFormatter(slots=[{"bos_token"}]), + chat_sep="\n", + suffix=[""], + stop_words=[""], + thought_words=("<|channel>thought\n", "\n"), + template_class=Llama2Template, +) diff --git a/paddleformers/transformers/gemma4_moe/modeling.py b/paddleformers/transformers/gemma4_moe/modeling.py index 766cf20d5d3..96b639f74c5 100644 --- a/paddleformers/transformers/gemma4_moe/modeling.py +++ b/paddleformers/transformers/gemma4_moe/modeling.py @@ -155,7 +155,7 @@ def _get_decoder_layers_spec(self, config): num_experts=None, use_qk_norm=True, normalization=getattr(config, "normalization", "RMSNorm"), - layer_number=i + 1, + layer_number=i, attention_layer_type="gemma4", ) for i in range(num_layers) @@ -380,7 +380,7 @@ def _gen_aoa_config(cls, config): f"{hf}.pre_feedforward_layernorm_2.weight -> {pf}.mlp.pre_feedforward_layernorm_2.weight", f"{hf}.post_feedforward_layernorm_2.weight -> {pf}.mlp.post_moe_layernorm.weight", f"{hf}.router.proj.weight -> {pf}.mlp.gate.weight, dtype='float32'", - f"{hf}.router.per_expert_scale -> {pf}.mlp.gate.per_expert_scale, dtype='float32'", + f"{hf}.router.per_expert_scale -> {pf}.mlp.gate.routed_scaling_factor_param, dtype='float32'", f"{hf}.router.scale -> {pf}.mlp.gate.router_input_scale, dtype='float32'", ] # Routed experts @@ -432,7 +432,7 @@ def _gen_inv_aoa_config(cls, config): ] # layer_scalar - aoa_statements.append(f"{pf}.layer_scalar -> {hf}.layer_scalar") + aoa_statements.append(f"{pf}.layer_scalar -> {hf}.layer_scalar, dtype='bfloat16'") # Attention: qkv_proj -> split q/k/v + transpose # Global layers (K=V tying): HF has no v_proj, skip v output @@ -469,8 +469,8 @@ def _gen_inv_aoa_config(cls, config): # Router aoa_statements += [ f"{pf}.mlp.gate.weight -> {hf}.router.proj.weight, dtype='bfloat16'", - f"{pf}.mlp.gate.per_expert_scale -> {hf}.router.per_expert_scale", - f"{pf}.mlp.gate.router_input_scale -> {hf}.router.scale", + f"{pf}.mlp.gate.routed_scaling_factor_param -> {hf}.router.per_expert_scale, dtype='bfloat16'", + f"{pf}.mlp.gate.router_input_scale -> {hf}.router.scale, dtype='bfloat16'", ] # Routed experts (inverse) From 10be004bec5f3a647f2fb70d0ac8d25bca46e03c Mon Sep 17 00:00:00 2001 From: xingmingyyj Date: Thu, 9 Jul 2026 21:58:16 +0800 Subject: [PATCH 3/3] fix --- paddleformers/transformers/__init__.py | 4 ++++ paddleformers/transformers/gemma4_moe/modeling.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/paddleformers/transformers/__init__.py b/paddleformers/transformers/__init__.py index 554336bbd9d..184368c00c9 100644 --- a/paddleformers/transformers/__init__.py +++ b/paddleformers/transformers/__init__.py @@ -348,6 +348,9 @@ ], "glm_ocr.processor": ["Glm46VProcessor"], "glm_ocr.image_processor": ["Glm46VImageProcessor"], + "gemma4_moe.configuration": ["Gemma4MoeConfig"], + "gemma4_moe.modeling": ["Gemma4MoeForCausalLM"], + "gemma4_moe": [], "phi4.configuration": ["Phi4Config"], "phi4.modeling": ["Phi4Model", "Phi4ForCausalLM"], "phi4.tokenizer": ["Phi4Tokenizer"], @@ -430,6 +433,7 @@ from .phi3 import * from .gemma3_text import * from .glm_ocr import * + from .gemma4_moe import * from .phi4 import * else: sys.modules[__name__] = _LazyModule( diff --git a/paddleformers/transformers/gemma4_moe/modeling.py b/paddleformers/transformers/gemma4_moe/modeling.py index 96b639f74c5..1cc291a677f 100644 --- a/paddleformers/transformers/gemma4_moe/modeling.py +++ b/paddleformers/transformers/gemma4_moe/modeling.py @@ -27,7 +27,7 @@ logger = logging.getLogger(__name__) -from paddlefleet.models.gpt.gemma4_layer_specs import Gemma4DualRotaryEmbedding +from paddlefleet.models.common.embeddings import Gemma4DualRotaryEmbedding from paddlefleet.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec from paddlefleet.transformer.transformer_layer import Gemma4TransformerLayer