【NPU】【bugfix】fix server error when mtp unquant (#26389)
Co-authored-by: cen121212 <luochen23@huawei.com> Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
This commit is contained in:
cen121212
co-authored by
cen121212
Even Zhou
cen121212
Even Zhou
parent
282c46133f
commit
b421e60eed
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union
|
||||
@@ -437,11 +436,8 @@ class _DeepEPDispatcherImplBase:
|
||||
# NVFP4 is supported on GPU, no adjustment needed
|
||||
|
||||
def _update_int8_quant_env(self) -> None:
|
||||
"""Update the DEEP_NORMAL_MODE_USE_INT8_QUANT environment variable."""
|
||||
if self.use_fp8:
|
||||
os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1"
|
||||
else:
|
||||
os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0"
|
||||
"""TODO adapt different quantization schemes for base model and draft model on NPU"""
|
||||
pass
|
||||
|
||||
def set_overlap_args(
|
||||
self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
|
||||
import logging
|
||||
import os
|
||||
from contextlib import ExitStack
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
@@ -169,70 +170,87 @@ class DeepseekModelNextN(nn.Module):
|
||||
forward_batch: ForwardBatch,
|
||||
input_embeds: torch.Tensor = None,
|
||||
) -> torch.Tensor:
|
||||
zero_allocator = BumpAllocator(
|
||||
buffer_size=2,
|
||||
dtype=torch.float32,
|
||||
device=(
|
||||
input_embeds.device if input_embeds is not None else input_ids.device
|
||||
),
|
||||
)
|
||||
exit_stack = ExitStack()
|
||||
if (
|
||||
_is_npu
|
||||
and self.quant_config is None
|
||||
and get_global_server_args().quantization is not None
|
||||
):
|
||||
# ascend mtp unquant
|
||||
exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
|
||||
exit_stack.enter_context(
|
||||
envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
|
||||
)
|
||||
|
||||
if input_embeds is None:
|
||||
hidden_states = self.embed_tokens(input_ids)
|
||||
else:
|
||||
hidden_states = input_embeds
|
||||
|
||||
if hidden_states.shape[0] > 0:
|
||||
eh_input = torch.cat(
|
||||
(
|
||||
self.enorm(hidden_states),
|
||||
self.hnorm(
|
||||
forward_batch.spec_info.hidden_states
|
||||
if self.rot_weight is None
|
||||
else torch.matmul(
|
||||
forward_batch.spec_info.hidden_states, self.rot_weight
|
||||
)
|
||||
),
|
||||
try:
|
||||
zero_allocator = BumpAllocator(
|
||||
buffer_size=2,
|
||||
dtype=torch.float32,
|
||||
device=(
|
||||
input_embeds.device
|
||||
if input_embeds is not None
|
||||
else input_ids.device
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
if isinstance(self.eh_proj, ReplicatedLinear):
|
||||
hidden_states, _ = self.eh_proj(eh_input)
|
||||
else:
|
||||
hidden_states = self.eh_proj(eh_input)
|
||||
|
||||
if dsa_use_prefill_cp(
|
||||
forward_batch, self.dsa_enable_prefill_cp
|
||||
) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp):
|
||||
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
|
||||
positions = cp_split_and_rebuild_position(forward_batch, positions)
|
||||
residual = None
|
||||
with get_global_expert_distribution_recorder().disable_this_region():
|
||||
hidden_states, residual, topk_indices = self.decoder(
|
||||
positions,
|
||||
hidden_states,
|
||||
forward_batch,
|
||||
residual,
|
||||
zero_allocator,
|
||||
)
|
||||
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
if residual is not None:
|
||||
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
|
||||
if input_embeds is None:
|
||||
hidden_states = self.embed_tokens(input_ids)
|
||||
else:
|
||||
hidden_states = self.shared_head.norm(hidden_states)
|
||||
hidden_states = input_embeds
|
||||
|
||||
if hidden_states.shape[0] > 0:
|
||||
eh_input = torch.cat(
|
||||
(
|
||||
self.enorm(hidden_states),
|
||||
self.hnorm(
|
||||
forward_batch.spec_info.hidden_states
|
||||
if self.rot_weight is None
|
||||
else torch.matmul(
|
||||
forward_batch.spec_info.hidden_states, self.rot_weight
|
||||
)
|
||||
),
|
||||
),
|
||||
dim=-1,
|
||||
)
|
||||
if isinstance(self.eh_proj, ReplicatedLinear):
|
||||
hidden_states, _ = self.eh_proj(eh_input)
|
||||
else:
|
||||
hidden_states = self.eh_proj(eh_input)
|
||||
|
||||
if dsa_use_prefill_cp(
|
||||
forward_batch, self.dsa_enable_prefill_cp
|
||||
) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp):
|
||||
# allgather + rerrange
|
||||
hidden_states = cp_all_gather_rerange_output(
|
||||
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
|
||||
positions = cp_split_and_rebuild_position(forward_batch, positions)
|
||||
residual = None
|
||||
with get_global_expert_distribution_recorder().disable_this_region():
|
||||
hidden_states, residual, topk_indices = self.decoder(
|
||||
positions,
|
||||
hidden_states,
|
||||
self.cp_size,
|
||||
forward_batch,
|
||||
torch.cuda.current_stream(),
|
||||
residual,
|
||||
zero_allocator,
|
||||
)
|
||||
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
if residual is not None:
|
||||
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
|
||||
else:
|
||||
hidden_states = self.shared_head.norm(hidden_states)
|
||||
|
||||
if dsa_use_prefill_cp(
|
||||
forward_batch, self.dsa_enable_prefill_cp
|
||||
) or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp):
|
||||
# allgather + rerrange
|
||||
hidden_states = cp_all_gather_rerange_output(
|
||||
hidden_states,
|
||||
self.cp_size,
|
||||
forward_batch,
|
||||
torch.cuda.current_stream(),
|
||||
)
|
||||
finally:
|
||||
exit_stack.close()
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from contextlib import ExitStack
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
@@ -23,6 +24,7 @@ from torch import nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
|
||||
from sglang.srt.layers.layernorm import GemmaRMSNorm
|
||||
@@ -140,38 +142,55 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
|
||||
input_embeds: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
assert input_embeds is None
|
||||
input_embeds = forward_batch.mm_input_embeds
|
||||
exit_stack = ExitStack()
|
||||
if (
|
||||
forward_batch.forward_mode.is_extend()
|
||||
and forward_batch.contains_mm_inputs()
|
||||
and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
|
||||
is_npu()
|
||||
and self.quant_config is None
|
||||
and get_global_server_args().quantization is not None
|
||||
):
|
||||
assert input_embeds is not None
|
||||
input_embeds = torch.cat(
|
||||
[input_embeds[:-1], self.model.embed_tokens(input_ids[-1].unsqueeze(0))]
|
||||
# ascend mtp unquant
|
||||
exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
|
||||
exit_stack.enter_context(
|
||||
envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
|
||||
)
|
||||
|
||||
if input_embeds is None:
|
||||
input_embeds = self.model.embed_tokens(input_ids)
|
||||
try:
|
||||
assert input_embeds is None
|
||||
input_embeds = forward_batch.mm_input_embeds
|
||||
if (
|
||||
forward_batch.forward_mode.is_extend()
|
||||
and forward_batch.contains_mm_inputs()
|
||||
and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
|
||||
):
|
||||
assert input_embeds is not None
|
||||
input_embeds = torch.cat(
|
||||
[
|
||||
input_embeds[:-1],
|
||||
self.model.embed_tokens(input_ids[-1].unsqueeze(0)),
|
||||
]
|
||||
)
|
||||
|
||||
hidden_states = forward_batch.spec_info.hidden_states
|
||||
if input_embeds is None:
|
||||
input_embeds = self.model.embed_tokens(input_ids)
|
||||
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
input_embeds = self.pre_fc_norm_embedding(input_embeds)
|
||||
hidden_states = self.pre_fc_norm_hidden(hidden_states)
|
||||
hidden_states = torch.cat([input_embeds, hidden_states], dim=-1)
|
||||
hidden_states = forward_batch.spec_info.hidden_states
|
||||
|
||||
hidden_states = self.fc(hidden_states)
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
input_embeds = self.pre_fc_norm_embedding(input_embeds)
|
||||
hidden_states = self.pre_fc_norm_hidden(hidden_states)
|
||||
hidden_states = torch.cat([input_embeds, hidden_states], dim=-1)
|
||||
|
||||
with get_global_expert_distribution_recorder().disable_this_region():
|
||||
hidden_states = self.model(
|
||||
input_ids,
|
||||
positions,
|
||||
forward_batch,
|
||||
hidden_states,
|
||||
)
|
||||
hidden_states = self.fc(hidden_states)
|
||||
|
||||
with get_global_expert_distribution_recorder().disable_this_region():
|
||||
hidden_states = self.model(
|
||||
input_ids,
|
||||
positions,
|
||||
forward_batch,
|
||||
hidden_states,
|
||||
)
|
||||
finally:
|
||||
exit_stack.close()
|
||||
|
||||
return self.logits_processor(
|
||||
input_ids, hidden_states, self.lm_head, forward_batch
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from contextlib import ExitStack
|
||||
from typing import Iterable, Optional, Tuple
|
||||
|
||||
import torch
|
||||
@@ -23,6 +24,7 @@ from torch import nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
|
||||
from sglang.srt.layers.layernorm import GemmaRMSNorm
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
@@ -94,25 +96,39 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
|
||||
input_embeds: Optional[torch.Tensor] = None,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
if input_embeds is None:
|
||||
input_embeds = self.model.embed_tokens(input_ids)
|
||||
|
||||
hidden_states = forward_batch.spec_info.hidden_states
|
||||
# Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0.
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
input_embeds = self.pre_fc_norm_embedding(input_embeds)
|
||||
hidden_states = self.pre_fc_norm_hidden(hidden_states)
|
||||
hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
|
||||
|
||||
with get_global_expert_distribution_recorder().disable_this_region():
|
||||
hidden_states = self.model(
|
||||
input_ids,
|
||||
positions,
|
||||
forward_batch,
|
||||
hidden_states,
|
||||
exit_stack = ExitStack()
|
||||
if (
|
||||
is_npu()
|
||||
and self.quant_config is None
|
||||
and get_global_server_args().quantization is not None
|
||||
):
|
||||
# ascend mtp unquant
|
||||
exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
|
||||
exit_stack.enter_context(
|
||||
envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
|
||||
)
|
||||
|
||||
try:
|
||||
if input_embeds is None:
|
||||
input_embeds = self.model.embed_tokens(input_ids)
|
||||
|
||||
hidden_states = forward_batch.spec_info.hidden_states
|
||||
# Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0.
|
||||
if not forward_batch.forward_mode.is_idle():
|
||||
input_embeds = self.pre_fc_norm_embedding(input_embeds)
|
||||
hidden_states = self.pre_fc_norm_hidden(hidden_states)
|
||||
hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))
|
||||
|
||||
with get_global_expert_distribution_recorder().disable_this_region():
|
||||
hidden_states = self.model(
|
||||
input_ids,
|
||||
positions,
|
||||
forward_batch,
|
||||
hidden_states,
|
||||
)
|
||||
finally:
|
||||
exit_stack.close()
|
||||
|
||||
return self.logits_processor(
|
||||
input_ids, hidden_states, self.lm_head, forward_batch
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user