Llama-2-7B 在昇腾 NPU 上的性能测评与部署指南

2. 性能基准测试

为了更准确地评估性能，我们编写了包含预热、同步及多批次统计的测评脚本。核心逻辑包括：

1. 预热：消除算子编译开销，确保正式测试数据稳定。
2. 同步：使用 torch.npu.synchronize() 避免计时漂移。
3. KV 缓存：开启 use_cache=True 优化长序列推理。
4. 批量处理：动态调整 batch_size 观察吞吐量变化。

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
PRECISION = "fp16"

def load_model_and_tokenizer(model_name, precision):
    print(f"===== 开始加载模型 {model_name}（精度：{precision}） =====")
    start_load = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    dtype = torch.float16 if precision == "fp16" else torch.int8
    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_name, torch_dtype=dtype, low_cpu_mem_usage=True
        ).to(DEVICE)
    except Exception as e:
        print(f"INT8 精度加载失败，自动 fallback 到 FP16：{str(e)[:50]}")
        dtype = torch.float16
        model = AutoModelForCausalLM.from_pretrained(
            model_name, torch_dtype=dtype, low_cpu_mem_usage=True
        ).to(DEVICE)
    model.eval()
    end_load = time.time()
    load_time = end_load - start_load
    mem_used = torch.npu.memory_allocated() / 1e9
    print(f"模型加载完成：耗时 {load_time:.2f} 秒，显存占用 {mem_used:.2f} GB")
    return model, tokenizer, load_time, mem_used, str(dtype)

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,
                eos_token_id=tokenizer.eos_token_id, use_cache=True,
                cache_implementation="npu_optimized"
            )

    # 正式测试
    latencies = []
    print(f"开始正式测试...（{TEST_RUNS}次，生成长度={max_new_tokens}）")
    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,
                eos_token_id=tokenizer.eos_token_id, use_cache=True,
                cache_implementation="npu_optimized",
                max_length=inputs.input_ids.shape[1] + max_new_tokens,
                return_dict_in_generate=False
            )
        torch.npu.synchronize()
        end = time.time()
        latency = end - start
        latencies.append(latency)
        print(f" 第{i+1}次：耗时 {latency:.2f} 秒 | 速度 {max_new_tokens/latency:.2f} tokens/秒")

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

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

核心性能数据

1. 单请求多场景性能

覆盖短文本、中文对话、代码生成、长文本叙事及多轮问答等典型任务。实测显示，单请求吞吐量稳定在 15.6 ~ 17.6 tokens/秒 区间，延迟波动极小。

2. 批量并发性能

模拟高并发场景，batch_size 从 1 增至 4 时，总吞吐量增长接近线性，显存可控。

3. 高并发极限测试

在 64GB 显存环境下，进一步测试 batch_size 至 70 的情况。结果显示，单请求吞吐量仍维持在 16 tokens/秒左右，总吞吐量随 batch_size 线性增长，70 batch 时达 1125.87 tokens/秒。显存峰值仅从 13.71GB 增至 17.68GB，资源占用极低，证明昇腾 NPU 对大 Batch 推理的支持无明显上限。

性能优化方案

1. 环境变量配置

通过调整 NPU 相关环境变量，可显著提升推理效率。

# 启用算子融合（核心提速）
export NPU_FUSION_ENABLE=1
# 预分配显存池，减少碎片
export ASCEND_GLOBAL_MEM_POOL_SIZE=2147483648
# 优化内存分配策略
export ASCEND_CACHE_CLEAR_INTERVAL=100
# 关闭冗余日志输出
export NPU_PRINT_TENSOR_SIZE=0

2. 代码级优化

• Pad Token 处理：Llama 默认无 pad_token，需在加载前手动指定为 eos_token。
• KV 缓存：显式开启 use_cache=True 和 cache_implementation="npu_optimized"。
• NPU 同步：在计时前后加入 torch.npu.synchronize() 确保时间准确。

常见问题排查

1. 镜像源安装失败

执行 pip install 时报 ConnectionTimeout 或 404 错误，建议切换镜像源或升级 pip。

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

2. torch_npu 版本不兼容

导入时报 AttributeError: 'module' object has no attribute 'npu'，需严格匹配版本。

pip install torch==2.1.0 torch_npu==2.1.0.post3
python -c "import torch_npu; print(torch_npu.npu.is_available())"

总结

本次测试表明，Llama-2-7B 在昇腾 NPU 上展现出极强的性能稳定性与线性增长能力。16GB 显存即可覆盖模型加载到 batch=4 并发的全流程，单请求吞吐量稳定在 15.6-17.6 tokens/秒，批量总吞吐量可达 63.33 tokens/秒。同时兼具低延迟波动、全场景适配及可复现部署方案，为国产算力下大模型落地提供了稳定可靠的性能保障。