[bug fix] Fix 3 issues when using Gemma4 MTP (#26026)

This commit is contained in:
Charles Chen
2026-05-23 03:16:47 -07:00
committed by GitHub
parent b0ce16d0c5
commit 89ff2bc111
3 changed files with 20 additions and 11 deletions
+11 -6
View File
@@ -1147,7 +1147,8 @@ class Gemma4ForCausalLM(PreTrainedModel):
("experts.w13_weight", "experts.gate_up_proj", ("w1", "w3")), ("experts.w13_weight", "experts.gate_up_proj", ("w1", "w3")),
("experts.w2_weight", "experts.down_proj", ("w2",)), ("experts.w2_weight", "experts.down_proj", ("w2",)),
] ]
num_experts = self.config.num_experts # Dense subclasses (e.g. the Gemma4 MTP assistant) reuse this.
num_experts = getattr(self.config, "num_experts", None) or 0
# Per-expert checkpoint format used by compressed-tensors / FP8 # Per-expert checkpoint format used by compressed-tensors / FP8
# (e.g. RedHatAI/*-FP8-Dynamic) and by ModelOpt NVFP4 # (e.g. RedHatAI/*-FP8-Dynamic) and by ModelOpt NVFP4
@@ -1159,11 +1160,15 @@ class Gemma4ForCausalLM(PreTrainedModel):
# in a trailing dot, so the standard `name.replace(weight_name, # in a trailing dot, so the standard `name.replace(weight_name,
# param_name)` collapses every suffix uniformly to the fused # param_name)` collapses every suffix uniformly to the fused
# FusedMoE params (experts.w13_*, experts.w2_*). # FusedMoE params (experts.w13_*, experts.w2_*).
per_expert_params_mapping = FusedMoE.make_expert_params_mapping( per_expert_params_mapping = (
ckpt_gate_proj_name="gate_proj", FusedMoE.make_expert_params_mapping(
ckpt_down_proj_name="down_proj", ckpt_gate_proj_name="gate_proj",
ckpt_up_proj_name="up_proj", ckpt_down_proj_name="down_proj",
num_experts=num_experts, ckpt_up_proj_name="up_proj",
num_experts=num_experts,
)
if num_experts
else []
) )
k_eq_v_layers = self._get_k_eq_v_layers() k_eq_v_layers = self._get_k_eq_v_layers()
+2
View File
@@ -21,6 +21,7 @@ import torch
from torch import nn from torch import nn
from transformers import PretrainedConfig, PreTrainedModel from transformers import PretrainedConfig, PreTrainedModel
from sglang.srt.distributed import get_pp_group
from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.logits_processor import ( from sglang.srt.layers.logits_processor import (
LogitsMetadata, LogitsMetadata,
@@ -72,6 +73,7 @@ class Gemma4AssistantForCausalLM(Gemma4ForCausalLM):
self.assistant_config = config self.assistant_config = config
self.config = text_config self.config = text_config
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group()
self.vocab_size = text_config.vocab_size self.vocab_size = text_config.vocab_size
self.hidden_size = text_config.hidden_size self.hidden_size = text_config.hidden_size
+7 -5
View File
@@ -2277,11 +2277,13 @@ class ServerArgs:
) )
if is_sm100_supported() and self.moe_runner_backend == "auto": if is_sm100_supported() and self.moe_runner_backend == "auto":
if self.get_model_config().quantization == "modelopt_fp4":
self.moe_runner_backend = "flashinfer_trtllm" self.quantization = "modelopt_fp4"
logger.info( self.moe_runner_backend = "flashinfer_trtllm"
"Use flashinfer_trtllm as MoE runner backend on SM100 for Gemma-4 NVFP4" logger.info(
) "Use flashinfer_trtllm as MoE runner backend on "
"SM100 for Gemma-4 (modelopt_fp4)"
)
elif model_arch == "MossVLForConditionalGeneration": elif model_arch == "MossVLForConditionalGeneration":
if self.is_attention_backend_not_set(): if self.is_attention_backend_not_set():
self.prefill_attention_backend = "flashinfer" self.prefill_attention_backend = "flashinfer"