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

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

随着大模型国产化需求的提升，如何在国产算力上高效部署开源模型成为关键。本文以 Llama-2-7B 为例，基于昇腾 NPU 环境完成从依赖安装到模型部署的全流程落地，并通过多维度测试验证其性能表现，提供可复现的部署方案、性能基准数据及硬件选型建议。

环境准备与配置

在开始之前，确保开发环境已适配昇腾 NPU。推荐使用云端 Notebook 或本地服务器，需确认操作系统、Python 版本及 PyTorch 与 torch_npu 的兼容性。

1. 检查核心环境

进入终端后，先验证基础组件版本：

# 检查系统版本
cat /etc/os-release
# 检查 Python 版本
python3 --version
# 检查 PyTorch 版本
python -c "import torch; print(f'PyTorch 版本：{torch.__version__}')"
# 检查 torch_npu
python -c "import torch_npu; print(f'torch_npu 版本：{torch_npu.__version__}')"

2. 安装依赖库

通过国内镜像源加速安装 transformers 和 accelerate，避免下载超时：

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

若遇到网络问题，可尝试阿里云或华为云镜像源兜底：

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

3. 设置 Hugging Face 镜像

为加快模型权重下载，临时切换 HF_ENDPOINT：

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

模型加载与推理测试

编写推理脚本 llama.py，重点在于正确加载模型至 NPU 设备并处理 Tokenizer 的 padding 问题（Llama 默认无 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)
# 补充 pad_token，避免推理报错
tokenizer.pad_token = tokenizer.eos_token

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", padding=True)
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")

运行脚本即可看到模型加载成功及首字生成延迟。此时显存占用通常在 13.6GB 左右。

性能基准测试

为了更准确地评估性能，我们编写了自动化测评脚本 Test.py，覆盖单请求多场景及批量并发测试。

1. 基础场景测试

脚本预设了英文短文本、中文对话、代码生成、长文本叙事等场景，并统计吞吐量、延迟及显存峰值。

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=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, truncation=True, max_length=512).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():
            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),
        "batch_size": batch_size
    }

if __name__ == "__main__":
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    tokenizer.pad_token = tokenizer.eos_token
    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/秒)']} tokens/s")

2. 测试结果分析

根据实测数据，Llama-2-7B 在昇腾 NPU 上的表现如下：

测试场景	batch_size	单请求吞吐量 (tokens/s)	批量总吞吐量 (tokens/s)	显存峰值 (GB)
英文短文本生成	1	15.60 ~ 17.40	16.08	13.71
中文对话	1	16.01 ~ 17.61	16.01	13.71
代码生成	1	15.69 ~ 17.17	15.69	13.71
高并发批量	4	15.60 ~ 17.40	63.33	16.04

核心结论：

1. 稳定性：单请求吞吐量稳定在 15.6~17.6 tokens/s，延迟标准差≤0.22 秒，无明显抖动。
2. 批量效率：batch_size 从 1 增至 4 时，总吞吐量接近线性增长（约 3.9 倍），适合高并发 API 服务。
3. 显存需求：FP16 精度下，模型加载显存 13.61GB，batch=4 峰值 16.04GB，16GB 显存卡即可支撑全流程。

3. 高并发极限测试

针对 64GB 显存的高配 NPU，进一步测试 batch_size 上限。通过开启算子融合与 KV 缓存，验证了 batch=70 时的性能表现。

# 高并发优化配置示例
os.environ["NPU_FUSION_ENABLE"] = "1"
os.environ["ASCEND_GLOBAL_MEM_POOL_SIZE"] = "8589934592"
os.environ["NPU_ENABLE_CACHE_OP"] = "1"

# 推理时开启 KV 缓存
outputs = model.generate(
    **inputs,
    use_cache=True,
    cache_implementation="npu_optimized",
    ... 
)

在 64GB 显存环境下，batch=70 时总吞吐量可达 1125.87 tokens/s，是单请求的 68.73 倍，且显存峰值仅 17.68GB，资源利用率极高。

性能优化方案

若在实际部署中遇到性能瓶颈，可参考以下优化策略：

1. 环境变量调优

在启动脚本前设置关键环境变量，启用算子融合与显存池管理：

export NPU_FUSION_ENABLE=1
export ASCEND_GLOBAL_MEM_POOL_SIZE=2147483648
export NPU_PRINT_TENSOR_SIZE=0

2. 推理逻辑调整

修改 generate 调用参数，显式开启 KV 缓存并固定输出长度，减少动态计算开销：

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,
    return_dict_in_generate=False
)

3. 量化支持

如需进一步降低显存，可尝试 INT8 量化，但需确保模型权重支持且验证精度损失。

常见问题排查

1. 镜像源连接失败

执行 pip install 时报 ConnectionTimeout，可切换多源镜像兜底：

pip install transformers accelerate -i https://repo.huaweicloud.com/repository/pypi/simple --trusted-host repo.huaweicloud.com

2. torch_npu 版本不兼容

导入报错 AttributeError: 'module' object has no attribute 'npu'，通常因版本不匹配导致。严格安装对应组合：

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/s，批量总吞吐量达 63.33 tokens/s。结合算子融合与 KV 缓存优化，高并发场景下性能衰减率可控，适合生产环境部署。对于需要更高并发的场景，64GB 显存卡可支撑 batch=70 以上的线性扩展，性价比突出。

官方资源参考：

• 昇腾官网：https://www.hiascend.com/
• 昇腾官方文档：https://www.hiascend.com/document
• 昇腾开源仓库：https://gitcode.com/ascend