[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.
|
# 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
|
import logging
|
||||||
from typing import Iterable, Optional, Tuple
|
from typing import Iterable, Optional, Tuple
|
||||||
|
|
||||||
@@ -23,7 +22,6 @@ from torch import nn
|
|||||||
from transformers import PretrainedConfig
|
from transformers import PretrainedConfig
|
||||||
|
|
||||||
from sglang.srt.distributed import get_tensor_model_parallel_world_size
|
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.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||||
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
|
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
|
||||||
from sglang.srt.layers.layernorm import RMSNorm
|
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.model_executor.forward_batch_info import ForwardBatch
|
||||||
from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM
|
from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM
|
||||||
from sglang.srt.server_args import get_global_server_args
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -51,7 +49,7 @@ class Glm4MoeModelNextN(nn.Module):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
|
if quant_config is not None and quant_config.get_name() == "modelopt_fp4":
|
||||||
logger.warning(
|
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
|
quant_config = None
|
||||||
|
|
||||||
@@ -128,11 +126,13 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
|
|||||||
nn.Module.__init__(self)
|
nn.Module.__init__(self)
|
||||||
self.config = config
|
self.config = config
|
||||||
self.tp_size = get_tensor_model_parallel_world_size()
|
self.tp_size = get_tensor_model_parallel_world_size()
|
||||||
self.needs_quant_draft = (
|
if (
|
||||||
get_global_server_args().speculative_draft_model_quantization is not None
|
is_npu()
|
||||||
or quant_config is not None
|
and get_global_server_args().speculative_draft_model_quantization is None
|
||||||
)
|
):
|
||||||
quant_config = quant_config if self.needs_quant_draft else None
|
quant_config = None
|
||||||
|
self.quant_config = quant_config
|
||||||
|
|
||||||
self.model = Glm4MoeModelNextN(
|
self.model = Glm4MoeModelNextN(
|
||||||
config, quant_config, prefix=add_prefix("model", prefix)
|
config, quant_config, prefix=add_prefix("model", prefix)
|
||||||
)
|
)
|
||||||
@@ -156,21 +156,7 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
|
|||||||
positions: torch.Tensor,
|
positions: torch.Tensor,
|
||||||
forward_batch: ForwardBatch,
|
forward_batch: ForwardBatch,
|
||||||
) -> torch.Tensor:
|
) -> 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(
|
return self.logits_processor(
|
||||||
input_ids, hidden_states, self.lm_head, forward_batch
|
input_ids, hidden_states, self.lm_head, forward_batch
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user