Llama-2-7b 昇腾 NPU 性能测评与部署优化指南

编写 llama.py 文件，配置模型名称并加载到 NPU。注意处理 Pad Token 以避免推理报错。

import torch
import torch_npu
from transformers import AutoModelForCausalLM, AutoTokenizer
import time

print("开始测试...")
MODEL_NAME = "NousResearch/Llama-2-7b-hf"
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
)

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.npu() 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"\n生成文本：{text}")
print(f"耗时：{(end-start)*1000:.2f}ms")
print(f"吞吐量：{50/(end-start):.2f} tokens/s")

运行前建议将 Hugging Face 下载源切换至国内镜像站，避免网络超时。

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

执行脚本后，若看到生成文本及速度统计，即表示部署成功。

性能基准测试

为了获得更准确的数据，我们需要编写专门的测评脚本 Test.py，涵盖预热、多批次统计及详细报告生成。

1. 测评脚本核心逻辑

脚本包含全局配置区，用户只需修改模型名称、精度等参数。核心函数 benchmark 负责带预热的性能测试，记录延迟、吞吐量及显存峰值。

import torch
import torch_npu
import time
import json
import pandas as pd
from datetime import datetime
from transformers import AutoModelForCausalLM, AutoTokenizer

# ===================== 全局配置区 =====================
MODEL_NAME = "NousResearch/Llama-2-7b-hf"
DEVICE = "npu:0"
WARMUP_RUNS = 5
TEST_RUNS = 10
SAVE_RESULT = True
PRECISION = "fp16"

TEST_CASES = [
    {"场景": "英文短文本生成", "输入": "The capital of France is", "生成长度": 50, "batch_size": 1},
    {"场景": "中文对话", "输入": "请解释什么是人工智能：", "生成长度": 100, "batch_size": 1},
    {"场景": "代码生成", "输入": "Write a Python function to calculate fibonacci:", "生成长度": 150, "batch_size": 1},
    {"场景": "批量推理（batch=2）", "输入": "The capital of France is", "生成长度": 50, "batch_size": 2},
    {"场景": "长文本叙事", "输入": "请写一篇关于人工智能未来的科幻短篇：", "生成长度": 200, "batch_size": 1},
    {"场景": "多轮问答", "输入": "Q: 什么是机器学习？\nA: ...", "生成长度": 100, "batch_size": 1},
    {"场景": "高并发批量（batch=4）", "输入": "The capital of France is", "生成长度": 50, "batch_size": 4},
]

# ===================== 核心测试函数 =====================
def benchmark(prompt, tokenizer, model, max_new_tokens, batch_size):
    # 构造批量输入
    batch_inputs = [prompt] * batch_size
    inputs = tokenizer(
        batch_inputs,
        return_tensors="pt",
        padding=True if batch_size > 1 else False,
        truncation=True,
        max_length=512
    ).to(DEVICE)

    # 预热
    print(f"预热中...（{WARMUP_RUNS}次，batch_size={batch_size}）")
    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 = []
    print(f"开始正式测试...（{TEST_RUNS}次）")
    for i in range(TEST_RUNS):
        torch.npu.synchronize()
        start = 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()
        latencies.append(end - start)

    avg_latency = sum(latencies) / len(latencies)
    throughput = max_new_tokens / avg_latency
    total_throughput = throughput * batch_size
    mem_peak = torch.npu.max_memory_allocated() / 1e9

    return {
        "平均延迟 (秒)": round(avg_latency, 3),
        "单请求吞吐量 (tokens/秒)": round(throughput, 2),
        "批量总吞吐量 (tokens/秒)": round(total_throughput, 2),
        "显存峰值 (GB)": round(mem_peak, 2),
        "生成长度": max_new_tokens,
        "batch_size": batch_size
    }

if __name__ == "__main__":
    # 加载模型
    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()

    results = []
    for case in TEST_CASES:
        res = benchmark(case["输入"], tokenizer, model, case["生成长度"], case["batch_size"])
        res.update({"场景": case["场景"]})
        results.append(res)
        print(f"场景 {case['场景']} 完成：总吞吐 {res['批量总吞吐量 (tokens/秒)']:.2f} tokens/s")

    # 保存结果
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    with open(f"llama_npu_benchmark_{timestamp}.json", "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)

运行 python test.py 即可获取详细数据。

2. 测试结果分析

基础能力测评

所有基础依赖版本固定，确保测试基准一致。实测显示模型加载耗时约 30-35 秒，加载后显存占用 13.61 GB，稳定性良好。

单请求多场景性能

覆盖短/长文本、中英文、代码、多轮对话，单请求吞吐量稳定在 15.6~17.6 tokens/秒。

批量并发性能

模拟多用户请求，batch_size 从 1 增至 4 时，总吞吐量增长近 4 倍，显存可控。

结论：NPU 对批量推理优化充分，适合高并发场景。

高并发极限测试

针对 64GB 显存卡进行 Batch 1~70 的极限测试。结果显示，Batch=70 时总吞吐量达 1125.87 tokens/秒，是单请求的 68.73 倍，增长衰减率仅 1.81%，验证了 64GB 卡支撑 60-70 Batch 的高稳定性。

性能优化方案

1. 环境配置优化

通过环境变量开启算子融合与显存池，减少碎片和编译开销。

export NPU_FUSION_ENABLE=1
export ASCEND_GLOBAL_MEM_POOL_SIZE=2147483648
export NPU_ENABLE_CACHE_OP=1

2. 推理逻辑优化

在 llama.py 中补充 pad_token，并显式开启 KV 缓存 (use_cache=True) 及定制缓存实现 (cache_implementation="npu_optimized")。

# 补充 pad_token
tokenizer.pad_token = tokenizer.eos_token

# 推理调用
outputs = model.generate(
    **inputs,
    max_new_tokens=50,
    use_cache=True,
    cache_implementation="npu_optimized",
    pad_token_id=tokenizer.pad_token_id,
    eos_token_id=tokenizer.eos_token_id
)

集成优化后，各类生成场景吞吐量均有明显增长，显存占用降低，延迟稳定性大幅改善。

常见问题与解决方案

1. 镜像源安装失败

现象：pip install 报 ConnectionTimeout。

解决：切换阿里云或华为源兜底。

pip install transformers accelerate -i https://mirrors.aliyun.com/pypi/simple --trusted-host mirrors.aliyun.com

2. torch_npu 版本不兼容

现象：导入时报 AttributeError。

解决：严格匹配版本。

pip install torch==2.1.0 torch_npu==2.1.0.post3

3. 依赖包冲突

现象：ImportError 无法导入 AutoModelForCausalLM。

解决：指定兼容版本并清理缓存。

pip uninstall transformers -y
pip cache purge
pip install transformers==4.39.2 accelerate==0.28.0

总结

本次测试验证了 Llama-2-7b 在昇腾 NPU 上的高效表现。16GB 显存即可覆盖模型加载到 Batch=4 并发的全流程，单请求吞吐量 15.6-17.6 tokens/秒，批量总吞吐量达 63.33 tokens/秒。高并发测试进一步证明，在 64GB 显存下，Batch 增至 70 时仍保持线性增长，衰减率极低。结合算子融合与显存池优化，该方案既降低了硬件选型成本，又为国产算力下大模型落地提供了稳定可靠的性能保障。