openPangu-Embedded（https://ai.gitcode.com/ascend-tribe/openpangu-embedded-7b-model/）是基于昇腾NPU从零训练的高效大语言模型，参数量为7B（不含词表Embedding）。openPangu-Embedded-7B训练了约19T tokens，具备快慢思考融合能力。本文主要介绍openpangu-embedded-7b-model模型在MindIE上的部署适配。
git clone https://gitee.com/ascend/MindIE-LLM.git创建模型目录位置：
cd MindIE-LLM/examples/atb_models/atb_llm/models/
mkdir pangu进入模型目录分别创建router_pangu.py、input_builder_pangu.py、tools_call_processor_pangu.py
cd pangu
touch router_pangu.py input_builder_pangu.py tools_call_processor_pangu.pyPARunner会调用
{llm_path}/atb_llm/models/{model}/router_{model}.py中的{model}Router来进行模型初始化及配置文件加载，其中*{model}*为模型名称，必须与模型配置文件中的“model_type”严格保持一致。
{model}Router类起到路由作用，用于告知模型从何处加载模型及对应的配置文件。
{model}Router类继承自基类BaseRouter，在模型迁移适配时，该类需实现“get_input_builder”和“get_tokenizer”方法。
“get_input_builder”用于在推理过程中套用prompt模板。
“get_tokenizer”用于初始化tokenizer。
openpangu的适配情况如下：
from dataclasses import dataclass
from .tools_call_processor_pangu import ToolsCallProcessorPangu
from .input_builder_pangu import PanguInputBuilder
from ..base.model_utils import safe_get_tokenizer_from_pretrained
from ..base.router import BaseRouter
@dataclass
class PanguRouter(BaseRouter):
@property
def model_version(self):
return "v1_7b"
def get_tokenizer(self):
return safe_get_tokenizer_from_pretrained(
self.model_name_or_path,
revision=self.revision,
padding_side="left",
truncation_side="left",
trust_remote_code=True,
use_fast=False
)
def get_input_builder(self):
return PanguInputBuilder(self.tokenizer, self.model_version)
def get_toolscallprocessor(self):
# 没有实现toolcall功能可以不写
return ToolsCallProcessorPangu(self.model_version)
openpangu使用[unused11]和[unused12]来包裹工具调用的回答部分，我们只需使用正则表达式提取该部分文本，然后返回对应的tool_call字段格式。
import string
import random
import json
import re
NAME = "name"
ARGUMENTS = "arguments"
class ToolsCallProcessorPangu:
def __init__(self, model_version):
self.model_version = model_version
self.name_key = "name"
self.arguments_key = "arguments"
def decode(self, content):
lines = content.strip()
pattern = re.compile(r'
$$unused11$$
\s*({.*?})\s*
$$unused12$$
', re.DOTALL)
matches = pattern.findall(lines)
is_tool_call = True
if matches:
tool_calls = []
try:
tool_calls = [json.loads(match) for match in matches]
for item in tool_calls:
_ = item[self.name_key]
_ = item[self.arguments_key]
except json.JSONDecodeError:
is_tool_call = False
if is_tool_call:
call_res = []
for item in tool_calls:
tool_call = {
self.name_key: item[self.name_key],
self.arguments_key: json.dumps(item[self.arguments_key], ensure_ascii=False)
}
characters = string.ascii_letters + string.digits
call_id = "call_" + ''.join(random.choice(characters) for _ in range(8))
res = {
"type": "function",
"id": call_id,
"function": tool_call
}
call_res.append(res)
return {"tool_calls": call_res}
return content.strip()因为openpangu-embedded-7b-model的"model_type"为"PanguEmbedded"，需要在atb_models/atb_llm/models/__init__.py中重新映射以加载我们的pangu模型
def router model type(model_type, config_dict, model_type_key, model_name_or_path):
model_type map = {
"kclgpt": "codeshell",
"internvl_chat": "internvl",
"LLava_next_video": "Llava_next"
"bunny-qwen2":"bunny",
"bunny-minicpm": "bunny",
'deepseek_v2': 'deepseekv2',
'deepseek_v3': 'deepseekv2',
"vita-qwen2": "vita",
"qwen2_5_vl": "qwen2_vl",
"ernie4_5_moe": "ernie_moe",
"sophon": "pangu",
"panguembedded": "pangu",
}
if model_type in model_type_map:
model_type = model_type_map[model_type]需要在atb_models/atb_llm/models/base/router.py中重新映射以加载我们的pangu模型
class BaseRouter:
def __post_init__(self):
setf.model_type = "Janus"
if self.model_type == "qwen2_5_vl":
self.model_type = "qwen2_vl"
if self.model_type "ernie4_5_moe":
self.model_type = "ernie_moe"
if self.model_type == "sophon" or self.model_type == "panguembedded":
self.model_type = "pangu"
self.model_type_cap = self.model_type.capitalize()
if self.model_type_cap == "Qwen2_moe":
self.model_type_cap = self.model_type_cap.replace('_','')服务端接收到请求后会调用tokenize()，当请求为OpenAI格式时，将调用模型侧InputBuilder类的make_context()接口。
{model}InputBuilder类继承自基类InputBuilder，在模型迁移适配过程中，该类需实现以下方法：
apply_chat_template_default_apply_chat_templatefrom ..base.input_builder import InputBuilder
from ...utils.log import logger
from ...utils.log.error_code import ErrorCode
class PanguInputBuilder(InputBuilder):
def __init__(self, tokenizer, model_version, **kwargs):
self.model_version = model_version
super().__init__(tokenizer, **kwargs)
def apply_chat_template_default(self, conversation, **kwargs):
role_field = "role"
content_field = "content"
bos_token = "<s>"
eos_token = "[unused10]"
start_token = "[unused9]"
formatted = bos_token
for idx, message in enumerate(conversation):
if idx == 0 and conversation[0]['role'] != 'system':
formatted += start_token + "系统：" + eos_token
if message[role_field] == "user":
formatted += start_token + "用户：" + message[content_field] + eos_token
elif message[role_field] == "assistant":
formatted += start_token + "助手：" + message[content_field] + eos_token
elif message[role_field] == "system":
formatted += start_token + "系统：" + message[content_field] + eos_token
elif message[role_field] == "tool":
formatted += start_token + "工具：" + message[content_field] + eos_token
else:
msg = "Only user/assistant/system roles are supported!"
logger.error(msg, ErrorCode.ATB_MODELS_PARAM_OUT_OF_RANGE)
raise ValueError(msg)
if "add_generation_prompt" in kwargs and kwargs["add_generation_prompt"]:
formatted += start_token + "助手："
# 已拼接bos_token, encode内add_special_tokens要设置False,如果为True,则会重复有bos_token
return self.tokenizer.encode(formatted, add_special_tokens=False)
def _apply_chat_template(self, conversation, **kwargs):
return self.apply_chat_template_default(conversation, **kwargs)
在router_pangu.py中可以看到，我们将当前模型版本设置为v1_7b，因此有关模型的配置文件、transformer文件、权重加载以及前向推理等配置文件均在v1_7b目录下进行处理。
mkdir v1_7b
cd v1_7b
touch config_pangu.py flash_causal_pangu.py modeling_pangu.pymodeling_pangu.py参考模型仓的实现进行修改。
ATB-models中的模型transformer类仅用于加载权重，不用于前向推理，因而不需要定义forward函数。
迁移过程中应尽量使用atb-models定义好的模型结构基类，如attention，mlp等，以免带来兼容性问题。
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright (c) Huawei Technologies Co., Ltd.2025. All rights reserved.
"""
PanguModel
"""
import torch
import torch.distributed
from torch import nn
from transformers.activations import ACT2FN
from atb_llm.nn import Module
from atb_llm.utils import OpBackend
from atb_llm.utils.layers import (
TensorParallelRowLinear,
TensorParallelColumnLinear,
PositionRotaryEmbedding,
TensorParallelEmbedding,
load_column_multi,
KvCache,
get_linear
)
from atb_llm.utils.quantize.pack_type import PackType, calc_linear_pack_type, ALL_PACK_LIST
def load_row(config, prefix: str, weights, bias: bool):
weight = weights.get_multi_weights_row(prefix, quantize=config.quantize)
if bias:
bias = weights.get_tensor(f"{prefix}.bias")
linear = get_linear(weight, bias, config.quantize)
# different all_reduce positions
return TensorParallelRowLinear(linear, process_group=weights.process_group)
class PanguRMSNorm(nn.Module):
def __init__(self, prefix, weights, eps=1e-6):
super().__init__()
weight = weights.get_tensor(f"{prefix}.weight")
self.weight = nn.Parameter(weight)
self.variance_epsilon = eps
class PanguRMSNormBias(nn.Module):
def __init__(self, prefix, weights, eps=1e-6):
super().__init__()
weight = weights.get_tensor(f"{prefix}.weight")
try:
bias = weights.get_tensor(f"{prefix}.bias")
except AssertionError:
bias = torch.zeros(weight.shape, dtype=weights.dtype)
self.weight = nn.Parameter(weight)
self.bias = nn.Parameter(bias)
self.variance_epsilon = eps
class PanguRMSNormWrapper(nn.Module):
def __init__(self, prefix, weights, eps=1e-6):
super().__init__()
self.ori = PanguRMSNorm(prefix, weights, eps)
self.anti = PanguRMSNormBias(f'{prefix}.module', weights, eps)
class PanguRMSNormAntiOutlierWrapper(nn.Module):
def __init__(self, prefix, weights, eps=1e-6):
super().__init__()
self.ori = PanguRMSNorm(f'{prefix}.ori', weights, eps)
self.anti = PanguRMSNormBias(f'{prefix}.anti', weights, eps)
class PanguMLP(nn.Module):
def __init__(self, prefix, config, weights):
super().__init__()
act = config.hidden_act
approximate = "tanh" if act in ["gelu_fast", "gelu_pytorch_tanh"] else "none"
self.act = (
ACT2FN[act]
if "gelu" not in act
else lambda x: torch.nn.functional.gelu(x, approximate=approximate)
)
linear_names = [f'{prefix}.up_proj', f'{prefix}.gate_proj']
layer_prefix = '.'.join(prefix.split('.')[:-1])
norm_name = f'{layer_prefix}.post_attention_layernorm'
pack_name = f'{prefix}.gate_up_proj'
self.pack_type = calc_linear_pack_type(weights, linear_names, norm_name, pack_name)
if self.pack_type in ALL_PACK_LIST:
self.gate_up_proj = load_column_multi(
config,
prefixes=[f"{prefix}.gate_proj", f"{prefix}.up_proj"],
weights=weights,
head_size=1,
)
else:
self.gate_proj = TensorParallelColumnLinear.load(
config,
prefix=f"{prefix}.gate_proj",
weights=weights,
bias=False,
)
self.up_proj = TensorParallelColumnLinear.load(
config,
prefix=f"{prefix}.up_proj",
weights=weights,
bias=False,
)
self.down_proj = TensorParallelRowLinear.load(
config,
prefix=f"{prefix}.down_proj",
weights=weights,
bias=False,
)
self.intermediate_size = (
(config.intermediate_size + weights.process_group.size() - 1) // weights.process_group.size()
)
class PanguAttention(nn.Module):
def __init__(
self,
prefix: str,
config,
weights,
**kwargs
):
super().__init__()
self.num_heads = config.num_attention_heads
self.hidden_size = config.hidden_size
self.head_size = self.hidden_size // self.num_heads
self.rotary_emb = PositionRotaryEmbedding.static(
dim=self.head_size, base=10000.0, device="cpu").to(weights.device)
self.softmax_scale = self.head_size ** -0.5
# can support self.num_heads % weights.process_group.size() != 0
if (config.num_attention_heads != config.num_key_value_heads
and (self.num_heads % weights.process_group.size() != 0)):
raise ValueError(
f"`num_heads` must be divisible by `num_shards` (got `num_heads`: {self.num_heads} "
f"and `num_shards`: {weights.process_group.size()}"
)
self.num_heads = (self.num_heads + weights.process_group.size() - 1) // weights.process_group.size()
if config.num_key_value_heads != config.num_attention_heads:
self.num_key_value_heads = config.num_key_value_heads
self.num_key_value_heads = self.num_key_value_heads // weights.process_group.size()
else:
self.num_key_value_heads = self.num_heads
linear_names = [f'{prefix}.q_proj', f'{prefix}.k_proj', f'{prefix}.v_proj']
self.kv_cache_quant = None
if config.quantization_config.kv_quant_type is not None:
if len(linear_names) >= 3:
k_name = linear_names[1]
v_name = linear_names[2]
else:
k_name = f"{linear_names[0]}.k_proj"
v_name = f"{linear_names[0]}.v_proj"
self.kv_cache_quant = KvCache.load(prefix_k=k_name,
prefix_v=v_name, weights=weights, gqa_size=self.head_size,
backend=kwargs.get("attn_decode_backend", OpBackend.ATB))
layer_prefix = '.'.join(prefix.split('.')[:-1])
norm_name = f'{layer_prefix}.input_layernorm'
pack_name = f'{prefix}.query_key_value'
self.pack_type = calc_linear_pack_type(weights, linear_names, norm_name, pack_name)
if self.pack_type in ALL_PACK_LIST:
self.query_key_value = TensorParallelColumnLinear.load_multi(
config,
prefixes=[f"{prefix}.q_proj",
f"{prefix}.k_proj",
f"{prefix}.v_proj"],
weights=weights,
bias=config.bias,
dim=0
)
else:
self.q_proj = TensorParallelColumnLinear.load(
config,
prefix=f"{prefix}.q_proj",
weights=weights,
bias=config.bias
)
self.k_proj = TensorParallelColumnLinear.load(
config,
prefix=f"{prefix}.k_proj",
weights=weights,
bias=config.bias
)
self.v_proj = TensorParallelColumnLinear.load(
config,
prefix=f"{prefix}.v_proj",
weights=weights,
bias=config.bias
)
self.o_proj = load_row(config,
prefix=f"{prefix}.o_proj",
weights=weights,
bias=config.bias)
self.num_groups = self.num_heads // self.num_key_value_heads
self.kv_head_mapping = torch.arange(
0, self.num_key_value_heads, dtype=torch.int32, device=weights.device
).repeat_interleave(self.num_groups)
self.prefix = prefix
class PanguDecoderLayer(nn.Module):
def __init__(self, layer_id, config, weights, attn_decode_backend):
super().__init__()
prefix = f"model.layers.{layer_id}"
self.self_attn = PanguAttention(prefix=f"{prefix}.self_attn", config=config, weights=weights,
attn_decode_backend=attn_decode_backend)
self.mlp = PanguMLP(prefix=f"{prefix}.mlp", config=config, weights=weights)
if self.self_attn.pack_type in [PackType.ALL_FP]:
self.input_layernorm = PanguRMSNorm(
prefix=f"{prefix}.input_layernorm", weights=weights,
eps=config.rms_norm_eps)
else:
raise AssertionError(f'self_attn.pack_type: {self.self_attn.pack_type} not supported')
if self.mlp.pack_type in [PackType.ALL_FP]:
self.post_attention_layernorm = PanguRMSNorm(prefix=f"{prefix}.post_attention_layernorm", weights=weights,
eps=config.rms_norm_eps)
else:
raise AssertionError(f'mlp.pack_type: {self.mlp.pack_type} not supported')
class PanguModel(Module):
def __init__(self, config, weights, **kwargs):
super().__init__()
process_group = weights.process_group
self.tp_rank = process_group.rank()
self.tp_world_size = process_group.size()
self.quantize = config.quantize
self.embed_tokens = TensorParallelEmbedding(
prefix="model.embed_tokens", weights=weights
)
layer_list = []
for idx in range(config.num_hidden_layers):
layer_list.append(PanguDecoderLayer(
idx, config, weights, attn_decode_backend=kwargs.get("attn_decode_backend", OpBackend.ATB)
))
self.layers = nn.ModuleList(layer_list)
self.norm = PanguRMSNorm(
prefix="model.norm", weights=weights, eps=config.rms_norm_eps
)
self.gradient_checkpointing = False
self.head_size = self.layers[0].self_attn.head_size
self.num_heads = self.layers[0].self_attn.num_heads
self.num_key_value_heads = self.layers[0].self_attn.num_key_value_heads
{model}Config类实现模型配置参数的加载，用于初始化模型。
openpangu适配如下：
from dataclasses import dataclass
from typing import Optional
from ...base.config import BaseConfig
@dataclass
class PanguConfig(BaseConfig):
bias: bool = True
bos_token_id: int = 1
eos_token_id: int = 2
hidden_act: str = "silu"
hidden_size: int = 4096
initializer_range: float = 0.02
intermediate_size: int = 12800
max_position_embeddings: int = 32768
model_type: str = "PanguEmbedded"
num_attention_heads: int = 32
num_hidden_layers: int = 34
num_key_value_heads: Optional[int] = None
pad_token_id: int = 0
rms_norm_eps: float = 1e-6
tie_word_embeddings: bool = False
use_cache: bool = True
vocab_size: int = 153376
rope_theta: float = 16000000.0
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.set_model_type(**kwargs)
if self.num_key_value_heads is None:
self.num_key_value_heads = self.num_attention_heads
def set_model_type(self, **kwargs):
if 'model_type' not in kwargs:
self.model_type = 'PanguEmbedded'Flash{model}ForCausalLM类实现了模型的初始化、权重加载以及前向推理。
forward在FlashForCausalLM中底座大模型已适配，所有的Flash{model}ForCausalLM继承于基类FlashForCausalLM。
openpangu适配如下：
import json
from typing import Optional, List, Tuple
import torch
import torch_npu
from atb_llm.models.base.flash_causal_lm import FlashForCausalLM
from atb_llm.models.pangu.v1_7b.config_pangu import PanguConfig
from atb_llm.models.pangu.v1_7b.modeling_pangu import PanguModel
from atb_llm.utils.data.weight_wrapper import WeightWrapper, AttnWrapper, MlpWrapper
from atb_llm.utils.layers import load_column_multi
from atb_llm.utils.layers.embedding.position_rotary_embedding import PositionEmbeddingType
from atb_llm.utils.layers.linear.linear_utils import LinearUtils
from atb_llm.utils.layers.norm.fast_layer_norm import NormType
from atb_llm.utils.log.logging import logger
from atb_llm.utils.op_backend import OpBackend
from atb_llm.utils.quantize.quant_type import QuantType
CPP_PANGU_MODEL_CLASS_NAME = "pangu_7b_DecoderModel"
class FlashPanguForCausalLM(FlashForCausalLM):
def __init__(self, config: PanguConfig, weights, **kwargs):
self.acl_decoder_regression_operation = None
super().__init__(config, weights, **kwargs)
self.attn_decode_backend = OpBackend.ATB
self.soc_info.matmul_nd_nz = (self.soc_info.soc_version == 225 or self.soc_info.soc_version == 223) and \
(config.quantize is None or config.quantize in [QuantType.FLOAT])
LinearUtils.soc_info = self.soc_info
self.aclnn_matmul_backend = (self.soc_info.soc_version == 225 or self.soc_info.soc_version == 223) and \
(config.quantize is None or config.quantize in [QuantType.FLOAT])
self.model = PanguModel(config, weights, attn_decode_backend=self.attn_decode_backend)
self.lm_head = load_column_multi(config, prefixes=["lm_head"], weights=weights, head_size=1, lm_head=True)
self.num_hidden_layers = config.num_hidden_layers
self.config = config
self.place_holder = torch.zeros(1, dtype=self.dtype, device="npu")
self.transdata_operation = torch.classes.OperationTorch.OperationTorch("TransdataOperation")
self.transdata_param = json.dumps({})
self.transdata_operation.set_param(self.transdata_param)
self.acl_param = None
self.acl_operation_inputs = None
self.ascend_weight = None
# prefix_cache是否启用自回归
self.prefix_cache_regression = False
self.mapping = weights.mapping
process_group = weights.process_group
self.tp_rank = process_group.rank()
self.tp_world_size = process_group.size()
self.decode_weight = None
# AddRmsNorm特性
self.enable_inter_layer_add_norm = True
self.enable_intra_layer_add_norm = True
def init_ascend_operations(self, config: PanguConfig):
# 初始化模型
logger.info(f"using {CPP_PANGU_MODEL_CLASS_NAME}")
self.acl_encoder_operation = torch.classes.ModelTorch.ModelTorch(CPP_PANGU_MODEL_CLASS_NAME)
self.acl_decoder_operation = torch.classes.ModelTorch.ModelTorch(CPP_PANGU_MODEL_CLASS_NAME)
if self.prefix_cache_enable:
self.acl_decoder_regression_operation = torch.classes.ModelTorch.ModelTorch(
CPP_PANGU_MODEL_CLASS_NAME)
def get_weights(self, quantize_type: QuantType = None):
quantize_type = self.quantize if quantize_type is None else quantize_type
attn_wrapper = AttnWrapper(
norm_name='input_layernorm', wrapper_name='self_attn',
pack_name='query_key_value', sep_names=['q_proj', 'k_proj', 'v_proj'], o_name='o_proj'
)
mlp_wrapper = MlpWrapper(
norm_name='post_attention_layernorm', wrapper_name='mlp',
pack_name='gate_up_proj', sep_names=['gate_proj', 'up_proj'], down_name='down_proj'
)
weight_wrapper = WeightWrapper(self.soc_info, self.tp_rank, attn_wrapper, mlp_wrapper)
weight_wrapper.register_embedding(self.model.embed_tokens)
for _, layer in enumerate(self.model.layers):
weight_wrapper.mlp_wrapper = mlp_wrapper
weight_wrapper.register_layer(layer, quantize_type)
if self.enable_intra_layer_add_norm or self.enable_inter_layer_add_norm:
weight_wrapper.register_layer_addrmsnormquant(layer, attn_wrapper, mlp_wrapper, self.quantize)
if self.config.quantization_config.kv_quant_type is not None:
weight_wrapper.register_layer_kvquant(layer)
if self.soc_info.need_nz:
del layer.self_attn
del layer.post_attention_layernorm
del layer.mlp
weight_wrapper.register_model_norm(self.model.norm)
weight_wrapper.register_model_lmhead(self.lm_head)
return weight_wrapper
def init_ascend_weight(self):
weight_wrapper = self.get_weights()
self.ascend_weight = weight_wrapper.weights
linear_types = weight_wrapper.linear_type
pack_quant_configs = weight_wrapper.pack_quant_type
linear_descs_configs = weight_wrapper.linear_descs
linear_transpose_types = weight_wrapper.linear_transpose_types
decode_linear_descs_configs = linear_descs_configs
acl_param_dict = self.get_acl_param_dict(pack_quant_configs, linear_types, linear_transpose_types)
acl_param_encoder = {**acl_param_dict, "isPrefill": True, "enableLcoc": self.lcoc_enable,
"enableSpeculate": False, "enableSplitFuse": self.split_fuse_enable,
"linearDescs": linear_descs_configs}
acl_param_decoder = {**acl_param_dict, "isPrefill": False,
"enableLcoc": False, "enableSpeculate": self.speculate_enable,
"enablePrefixCache": self.prefix_cache_enable, "linearDescs": decode_linear_descs_configs}
self.acl_encoder_operation.set_param(json.dumps({**acl_param_encoder}))
self.acl_decoder_operation.set_param(json.dumps({**acl_param_decoder}))
self.acl_encoder_operation.set_weight(self.ascend_weight)
self.acl_decoder_operation.set_weight(self.ascend_weight)
if self.prefix_cache_enable:
self.init_prefix_cache_regression_weight(acl_param_dict)
def get_acl_param_dict(self, pack_quant_configs, linear_types, linear_transpose_types):
"""
get acl_param_dict
"""
return {
"normEps": self.config.rms_norm_eps,
"normType": NormType.RMS_NORM,
"isUnpadInputs": True,
"numAttentionHeadsPerRank": self.num_attention_heads,
"hiddenSizePerAttentionHead": self.head_size,
"numHiddenLayers": self.config.num_hidden_layers,
"numKeyValueHeadsPerRank": self.num_key_value_heads,
"isFA": False,
"isBF16": self.dtype == torch.bfloat16 and not self.soc_info.need_nz,
"packQuantType": pack_quant_configs,
"linearQuantType": linear_types,
"linearTransposeType": linear_transpose_types,
"isEmbeddingParallel": True,
"isLmHeadParallel": True,
"lmHeadTransposeType": self.lm_head.linear.trans_flag,
"enableSwiGLU": False if self.soc_info.need_nz else True,
"enableSwigluQuant": self.enable_swiglu_quant,
"enableKvQuant": self.config.quantization_config.kv_quant_type is not None,
"rank": self.tp_rank,
"worldSize": self.tp_world_size,
"backend": self.soc_info.communication_backend,
"attnBackend": self.attn_decode_backend,
"enableInterLayerAddNorm": self.enable_inter_layer_add_norm,
"enableIntraLayerAddNorm": self.enable_intra_layer_add_norm,
"positionEmbeddingType": PositionEmbeddingType.ROPE,
"linearHasBias": [[True, True, False, False]] * self.config.num_hidden_layers,
"matmulBackend": OpBackend.ACLNN if self.aclnn_matmul_backend else OpBackend.ATB,
}
def init_prefix_cache_regression_weight(self, coder_param):
# prefix cache特性多加一张图用于自回归decode
decoder_regression_param = {
**coder_param, "isPrefill": False,
"enableLcoc": False,
"enableSpeculate": False
}
self.acl_decoder_regression_operation.set_param(json.dumps({**decoder_regression_param}))
self.acl_decoder_regression_operation.set_weight(self.ascend_weight)
def prepare_inputs_for_ascend(
self, input_ids: torch.Tensor,
position_ids: torch.Tensor,
is_prefill: bool,
kv_cache: List[Tuple[torch.Tensor, torch.Tensor]],
block_tables: torch.Tensor,
slots: torch.Tensor,
input_lengths: torch.Tensor,
max_seq_len: int,
lm_head_indices: Optional[torch.Tensor] = None,
**kwargs
):
q_lens = None
seq_len = "seqLen"
self.acl_param = json.dumps({seq_len: input_lengths.tolist()})
cos_table, sin_table = self._get_cos_sin_tables()
self.prefix_cache_regression = False
attention_mask = self._get_attention_mask(is_prefill, max_seq_len, **kwargs)
if is_prefill:
if self.split_fuse_enable:
q_lens = kwargs.get('q_lens', [])
self.acl_param = json.dumps({
seq_len: input_lengths.tolist(),
"qLen": q_lens
})
q_lens = torch.tensor(q_lens, dtype=torch.int32, device=self.device)
else:
if self.speculate_enable and not self.prefix_cache_regression:
q_lens = kwargs.get('q_lens', [])
self.acl_param = json.dumps({
seq_len: input_lengths.tolist(),
"qLen": q_lens
})
q_lens = torch.tensor(q_lens, dtype=torch.int32, device=self.device)
lm_head_indices = self._update_lm_head_indices(input_ids, is_prefill, lm_head_indices)
self.acl_operation_inputs = [
input_ids, # IN_TENSOR_INPUTIDS
position_ids, # IN_TENSOR_POSITIONIDS
cos_table, # IN_TENSOR_COSEMBED
sin_table, # IN_TENSOR_SINEMBED
attention_mask, # IN_TENSOR_ATTENTIONMASK
block_tables.to(torch.int32), # IN_TENSOR_BLOCK_TABLES
slots.to(torch.int32), # IN_TENSOR_SLOTS
self.place_holder,
self.place_holder,
self.place_holder,
input_lengths.to(torch.int32), # IN_TENSOR_SEQ_LENGTHS
lm_head_indices, # IN_TENSOR_LOGTIS_INDICES
]
self._update_acl_operation(is_prefill, q_lens)
return self.acl_operation_inputs, self.acl_param
def execute_ascend_operator(self, acl_inputs, acl_param, is_prefill):
if is_prefill:
acl_model_out = self.acl_encoder_operation.execute(acl_inputs, acl_param)
else:
# prefix cache自回归decode
if self.prefix_cache_enable and self.prefix_cache_regression:
acl_model_out = self.acl_decoder_regression_operation.execute(acl_inputs, acl_param)
else:
acl_model_out = self.acl_decoder_operation.execute(acl_inputs, acl_param)
try:
acl_hidden_state = acl_model_out[0]
except IndexError as e:
raise RuntimeError("运行时报错，请开启日志进一步定位问题") from e
return acl_hidden_state
def init_kvcache(self, kv_cache):
kcache_id_diff = self.ascend_kcache_id != id(kv_cache[0][0])
vcache_id_diff = self.ascend_vcache_id != id(kv_cache[0][1])
kcache_shape_diff = self.ascend_kcache_shape != kv_cache[0][0].shape
vcache_shape_diff = self.ascend_vcache_shape != kv_cache[0][1].shape
kcache_diff = not self.ascend_kcache_id or kcache_id_diff or kcache_shape_diff
vcache_diff = not self.ascend_vcache_id or vcache_id_diff or vcache_shape_diff
if kcache_diff or vcache_diff:
k_caches, v_caches = map(lambda x: list(x), zip(*kv_cache))
logger.debug(f"k_caches[0] original shape: {k_caches[0].shape=}")
if self.soc_info.need_nz:
k_caches = [torch_npu.npu_format_cast_(k_cache, 29) for k_cache in k_caches]
v_caches = [torch_npu.npu_format_cast_(v_cache, 29) for v_cache in v_caches]
logger.debug(f"k_caches[0] shape after transdata: {k_caches[0].shape=}")
self.acl_encoder_operation.set_kv_cache(k_caches, v_caches)
self.acl_decoder_operation.set_kv_cache(k_caches, v_caches)
# 加入prefix-cache-kvcache初始化
if self.prefix_cache_enable:
self.acl_decoder_regression_operation.set_kv_cache(k_caches, v_caches)
self.ascend_kcache_id = id(kv_cache[0][0])
self.ascend_vcache_id = id(kv_cache[0][1])
self.ascend_kcache_shape = kv_cache[0][0].shape
self.ascend_vcache_shape = kv_cache[0][1].shape
logger.debug(f"id of kcache is {self.ascend_kcache_id}, id of vcache is {self.ascend_vcache_id}")
def _update_lm_head_indices(self, input_ids: torch.Tensor, is_prefill: bool,
lm_head_indices: Optional[torch.Tensor] = None):
if lm_head_indices is None:
lm_head_indices = torch.tensor(range(input_ids.shape[0]), dtype=torch.int64, device=input_ids.device)
if self.prefix_cache_regression or not (is_prefill or self.prefix_cache_enable):
lm_head_indices = self.lm_head_indices_fake
return lm_head_indices
def _update_acl_operation(self, is_prefill: bool, q_lens):
is_speculate = self.speculate_enable and not is_prefill and not self.prefix_cache_regression
is_split_fuse = self.split_fuse_enable and is_prefill
if is_speculate or is_split_fuse:
self.acl_operation_inputs.append(q_lens)
def _get_cos_sin_tables(self):
self.rotary_embedding.update_cos_sin_cache_total(
self.dtype,
self.device,
self.max_position_embeddings
)
cos_table = self.rotary_embedding.get_cos_cached_total()
sin_table = self.rotary_embedding.get_sin_cached_total()
return cos_table, sin_table
def _get_attention_mask(self, is_prefill: bool, max_seq_len: int, **kwargs):
q_lens = kwargs.get('q_lens', [])
if is_prefill:
attention_mask = self.attn_mask.get_attn_mask(
max_seq_len if self.split_fuse_enable else self.max_base_len, self.dtype, self.device)
# BF16 PA算子需要使用-10000.0表示mask掉的部分
if self.split_fuse_enable and self.dtype == torch.bfloat16:
attention_mask = attention_mask * -10000.0
if self.soc_info.need_nz:
attention_mask = self.transdata_operation.execute([attention_mask])[0]
return attention_mask
else:
if self.prefix_cache_enable and q_lens == []: # 开启prefix cache时q_lens为空时使用自回归
self.prefix_cache_regression = True
attention_mask = self.attn_mask_fake
if self.speculate_enable and not self.prefix_cache_regression:
spec_mask = kwargs.get('spec_mask', None)
if self.soc_info.need_nz:
spec_mask = self.transdata_operation.execute([spec_mask])[0]
attention_mask = spec_mask
return attention_maskcd MindIE-LLM/examples/atb_models/atb_framework/models
mkdir -p pangu/base
cd pangu/base
mkdir layer
mkdir model在pangu/base/model中分别创建decoder_model.h和decoder_model.cpp decoder_model.h
#ifndef ATB_SPEED_MODELS_PANGU_BASE_DECODER_MODEL_H
#define ATB_SPEED_MODELS_PANGU_BASE_DECODER_MODEL_H
#include <vector>
#include "atb_speed/utils/model_factory.h"
#include "models/pangu/base/layer/decoder_layer.h"
#include "models/base/model/decoder_model.h"
namespace atb_speed {
namespace pangu_base {
class PanguModelParam : public atb_speed::base::ModelParam {
public:
void PrintParam() override;
uint32_t numVanillaLayers = 0;
uint32_t numAugLayers = 0;
protected:
void ParseParam(const nlohmann::json ¶mJson) override;
};
class PanguBaseDecoderModel : public atb_speed::base::DecoderModel {
public:
explicit PanguBaseDecoderModel(const std::string ¶m);
private:
atb::Status CreateLayerOperation(atb::Operation **op, uint32_t layerId) override;
PanguModelParam param;
};
} // namespace pangu_base
} // namespace atb_speed
#endifdecoder_model.cpp：
#include "models/pangu/base/layer/decoder_layer.h"
#include "models/pangu/base/model/decoder_model.h"
namespace atb_speed {
namespace pangu_base {
void PanguModelParam::ParseParam(const nlohmann::json ¶mJson)
{
atb_speed::base::ModelParam::ParseParam(paramJson);
if (paramJson.contains("numVanillaLayers")) {
this->numVanillaLayers = paramJson["numVanillaLayers"].get<uint32_t>();
}
if (paramJson.contains("numAugLayers")) {
this->numAugLayers = paramJson["numAugLayers"].get<uint32_t>();
}
}
void PanguModelParam::PrintParam()
{
atb_speed::base::ModelParam::PrintParam();
ATB_SPEED_LOG_DEBUG(
"PanguModelParam "
<< ", enableInterLayerAddNorm: " << this->enableInterLayerAddNorm
<< ", enableIntraLayerAddNorm: " << this->enableIntraLayerAddNorm
<< ", matmulBackend" << this->matmulBackend
<< ", enableSwiGLU: " << this->enableSwiGLU
<< ", enableSwigluQuant" << this->enableSwigluQuant
<< ", enableRopeQuantKvcache" << this->enableRopeQuantKvcache
<< ", enablePrefixCache: " << this->enablePrefixCache
<< ", enableSpeculate: " << this->enableSpeculate
<< ", enableSplitFuse: " << this->enableSplitFuse
<< ", numHiddenLayers: " << this->numHiddenLayers
<< ", numVanillaLayers: " << this->numVanillaLayers
<< ", numAugLayers: " << this->numAugLayers
);
}
PanguBaseDecoderModel::PanguBaseDecoderModel(const std::string ¶m) : DecoderModel(param)
{
ATB_SPEED_LOG_DEBUG("PanguBaseDecoderModel Graph");
this->param.FromString(param);
}
atb::Status PanguBaseDecoderModel::CreateLayerOperation(atb::Operation **op, uint32_t layerId)
{
ATB_SPEED_LOG_DEBUG("CreateLayerOperation, graph PanguBaseDecoderLayer, layerId:" << layerId);
PanguLayerParam layerParam;
this->SetLayerParam(layerParam, layerId);
PanguBaseDecoderLayer decoderLayer(layerParam);
CHECK_OPERATION_STATUS_RETURN(decoderLayer.BuildGraph(op));
return atb::NO_ERROR;
}
} // namespace pangu_base
} // namespace atb_speed在pangu/base/layer中分别创建decoder_layer.h和decoder_layer.cpp decoder_layer.h：
#ifndef ATB_SPEED_MODELS_PANGU_BASE_DECODER_LAYER_H
#define ATB_SPEED_MODELS_PANGU_BASE_DECODER_LAYER_H
#include "models/base/layer/decoder_layer.h"
#include "models/base/param/layer_param.h"
namespace atb_speed {
namespace pangu_base {
class PanguLayerParam : public atb_speed::base::LayerParam {
public:
void PrintParam() override;
};
class PanguBaseDecoderLayer : public atb_speed::base::DecoderLayer<atb::infer::RmsNormParam> {
public:
explicit PanguBaseDecoderLayer(const PanguLayerParam ¶m);
~PanguBaseDecoderLayer() override{};
protected:
PanguLayerParam param;
};
} // namespace pangu_base
} // namespace atb_speed
#endifdecoder_layer.cpp：
#include "models/pangu/base/layer/decoder_layer.h"
namespace atb_speed {
namespace pangu_base {
void PanguLayerParam::PrintParam()
{
LayerParam::PrintParam();
}
PanguBaseDecoderLayer::PanguBaseDecoderLayer(
const PanguLayerParam ¶m) : base::DecoderLayer<atb::infer::RmsNormParam>(param)
{
this->param = param;
this->param.PrintParam();
};
} // namespace pangu_base
} // namespace atb_speedcd MindIE-LLM/examples/atb_models/atb_framework/models
mkdir -p pangu/7b
cd pangu/7b
mkdir layer
mkdir model在pangu/7b/model中分别创建decoder_model.h和decoder_model.cpp decoder_model.h
#ifndef ATB_SPEED_MODELS_PANGU_7B_DECODER_MODEL_H
#define ATB_SPEED_MODELS_PANGU_7B_DECODER_MODEL_H
#include "models/pangu/base/model/decoder_model.h"
namespace atb_speed {
namespace pangu_7b {
class DecoderModel : public atb_speed::pangu_base::PanguBaseDecoderModel {
public:
explicit DecoderModel(const std::string ¶m);
};
REGISTER_MODEL(pangu_7b, DecoderModel);
} // namespace pangu_7b
} // namespace atb_speed
#endifdecoder_model.cpp
#include "models/pangu/7b/model/decoder_model.h"
namespace atb_speed {
namespace pangu_7b {
DecoderModel::DecoderModel(const std::string ¶m) : atb_speed::pangu_base::PanguBaseDecoderModel(param)
{
ATB_SPEED_LOG_DEBUG("DecoderModel Graph");
}
} // namespace pangu_7b
} // namespace atb_speed
在pangu/7b/model中分别创建decoder_layer.h和decoder_layer.cpp decoder_layer.h
#ifndef ATB_SPEED_MODELS_PANGU_7B_DECODER_LAYER_H
#define ATB_SPEED_MODELS_PANGU_7B_DECODER_LAYER_H
#include "models/pangu/base/layer/decoder_layer.h"
namespace atb_speed {
namespace pangu_7b {
class DecoderLayer : public atb_speed::pangu_base::PanguBaseDecoderLayer {
public:
explicit DecoderLayer(const atb_speed::pangu_base::PanguLayerParam ¶m);
~DecoderLayer() override{};
};
} // namespace pangu_7b
} // namespace atb_speed
#endifdecoder_layer.cpp
#include "models/pangu/7b/layer/decoder_layer.h"
namespace atb_speed {
namespace pangu_7b {
DecoderLayer::DecoderLayer(
const atb_speed::pangu_base::PanguLayerParam ¶m) : atb_speed::pangu_base::PanguBaseDecoderLayer(param)
{
ATB_SPEED_LOG_DEBUG("DecoderLayer Graph");
};
} // namespace pangu_7b
} // namespace atb_speedsource /usr/local/Ascend/ascend-toolkit/set_env.sh
source /usr/local/Ascend/nnal/atb/set_env.shcd MindIE-LLM/examples/atb_models
bash scripts/build.sh# 编译成功后设置环境变量
source output/atb_models/set_env.sh && \
source /usr/local/Ascend/mindie/set_env.sh && \
source /usr/local/Ascend/mindie/latest/mindie-service/set_env.sh修改服务化配置文件
cd /usr/local/Ascend/mindie/latest/mindie-service
vim conf/config.json主要需要修改模型名称和模型路径，其余的根据场景需要设定 "modelName" : "openpangu-embedded-7b-model", "modelWeightPath" : "/data/weights/openpangu-embedded-7b-model/",
拉起服务化
./bin/mindieservice_daemon出现Daemon start success!说明服务化拉起成功
The old environment variable ATB_LOG_TO_STDOUT will be deprecated on 2025/12/31.Please use the new environment variable MINDIE_LOG_TO_STDOUT as soon as possible.
Daemon start success!首先使用curl发送单条命令测试服务化是否正常拉起：
curl -X POST -s http://127.0.0.1:1025/v1/completions -H "Content-Type: application/json" -d '{
"model": "openpangu-embedded-7b-model",
"prompt": "San Francisco is a",
"max_tokens": 28,
"temperature": 0,
"stream": true
}'注意这里model名与服务化配置保持一致，可正常得到模型回答：
root@localhost:/usr/local/Ascend/mindie/latest/mindie-service# curl -X POST -s http://127.0.0.1:1025/v1/completions -H "Content-Type: application/json" -d '{
"model": "openpangu-embedded-7b-model",
"prompt": "San Francisco is a",
"max tokens": 28,
"temperature": 0,
"stream": true
}'
data: {"id":"endpoint_common_1","object":"text_completion","created":1762411753,"model":"openpangu-embedded-7b-model","choices":[{"index":O,"text":" beautiful","logprobs":null,"stop_reason":null,"finish_reason":null}]}
data: {"id":"endpoint_common_1","object":"text_completion","created":1762411753,"model":"openpangu-embedded-7b-model","choices":[{"index":O,"text":" city","logprobs":null,"stop_reason":null,"finish_reason":null}]}
data: {"id":"endpoint_common_1","object":"text_completion","created":1762411754,"model":"openpangu-embedded-7b-model","choices":[{"index":O,"text":",","logprobs":null,"stop_reason":null,"finish_reason":null}]}
data: {"id":"endpoint_common_1","object":"text_completion","created":1762411754,"model":"openpangu-embedded-7b-model","choices":[{"index":O,"text":" but","logprobs":null,"stop_reason":null,"finish_reason":null}]}最后使用 ais_bench 进行服务化性能测试
ais_bench -m perf \
--models vllm_api_general_stream --datasets synthetic_gen \
--debug --dump-eval-details --summarizer example --num-prompts 8
注意执行前需先修改vllm_api_general_stream和synthetic_gen文件，文件位置为：/opt/package/benchmark/ais_bench/benchmark/configs/./models/vllm_api/vllm_api_general_stream.py 和 /opt/package/benchmark/ais_bench/datasets/synthetic/synthetic_config.py
其中vllm_api_general_stream.py文件需要对模型名称和路径以及ip和端口号做修改，和服务化配置保持一致，trust_remote_code改为True
from ais_bench.benchmark.models import VLLMCustomAPIStream
models = [
dict(
attr="service",
type=VLLMCustomAPIStream,
abbr='vllm-api-general-stream',
path="/data/weights/openpangu-embedded-7b-model/",
model="openpangu-embedded-7b-model",
request_rate = 0,
retry = 2,
host_ip = "localhost",
host_port = 1025,
max_out_len = 512,
batch_size=1,
trust_remote_code=True,
generation_kwargs = dict(
temperature = 0.5,
top_k = 10,
top_p = 0.95,
seed = None,
repetition_penalty = 1.03,
)
)
]
synthetic_config.py代码适配如下，需将TrustRemoteCode设置为True
synthetic_config = {
"Type":"string", # [tokenid/string]，生成的随机数据集类型，支持固定长度的随机tokenid，和随机长度的string，两种类型的数据集
"RequestCount": 9999, # 生成的请求条数，应与模型侧配置文件中的 decode_batch_size 一致
"TrustRemoteCode": True, #是否信任远端代码，tokenid模式下需要加载tokenizer生成tokenid，默认为Fasle
"StringConfig" : { # string类型的随机数据集的配置相关项，请参考以上注释处："StringConfig中的随机生成方法参数说明"
"Input" : { # 每条请求的输入长度
"Method": "uniform",
"Params": {"MinValue": 128, "MaxValue": 128}
},
"Output" : { # 每条请求的输出长度
"Method": "gaussian",
"Params": {"Mean": 128, "Var": 0, "MinValue": 128, "MaxValue": 128}
}
},
"TokenIdConfig" : { # tokenid类型的随机数据集的配置相关项
"RequestSize": 10 # 每条请求的长度，即每条请求中token id的个数，应与模型侧配置文件中的 input_seq_len 一致
}
}
配置完成后执行测试命令行，测试结果正常
| 通用指标 | 阶段 | 数值 |
|---|---|---|
| 基准测试时长 | 总计 | 50312.7944 毫秒 |
| 总请求数 | 总计 | 8 |
| 失败请求数 | 总计 | 0 |
| 成功请求数 | 总计 | 8 |
| 并发数 | 总计 | 1.0 |
| 最大并发数 | 总计 | 1 |
| 请求吞吐量 | 总计 | 0.159 请求/秒 |
| 总输入 tokens 数 | 总计 | 1032 |
| 预填充 Token 吞吐量 | 总计 | 1635.5591 Token/秒 |
| 总生成 tokens 数 | 总计 | 1024 |
| 输入 Token 吞吐量 | 总计 | 20.5117 Token/秒 |
| 输出 Token 吞吐量 | 总计 | 20.3527 Token/秒 |
| 总 Token 吞吐量 | 总计 | 40.8644 Token/秒 |
从测试数据上看，MindIE TPOT 性能优于 vllm-ascend 33%以上。
| 输入 tokens 数 | 最大输出 tokens 数 | 并发数 | 并行策略 | 浮点类型 | MindIE（TTFT 单位:毫秒） | vLLM（TTFT 单位:毫秒） | MindIE（TPOT 单位:毫秒） | vLLM（TPOT 单位:毫秒） |
|---|---|---|---|---|---|---|---|---|
| 128 | 1024 | 1 | TP2 | FP16 | 93.68 | 107.99 | 47.19 | 63.44 |
| 128 | 1024 | 4 | TP2 | FP16 | 187.72 | 210.29 | 48.62 | 67.30 |
| 128 | 1024 | 16 | TP2 | FP16 | 429.91 | 461.64 | 55.63 | 67.33 |
| 128 | 1024 | 32 | TP2 | FP16 | 736.14 | 715.01 | 65.37 | 78.15 |
上线 openpangu-embedded-7b 模型，满足业务使用。