在昇腾 NPU 上跑 Llama 大模型:从踩坑到通关的全程实战记
一、选择昇腾 NPU 的原因
- 自主可控:昇腾使用华为自研的达芬奇架构,供应链安全更有保障。
- 生态完善:昇腾 GitCode 仓库拥有多个开源项目,支持 PyTorch、TensorFlow 及 MindSpeed-LLM 框架。
- 资源获取:GitCode 提供免费的昇腾 Notebook 实例,适合低成本测试验证。
二、环境准备
1. 创建 Notebook 实例
进入 GitCode 控制台创建实例,关键配置如下:
- 计算类型:必选 NPU。
- 规格:推荐 1*NPU 800T A2、32v CPU、64GB 内存。
- 镜像:选择
euler2.9-py38-torch2.1.0-cann8.0-openmind0.6-notebook。
- 存储:50G 免费存储足够(模型约 13GB)。
2. 环境配置
预装工具包括 PyTorch 2.1.0、CANN 8.0、Python 3.8 及 torch_npu 2.1.0。
三、验证环境
1. 检查版本
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__}')"
2. 验证 NPU 可用性
注意必须先导入 torch_npu 插件:
import torch
import torch_npu
print(torch.npu.is_available())
四、安装依赖
手动安装 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. 核心代码
import torch
import torch_npu
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
MODEL_NAME = "NousResearch/Llama-2-7b-hf"
print("开始测试...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME, torch_dtype=torch.float16, low_cpu_mem_usage=True
)
print("加载到 NPU...")
model = model.npu()
model.eval()
print(f"显存占用:{torch.npu.memory_allocated()/1e9:.2f} GB")
prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt")
inputs = {k: v.to('npu:0') for k, v in inputs.items()}
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")
六、性能测试
1. 测试脚本优化
包含预热、同步计时及批量统计逻辑:
import torch
import torch_npu
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
DEVICE = "npu:0"
WARMUP_RUNS = 3
TEST_RUNS = 10
MODEL_NAME = "NousResearch/Llama-2-7b-hf"
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)
()
avg_latency = (latencies)/(latencies)
throughput = max_new_tokens / avg_latency
{
: prompt,
: (avg_latency, ),
: (throughput, ),
: (torch.npu.memory_allocated()/, )
}
__name__ == :
model, tokenizer = load_model_and_tokenizer(MODEL_NAME)
test_cases = [
{: , : , : },
{: , : , : },
{: , : , : }
]
results = []
test_cases:
result = benchmark([], tokenizer, model, [])
result[] = []
results.append(result)
()
res results:
()
2. 测试结果
- 英文短文本生成:16.75 tokens/秒
- 中文对话:16.58 tokens/秒
- 代码生成:16.59 tokens/秒
七、常见问题排查
- torch.npu 找不到:必须显式
import torch_npu。
- tokenizer.npu() 不存在:字典类型无此方法,应使用
.to('npu:0')。
- 模型下载权限:使用开源镜像版本
NousResearch/Llama-2-7b-hf。
- 网络超时:设置
HF_ENDPOINT 或使用 ModelScope 镜像。
八、优化建议
- MindSpeed-LLM 框架:昇腾官方框架,针对 NPU 深度优化。
- INT8 量化:减少显存占用并提升速度,配置
BitsAndBytesConfig。
- 批处理推理:增加 batch size 可显著提升吞吐量。
九、总结
昇腾 NPU 适合对供应链自主可控有要求的政企项目及预算有限的团队。云上测试成本低,生态逐步完善。部署时需注意导入 torch_npu 及设备转移细节,合理预期性能表现。