背景与选型
在昇腾 NPU 上部署 Llama-2-7B 大模型,主要考量自主可控的达芬奇架构及开源生态。相比昂贵的 GPU 服务器,利用云上的昇腾 Notebook 实例进行验证是性价比极高的方案。
环境准备
创建 Notebook 实例时,计算类型必须选择 NPU。规格建议 1*NPU 800T A2、32v CPU、64GB 内存。镜像需选用预装环境的版本,如 euler2.9-py38-torch2.1.0-cann8.0-openmind0.6-notebook,这能省去手动安装 PyTorch 和 CANN 的麻烦。存储 50G 足以容纳约 13GB 的模型文件。
环境验证
进入终端后,首先检查系统版本和 Python 版本。验证 NPU 可用性的关键步骤是必须先导入 torch_npu 插件。直接运行 torch.npu.is_available() 会报错,正确做法是先执行 import torch_npu。
python -c "import torch; import torch_npu; print(torch.npu.is_available())"
看到输出为 True 时,确认 NPU 已就绪。
依赖安装与模型下载
安装 transformers 和 accelerate 库时,建议使用国内镜像源加速。
pip install transformers accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple
模型下载方面,官方仓库可能需要权限且国内访问不稳定,推荐使用开源社区维护的镜像版本(如 NousResearch/Llama-2-7b-hf),并设置环境变量 HF_ENDPOINT 指向镜像地址。
核心代码实现
加载模型时需注意设备转移。字典类型的输入不能直接调用 .npu() 方法,应使用 .to('npu:0')。以下是完整的推理脚本结构:
import torch
import torch_npu
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
MODEL_NAME = "NousResearch/Llama-2-7b-hf"
DEVICE = "npu:0"
print("开始测试...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float16,
low_cpu_mem_usage=True
).to(DEVICE)
model.eval()
prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
start = time.time()
outputs = model.generate(**inputs, max_new_tokens=50)
end = time.time()
text = tokenizer.decode(outputs[0])
print()
()
()


