[NPU] fix model ERNIE-4.5-21B-A3B-PT bias need 1D error (#26038)
This commit is contained in:
@@ -42,9 +42,11 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
|||||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||||
from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4MLP
|
from sglang.srt.models.deepseek_v2 import DeepseekV2MLP as Ernie4MLP
|
||||||
from sglang.srt.models.llama import LlamaAttention as Ernie4Attention
|
from sglang.srt.models.llama import LlamaAttention as Ernie4Attention
|
||||||
from sglang.srt.utils import add_prefix, make_layers
|
from sglang.srt.utils import add_prefix, is_npu, make_layers
|
||||||
from sglang.srt.utils.hf_transformers_utils import get_rope_config
|
from sglang.srt.utils.hf_transformers_utils import get_rope_config
|
||||||
|
|
||||||
|
_is_npu = is_npu()
|
||||||
|
|
||||||
|
|
||||||
class MoEGate(nn.Module):
|
class MoEGate(nn.Module):
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -86,12 +88,16 @@ class Ernie4Moe(nn.Module):
|
|||||||
|
|
||||||
self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix))
|
self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix))
|
||||||
|
|
||||||
|
correction_bias = self.gate.e_score_correction_bias
|
||||||
|
# npu only supports 1D, but current correction_bias is 2D
|
||||||
|
if _is_npu:
|
||||||
|
correction_bias = correction_bias.squeeze(0)
|
||||||
self.topk = TopK(
|
self.topk = TopK(
|
||||||
top_k=config.moe_k,
|
top_k=config.moe_k,
|
||||||
layer_id=layer_id,
|
layer_id=layer_id,
|
||||||
renormalize=True,
|
renormalize=True,
|
||||||
use_grouped_topk=False,
|
use_grouped_topk=False,
|
||||||
correction_bias=self.gate.e_score_correction_bias,
|
correction_bias=correction_bias,
|
||||||
)
|
)
|
||||||
|
|
||||||
self.experts = get_moe_impl_class(quant_config)(
|
self.experts = get_moe_impl_class(quant_config)(
|
||||||
|
|||||||
@@ -222,6 +222,7 @@ class LlamaAttention(nn.Module):
|
|||||||
self.q_size,
|
self.q_size,
|
||||||
self.kv_size,
|
self.kv_size,
|
||||||
self.head_dim,
|
self.head_dim,
|
||||||
|
is_neox_style=self.rotary_emb.is_neox_style,
|
||||||
)
|
)
|
||||||
return q, k, v
|
return q, k, v
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user