昇腾 NPU 部署 Llama-2-7B 的踩坑与经验

验证环境

在 Notebook 里打开终端，先检查版本，都没毛病：

cat /etc/os-release            # openEuler 22.03
python3 --version              # Python 3.8.19
python -c "import torch; print(torch.__version__)"           # 2.1.0
python -c "import torch_npu; print(torch_npu.__version__)"   # 2.1.0.post3

然后想验证 NPU 是否可用，直接 torch.npu.is_available() 报错了——原来必须显式 import torch_npu。正确的检查姿势：

import torch
import torch_npu
print(torch.npu.is_available())  # True

看到 True 才放心。另外如果环境里没有 torch，可以这样装：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

不过镜像里已经预置了，通常不需要。

安装依赖

镜像里没有 transformers，需要手动安装：

pip install transformers accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple

加清华源，秒装完。

部署 Llama-2-7B

模型下载：官方仓库 meta-llama/Llama-2-7b-hf 需要申请权限，国内还容易超时。换成开源镜像 NousResearch/Llama-2-7b-hf，不需要额外操作。再设置下 HuggingFace 的国内镜像加速：

export HF_ENDPOINT=https://hf-mirror.com

核心加载代码：

import torch
import torch_npu
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "NousResearch/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
).npu()
model.eval()

注意两点：必须 import torch_npu；移动模型用 .npu()，移动输入（tokenizer 返回的字典）要用 .to('npu:0')，不能 .npu()。

测试推理：

prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to('npu:0')
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0]))

模型下载约 13GB，5 分钟左右完成。加载后显存占用 13.61GB，和 FP16 理论值接近。

性能测试

第一次写的脚本把计时放在 model.generate() 外面，忘记预热，速度虚高。后来用 torch.npu.synchronize() 严格同步，并加了预热，结果靠谱多了。

完整的测试脚本（已解决 NameError 等常见问题）：

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"模型已加载，显存占用：{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()
        with torch.no_grad():
            _ = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=tokenizer.eos_token_id)
        torch.npu.synchronize()
        latencies.append(time.time() - start)
        print(f"第{i+1}次: {latencies[-1]:.2f}s")
    avg_latency = sum(latencies)/len(latencies)
    return avg_latency, max_new_tokens / avg_latency

model, tokenizer = load_model_and_tokenizer(MODEL_NAME)
avg_latency, throughput = benchmark("The capital of France is", tokenizer, model, 100)
print(f"平均延迟: {avg_latency:.2f}s, 吞吐量: {throughput:.2f} tokens/s")

运行后，三个场景（英文短文本、中文对话、代码生成）的吞吐量都在 16～17 tokens/s，虽然不算快，但胜在稳定。

常见问题

坑1：torch.npu 找不到 原因：没有导入 torch_npu。加上 import torch_npu 即可。

坑2：dict 对象没有 .npu() 方法 tokenizer 返回的是字典，要用 .to('npu:0')。

坑3：模型下载需要权限 换用 NousResearch/Llama-2-7b-hf 镜像。

坑4：下载超时 设置镜像 export HF_ENDPOINT=https://hf-mirror.com，或用 ModelScope 下载。

优化方向

• 用 mindie 框架，昇腾官方优化，比原生 transformers 快。
• INT8 量化，显存减半，速度提升 20-30%，只需加个 BitsAndBytesConfig。
• 批处理，我现在只测了 batch=1，调大 batch 吞吐量能翻倍。

总结

这次折腾下来，我觉得昇腾 NPU 比预想的靠谱。开源生态够用，云上免费实例让验证零成本。如果你是做离线批量推理，或者对自主可控有要求，完全值得一试。性能虽然没到顶尖，但搭配优化手段，很多场景已经够用。别光看参数，动手试一把就知道了。