本文档记录 01ai/Yi-1.5-34B-Chat-16K 在昇腾 NPU 环境的适配与验证结果。
Yi-1.5-34B-Chat-16K 是一个基于 LLaMA 架构的 340 亿参数规模的大语言模型，支持 16K 上下文长度，采用分组查询注意力（Grouped Query Attention, GQA）机制，并使用 bfloat16 精度。
相关获取地址：
| 组件 | 版本 |
|---|---|
| NPU | Ascend910B |
| CANN | 25.5.2 |
| PyTorch | 2.9.0 |
| torch_npu | 2.9.0.post1 |
| transformers | 4.40.0+ |
| vllm | 0.18.0+empty |
| vllm_ascend | 0.18.0rc1 |
2 逻辑卡，每卡 65.8 GB/opt/atomgit/Yi-1.5-34B-Chat-16K/model/01ai/Yi-1.5-34B-Chat-16K| 参数 | 值 |
|---|---|
| 模型类型 | llama |
| 架构 | LlamaForCausalLM |
| 参数量 | 34B |
| 隐藏层大小 | 7168 |
| 注意力头数 | 56 |
| KV 头数 | 8 (GQA) |
| 层数 | 60 |
| 中间层大小 | 20480 |
| 最大位置编码 | 16384 |
| 词表大小 | 64000 |
| 激活函数 | silu |
| 精度 | bfloat16 |
使用 transformers + torch_npu 直接推理：
python3 inference.py使用 vLLM-Ascend 部署：
export ASCEND_RT_VISIBLE_DEVICES=0,1
export VLLM_USE_MODELSCOPE=true
export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
vllm serve /opt/atomgit/Yi-1.5-34B-Chat-16K/model/01ai/Yi-1.5-34B-Chat-16K \
--host 0.0.0.0 \
--port 8000 \
--tensor-parallel-size 2 \
--dtype bfloat16 \
--trust-remote-code \
--gpu-memory-utilization 0.90 \
--max-model-len 16384# 检查 NPU 状态
npu-smi info
# 推理测试
python3 -c "
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch, torch_npu
tokenizer = AutoTokenizer.from_pretrained('/opt/atomgit/Yi-1.5-34B-Chat-16K/model/01ai/Yi-1.5-34B-Chat-16K', trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
'/opt/atomgit/Yi-1.5-34B-Chat-16K/model/01ai/Yi-1.5-34B-Chat-16K',
torch_dtype=torch.bfloat16,
trust_remote_code=True,
device_map='auto'
)
inputs = tokenizer('Hello, how are you?', return_tensors='pt')
if 'npu' in str(model.device):
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model.generate(**inputs, max_length=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
"测试条件：单次推理，bfloat16 精度，max_length=50。
| 指标 | 数值 |
|---|---|
| 推理时间 | 6.48 秒 |
| 硬件 | Ascend910B |
| 精度 | bfloat16 |
| 设备映射 | device_map="auto" |
| 测试输入 | 模型输出 |
|---|---|
| Hello, how are you? | Hello, how are you? |
| What is 1+1? | 2 |
| 中国的首都是哪里？ | 北京 |
| Explain quantum computing. | Quantum computing uses quantum bits... |