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

性能基准测试

为了获得更准确的数据，我们编写了自动化测评脚本 Test.py，覆盖多种场景（短文本、长文本、代码、对话等）和批量大小（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"

# 定义测试用例
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},
    {"场景": "高并发批量（batch=4）", "输入": "The capital of France is", "生成长度": 50, "batch_size": 4},
]

def load_model_and_tokenizer(model_name, precision):
    print(f"===== 开始加载模型 {model_name} =====")
    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()
    mem_used = torch.npu.memory_allocated() / 1e9
    print(f"模型加载完成：耗时 {end_load - start_load:.2f} 秒，显存占用 {mem_used:.2f} GB")
    return model, tokenizer, end_load - start_load, 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,
                use_cache=True, cache_implementation="npu_optimized"
            )

    # 正式测试
    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,
                use_cache=True, cache_implementation="npu_optimized"
            )
        torch.npu.synchronize()
        latency = time.time() - 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
    }

if __name__ == "__main__":
    model, tokenizer, load_time, load_mem, actual_dtype = load_model_and_tokenizer(MODEL_NAME, PRECISION)
    results = []
    for case in TEST_CASES:
        print(f"\n===== 开始测试场景：{case['场景']} =====")
        res = benchmark(
            prompt=case["输入"], tokenizer=tokenizer, model=model,
            max_new_tokens=case["生成长度"], batch_size=case["batch_size"]
        )
        res.update({"场景": case["场景"], "输入示例": case["输入"][:50]})
        results.append(res)
        print(f"场景测试完成：{case['场景']} | 批量总吞吐量：{res['批量总吞吐量 (tokens/秒)']:.2f} tokens/秒")

    df = pd.DataFrame(results)
    print("\n===== 测试结果概览 =====")
    print(df[["场景", "batch_size", "单请求吞吐量 (tokens/秒)", "批量总吞吐量 (tokens/秒)"]].to_string())

核心性能数据

经过实测，Llama-2-7B 在昇腾 NPU 上的表现如下：

结论：单请求吞吐量稳定在 15.6-17.6 tokens/秒，不同任务类型适配性优异。批量推理时，batch_size 从 1 增至 4，总吞吐量接近线性增长，说明 NPU 算力未饱和，适合高并发场景。

高并发极限测试

针对生产环境可能面临的大规模并发，我们进一步测试了 batch_size 从 1 到 70 的性能表现。配置 NPU 算子融合与显存池优化后，在 64GB 显存设备上验证了线性增长的上限。

测试发现，batch_size 达到 60 时，总吞吐量约为单请求的 60 倍左右，衰减率控制在 5% 以内。当 batch_size 超过 60 后，由于显存带宽限制，衰减率略有上升，但整体仍保持较高效率。这表明昇腾 NPU 在处理大规模批量推理时具有显著优势。

性能优化方案

为了进一步提升性能，建议在部署前进行以下环境配置与代码调整。

1. NPU 环境配置

通过环境变量开启算子融合与显存池管理：

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

2. 推理逻辑优化

在推理脚本中显式开启 KV 缓存，并固定输出长度以减少动态计算开销：

outputs = model.generate(
    **inputs,
    max_new_tokens=50,
    use_cache=True,
    cache_implementation="npu_optimized",
    num_beams=1,
    pad_token_id=tokenizer.pad_token_id,
    max_length=inputs.input_ids.shape[1] + 50
)

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

常见问题与解决方案

1. 国内镜像源安装依赖失败

执行 pip install 时报 ConnectionTimeout 或 404 错误，可切换镜像源或升级 pip：

python -m pip install --upgrade pip
pip install transformers accelerate -i https://mirrors.aliyun.com/pypi/simple/

2. torch_npu 版本不兼容

导入 torch_npu 时报 AttributeError，需严格匹配版本：

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

显示 True 说明安装正常。

3. 依赖包版本冲突

若出现 ImportError，指定兼容版本安装：

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

总结

本次测试表明，Llama-2-7B 在昇腾 NPU 上展现出极强的性能稳定性与线性增长能力。16GB 显存即可覆盖模型加载到 batch=4 并发的全流程，单请求吞吐量稳定在 15.6-17.6 tokens/秒。通过合理的算子融合与显存优化，该组合既降低了硬件选型成本，又为国产算力下大模型落地提供了稳定可靠的性能保障。对于实时对话场景，推荐 batch_size=1 以保证低延迟；对于离线批量场景，可采用 batch_size=60 以上以最大化吞吐量。