[Bug Fix] Align glm4_moe_nextn NPU MTP loading with qwen3 MTP (#25524)
This commit is contained in:
@@ -12,9 +12,8 @@
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
"""Inference-only GLM-4.5, GLM-4.6 Speculative Decoding."""
|
||||
"""Inference-only GLM-4.5, GLM-4.6 and GLM-4.7 Speculative Decoding."""
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
@@ -23,7 +22,6 @@ from torch import nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from sglang.srt.distributed import get_tensor_model_parallel_world_size
|
||||
from sglang.srt.environ import temp_set_env
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
|
||||
from sglang.srt.layers.layernorm import RMSNorm
|
||||
@@ -36,7 +34,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
||||
from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.utils import add_prefix
|
||||
from sglang.srt.utils import add_prefix, is_npu
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -51,7 +49,7 @@ class Glm4MoeModelNextN(nn.Module):
|
||||
super().__init__()
|
||||
if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
|
||||
logger.warning(
|
||||
"Overriding Glm4MoeForCausalLMNextN quant config for modelopt_fp4 GLM-4.5 / GLM-4.6 model."
|
||||
"Overriding Glm4MoeForCausalLMNextN quant config for modelopt_fp4 GLM-4.5 / GLM-4.6 / GLM-4.7 model."
|
||||
)
|
||||
quant_config = None
|
||||
|
||||
@@ -128,11 +126,13 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
|
||||
nn.Module.__init__(self)
|
||||
self.config = config
|
||||
self.tp_size = get_tensor_model_parallel_world_size()
|
||||
self.needs_quant_draft = (
|
||||
get_global_server_args().speculative_draft_model_quantization is not None
|
||||
or quant_config is not None
|
||||
)
|
||||
quant_config = quant_config if self.needs_quant_draft else None
|
||||
if (
|
||||
is_npu()
|
||||
and get_global_server_args().speculative_draft_model_quantization is None
|
||||
):
|
||||
quant_config = None
|
||||
self.quant_config = quant_config
|
||||
|
||||
self.model = Glm4MoeModelNextN(
|
||||
config, quant_config, prefix=add_prefix("model", prefix)
|
||||
)
|
||||
@@ -156,21 +156,7 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
) -> torch.Tensor:
|
||||
# Support unquant speculative draft model
|
||||
if self.needs_quant_draft:
|
||||
cxt = contextlib.nullcontext()
|
||||
else:
|
||||
# SGLANG_DEEPEP_BF16_DISPATCH is deprecated, will need
|
||||
# to be removed in the future and moved to a new
|
||||
# --deepep-dispatcher-output-dtype server argument.
|
||||
unquant_patch = {
|
||||
"SGLANG_DEEPEP_BF16_DISPATCH": "1",
|
||||
}
|
||||
cxt = temp_set_env(allow_sglang=True, **unquant_patch)
|
||||
|
||||
with cxt:
|
||||
hidden_states = self.model(input_ids, positions, forward_batch)
|
||||
|
||||
hidden_states = self.model(input_ids, positions, forward_batch)
|
||||
return self.logits_processor(
|
||||
input_ids, hidden_states, self.lm_head, forward_batch
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user