昇腾 NPU 部署 Llama 2 模型：性能测试与优化实践

1. 设置环境变量使用国内 HF 镜像源。
2. 必须先 import torch_npu。
3. 输入张量需迁移至 NPU 设备，而非调用不存在的 .npu() 方法。

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

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

MODEL_NAME = "NousResearch/Llama-2-7b-hf"
DEVICE = "npu:0"

print("开始加载模型...")
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()

prompt = "The capital of France is"
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

with torch.no_grad():
    start_time = time.time()
    outputs = model.generate(**inputs, max_new_tokens=50)
    end_time = time.time()

    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"生成内容：{generated_text}")
    print(f"推理耗时：{end_time - start_time:.2f} 秒")

性能测试

为评估真实性能，设计包含预热与多次测试的基准脚本。使用 torch.npu.synchronize() 确保计时准确。

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

MODEL_NAME = "NousResearch/Llama-2-7b-hf"
DEVICE = "npu:0"
WARMUP_RUNS = 3
TEST_RUNS = 5

def load_model():
    print("加载模型与分词器...")
    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()
    return model, tokenizer

def benchmark(prompt, model, tokenizer, 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)
    
    # 正式测试
    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)
        torch.npu.synchronize()
        end = time.time()
        latency = end - start
        latencies.append(latency)
        print(f"第{i + 1}次耗时：{latency:.2f}s")
    
    avg_latency = sum(latencies) / len(latencies)
    throughput = max_new_tokens / avg_latency
    return throughput, avg_latency

if __name__ == "__main__":
    model, tokenizer = load_model()
    test_cases = [
        {"场景": "英文生成", "提示": "The future of artificial intelligence is", "长度": 100},
        {"场景": "中文问答", "提示": "请用简单的话解释量子计算：", "长度": 100},
        {"场景": "代码生成", "提示": "Write a Python function to reverse a string:", "长度": 150},
    ]
    
    print("\n" + "=" * 50)
    print("性能测试结果")
    print("=" * 50)
    for case in test_cases:
        throughput, avg_latency = benchmark(case["提示"], model, tokenizer, case["长度"])
        print(f"- {case['场景']}:")
        print(f"  平均延迟：{avg_latency:.2f}s")
        print(f"  吞吐量：{throughput:.2f} tokens/s")
    print("=" * 50)

测试结果分析

在 Ascend 910B 实例上的实测数据如下：

结论：FP16 精度下吞吐量稳定在 20 tokens/s 左右。对于离线批处理或内部工具开发，该性能表现满足需求，但在实时性要求极高的场景下仍有提升空间。

性能优化

1. INT8 量化

使用昇腾 ModelArts 或本地工具对模型进行 W8A8 量化，可显著降低显存占用并提升速度。修改加载配置如下：

from transformers import BitsAndBytesConfig

quantization_config = BitsAndBytesConfig(load_in_8bit=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=quantization_config,
    device_map="auto"
)

量化后的性能对比：

2. 批处理（Batch Inference）

同时处理多个请求可大幅提升吞吐量：

prompts = ["Prompt 1", "Prompt 2", "Prompt 3", "Prompt 4"]
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(DEVICE)
outputs = model.generate(**inputs, max_new_tokens=100)

总结

本次实践验证了基于昇腾 NPU 部署和运行 Llama 2 大模型的可行性。虽然绝对性能并非顶尖，但其在成本、自主可控和稳定性方面的优势，使其在 AI 算力多元化的今天成为一个不容忽视的选择。建议开发者先利用云资源验证方案，并密切关注 CANN 及 PyTorch 适配版本的匹配文档。