在昇腾 NPU 上跑 Llama 大模型:从踩坑到通关的全程实战记
本文分享了在昇腾 NPU 上部署测试 Llama-2-7B 大模型的全过程。提供踩坑经验。
一、为什么选择昇腾?
国外 GPU 价格高昂,Atlas 服务器成本更是让个人开发者望而却步。昇腾 NPU 凭借自主可控的达芬奇架构、完善的开源生态以及 GitCode 提供的免费测试资源,成为了高性价比的选择。
- 自主可控:华为自研达芬奇架构,供应链安全更有保障。
- 生态完善:GitCode 仓库中有大量开源项目,支持 PyTorch、TensorFlow,还有专门的 MindSpeed-LLM 框架。
- 成本优势:通过 GitCode 申请免费的昇腾 Notebook 实例,零成本验证流程,比购买硬件划算得多。
二、环境准备:GitCode 配置要点
创建实例看似简单,实则暗藏陷阱。建议按以下配置操作:
1. 计算类型
务必选择 NPU。误选 CPU 会导致性能极差,甚至无法运行。
2. 规格选择
推荐配置:1*NPU 800T A2、32v CPU、64GB 内存。此配置足以支撑 Llama-2-7B 运行,过高浪费,过低则不够用。
3. 镜像选择
必须使用预装环境的镜像:euler2.9-py38-torch2.1.0-cann8.0-openmind0.6-notebook。该镜像已集成 PyTorch 2.1.0、CANN 8.0 等关键工具。
4. 存储设置
默认 50G 免费存储足够(模型约 13GB),无需额外付费扩容。
三、环境验证:第一个坑来得猝不及防
进入 Notebook 后,首先验证 NPU 可用性。打开 Terminal 输入以下命令检查版本:
cat /etc/os-release
python3 --version
python -c "import torch; print(f'PyTorch 版本:{torch.__version__}')"
python -c "import torch_npu; print(f'torch_npu 版本:{torch_npu.__version__}')"
接着验证 NPU 是否可用。直接运行 torch.npu.is_available() 会报错 AttributeError: module 'torch' has no attribute 'npu'。必须先导入 torch_npu 插件。
正确写法:
import torch
import torch_npu
print(torch.npu.is_available()) # 输出:True
确认输出为 True 且显示有 1 个 NPU 设备后,才算环境就绪。
四、安装依赖
虽然镜像预装了 PyTorch,但运行 Llama 大模型仍需手动安装 transformers 库。建议使用清华源加速:
pip install transformers accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple
五、部署 Llama:下载与运行
1. 模型下载
官方 Llama-2 仓库需要权限且国内访问不稳定。推荐使用开源镜像版本 NousResearch/Llama-2-7b-hf,无需申请权限,下载更稳定。
设置环境变量加速下载:
export HF_ENDPOINT=https://hf-mirror.com
2. 核心代码
注意两个易错点:一是必须导入 torch_npu,二是张量转移需使用 .to('npu:0') 而非 .npu()。
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()
print(f"显存占用:{torch.npu.memory_allocated()/1e9:.2f} GB")
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(f"生成文本:{text}")
print(f"耗时:{(end-start)*1000:.2f}ms")
print(f"吞吐量:{50/(end-start):.2f} tokens/s")
执行 python llama.py 即可下载并运行模型。文件约 13GB,分片加载,显存占用约 13.61GB。
六、性能测试
为了获得准确数据,需加入预热和同步机制。以下是优化后的测试脚本:
import torch
import torch_npu
import time
import json
from datetime import datetime
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_NAME = "NousResearch/Llama-2-7b-hf"
DEVICE = "npu:0"
WARMUP_RUNS = 3
TEST_RUNS = 10
def load_model_and_tokenizer(model_name):
print(f"正在加载模型 {model_name}...")
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()
print(f"模型已加载到{DEVICE},显存占用:{torch.npu.memory_allocated()/1e9:.2f}GB")
return model, tokenizer
def benchmark(prompt, tokenizer, model, max_new_tokens=100):
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
for _ in range(WARMUP_RUNS):
with torch.no_grad():
_ = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.eos_token_id)
latencies = []
for i in range(TEST_RUNS):
torch.npu.synchronize()
start_time = time.time()
with torch.no_grad():
outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.eos_token_id)
torch.npu.synchronize()
end_time = time.time()
latency = end_time - start_time
latencies.append(latency)
print(f"第{i+1}次测试:耗时{latency:.2f}秒,速度{max_new_tokens/latency:.2f}tokens/秒")
avg_latency = sum(latencies) / len(latencies)
throughput = max_new_tokens / avg_latency
return {
"prompt": prompt,
"max_new_tokens": max_new_tokens,
"平均延迟 (秒)": round(avg_latency, 2),
"平均吞吐量 (tokens/秒)": round(throughput, 2),
"显存占用 (GB)": round(torch.npu.memory_allocated()/1e9, 2)
}
if __name__ == "__main__":
model, tokenizer = load_model_and_tokenizer(MODEL_NAME)
test_cases = [
{"场景": "英文短文本生成", "输入": "The capital of France is", "生成长度": 100},
{"场景": "中文对话", "输入": "请解释什么是人工智能:", "生成长度": 100},
{"场景": "代码生成", "输入": "Write a Python function to calculate fibonacci:", "生成长度": 150}
]
results = []
for case in test_cases:
print(f"\n===== 测试场景:{case['场景']} =====")
result = benchmark(prompt=case["输入"], tokenizer=tokenizer, model=model, max_new_tokens=case["生成长度"])
result["场景"] = case["场景"]
results.append(result)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
filename = f"llama_benchmark_{timestamp}.json"
with open(filename, "w", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
print(f"测试结果已保存到 {filename}")
print("\n===== 测试完成 =====")
print("性能总结:")
for res in results:
print(f"{res['场景']}:{res['平均吞吐量 (tokens/秒)']} tokens/秒")
实测结果约为 16-17 tokens/s,虽低于预期但性能稳定。
七、常见问题汇总
-
torch.npu 找不到
- 现象:
AttributeError: module 'torch' has no attribute 'npu' - 解决:代码开头显式添加
import torch_npu。
- 现象:
-
tokenizer.npu() 不存在
- 现象:调用
.npu()报错 - 原因:Tokenizer 返回的是字典,不支持该方法。
- 解决:使用
.to('npu:0')进行设备映射。
- 现象:调用
-
Llama 下载需要权限
- 现象:访问官方仓库被拒绝
- 解决:改用开源镜像
NousResearch/Llama-2-7b-hf。
-
网络超时
- 现象:下载中断
- 解决:设置
HF_ENDPOINT环境变量或使用 ModelScope 镜像。
八、性能优化建议
- MindSpeed-LLM 框架:昇腾官方框架,针对 NPU 深度优化,速度更快。
- INT8 量化:使用
BitsAndBytesConfig进行 INT8 量化,可提升 20%-30% 速度并降低显存占用。 - 批处理推理:增加 batch size,如一次处理 4 个请求,显著提升吞吐量。
九、总结与建议
昇腾 NPU 适合对供应链自主可控有要求的政企项目,或预算有限想跑大模型的团队。云上测试几十块即可验证方案,性价比不错。
注意事项:
- 先云上测试,别盲目买硬件。
- 镜像别乱选,认准 NPU+PyTorch 预装版。
- 细节别忽略,
import torch_npu和to('npu:0')是关键。 - 合理预期性能,满足日常需求没问题。
多逛社区,遇到问题先搜索,往往能找到前人踩坑的经验。

