[bug fix] Fix 3 issues when using Gemma4 MTP (#26026)
This commit is contained in:
@@ -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()
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|||||||
Reference in New Issue
Block a user