在昇腾 NPU 上运行 Llama 2 模型
面对高昂的 GPU 成本，华为昇腾（Ascend）NPU 凭借其自主可控的达芬奇架构和日益完善的软件生态，成为大模型部署的高性价比选择。本文记录了在昇腾 NPU 环境下从环境配置、模型部署到性能测试与优化的全过程，重点分析实际推理延迟与吞吐量表现。
环境搭建与验证
在云平台或本地部署时，确保计算资源正确识别是第一步。建议使用预装了 CANN 和 PyTorch 适配器的镜像，例如 euler2.9-py38-torch2.1.0-cann8.0-openmind0.6-notebook，可大幅减少手动配置时间。
实例启动后，需确认 NPU 可用。依次执行以下检查命令：
cat /etc/os-release
python3 --version
接着验证 PyTorch 及 torch_npu 插件版本：
import torch
print(f'PyTorch 版本：{torch.__version__}')
import torch_npu
print(f'torch_npu 版本：{torch_npu.__version__}')
注意：直接调用 torch.npu.is_available() 可能会报错，因为 torch_npu 需要显式导入才能注册后端。正确的验证方式是先 import torch_npu 再检查可用性。
模型部署与加载
安装必要的依赖库，国内用户建议使用镜像加速：
pip install transformers accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple
由于 Meta 官方仓库访问受限，推荐使用社区维护的镜像版本如 NousResearch/Llama-2-7b-hf。创建 Python 脚本进行核心部署，关键代码如下：
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
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)
with torch.no_grad():
start_time = time.time()
outputs = model.generate(**inputs, max_new_tokens=50)
end_time = time.time()
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(f"生成内容：{generated_text}")
print(f"推理耗时：{end_time - start_time:.2f} 秒")
关键点说明：
- 环境变量
HF_ENDPOINT用于指定国内镜像源，避免下载超时。 import torch_npu必须在任何 NPU 操作之前执行。- 输入张量迁移使用
.to(DEVICE)，而非不存在的.npu()方法。 - 模型加载时使用
float16精度以节省显存。
性能测试与分析
为了评估真实性能，设计了一个包含预热和多次测试的基准脚本。以下是核心逻辑：
import time
import torch
import torch_npu
from transformers import AutoModelForCausalLM, AutoTokenizer
WARMUP_RUNS = 3
TEST_RUNS = 5
def benchmark(prompt, model, tokenizer, max_new_tokens=100):
inputs = tokenizer(prompt, return_tensors="pt").to("npu:0")
# 预热
for _ in range(WARMUP_RUNS):
with torch.no_grad():
_ = model.generate(**inputs, max_new_tokens=max_new_tokens)
latencies = []
for i in range(TEST_RUNS):
torch.npu.synchronize() # 同步以确保计时准确
start = time.time()
with torch.no_grad():
_ = model.generate(**inputs, max_new_tokens=max_new_tokens)
torch.npu.synchronize()
end = time.time()
latencies.append(end - start)
avg_latency = sum(latencies) / len(latencies)
throughput = max_new_tokens / avg_latency
return throughput, avg_latency
在 NPU Basic 实例上的测试结果如下表所示：
| 测试类型 | 平均延迟 | 吞吐量 |
|---|---|---|
| 英文生成 | 4.94s | 20.24 tokens/s |
| 中文问答 | 4.87s | 20.55 tokens/s |
| 代码生成 | 7.24s | 20.73 tokens/s |
结果分析：
- 性能表现：吞吐量稳定在 20-30 tokens/秒左右。对于离线批处理、内部工具开发及对实时性要求不高的场景完全够用。
- 稳定性：测试过程中未出现崩溃或显著波动，表明硬件调度较为稳健。
- 结论：昇腾 NPU 为运行 Llama 2 提供了可行且具成本优势的算力选项。
性能优化策略
若对默认性能不满意，可通过以下手段进一步提升：
1. INT8 量化
使用昇腾 Modelslim 工具将模型从 FP16 量化至 INT8（W8A8）。这能显著降低内存占用并提升推理速度。实测数据显示，量化后吞吐量可提升约 1.2 倍。
from transformers import BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
quantization_config=quantization_config,
device_map="auto"
)
2. 批处理（Batch Inference）
同时处理多个请求可大幅提升吞吐量。只需将多个 prompt 传入 tokenizer 即可：
prompts = ["Prompt 1", "Prompt 2", "Prompt 3"]
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to("npu:0")
outputs = model.generate(**inputs, max_new_tokens=100)
总结
本次实战证明，基于昇腾 NPU 部署和运行 Llama 2 大模型是一条完全可行的技术路径。虽然绝对性能与顶级消费级 GPU 仍有差距，但其在成本、自主可控和稳定性方面的优势，使其在 AI 算力多元化的今天成为一个不容忽视的选择。建议开发者先从云资源验证方案，并密切关注官方文档的版本匹配问题。
