昇腾 NPU 部署 Llama 大模型全流程实战与性能测试

六、性能测试

1. 测试脚本优化

包含预热、同步计时及批量统计逻辑：

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

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

def load_model_and_tokenizer(model_name):
    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
    ).to(DEVICE)
    model.eval()
    print(f"模型已加载到{DEVICE}，显存占用：{torch.npu.memory_allocated()/1e9:.2f}GB")
    return model, tokenizer

def benchmark(prompt, tokenizer, model, 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, do_sample=False, pad_token_id=tokenizer.eos_token_id)
    
    latencies = []
    for i in range(TEST_RUNS):
        torch.npu.synchronize()
        start_time = 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.time()
        latency = end_time - start_time
        latencies.append(latency)
        print(f"第{i+1}次测试：耗时{latency:.2f}秒，速度{max_new_tokens/latency:.2f}tokens/秒")
    
    avg_latency = sum(latencies)/len(latencies)
    throughput = max_new_tokens / avg_latency
    return {
        "prompt": prompt,
        "平均延迟 (秒)": round(avg_latency, 2),
        "平均吞吐量 (tokens/秒)": round(throughput, 2),
        "显存占用 (GB)": round(torch.npu.memory_allocated()/1e9, 2)
    }

if __name__ == "__main__":
    model, tokenizer = load_model_and_tokenizer(MODEL_NAME)
    test_cases = [
        {"场景": "英文短文本生成", "输入": "The capital of France is", "生成长度": 100},
        {"场景": "中文对话", "输入": "请解释什么是人工智能：", "生成长度": 100},
        {"场景": "代码生成", "输入": "Write a Python function to calculate fibonacci:", "生成长度": 150}
    ]
    results = []
    for case in test_cases:
        result = benchmark(case["输入"], tokenizer, model, case["生成长度"])
        result["场景"] = case["场景"]
        results.append(result)
    print("===== 测试完成 =====")
    for res in results:
        print(f"{res['场景']}：{res['平均吞吐量 (tokens/秒)']} tokens/秒")

2. 测试结果

• 英文短文本生成：16.75 tokens/秒
• 中文对话：16.58 tokens/秒
• 代码生成：16.59 tokens/秒

七、常见问题排查

1. torch.npu 找不到：必须显式 import torch_npu。
2. tokenizer.npu() 不存在：字典类型无此方法，应使用 .to('npu:0')。
3. 模型下载权限：使用开源镜像版本 NousResearch/Llama-2-7b-hf。
4. 网络超时：设置 HF_ENDPOINT 或使用 ModelScope 镜像。

八、优化建议

1. MindSpeed-LLM 框架：昇腾官方框架，针对 NPU 深度优化。
2. INT8 量化：减少显存占用并提升速度，配置 BitsAndBytesConfig。
3. 批处理推理：增加 batch size 可显著提升吞吐量。

九、总结

昇腾 NPU 适合对供应链自主可控有要求的政企项目及预算有限的团队。云上测试成本低，生态逐步完善。部署时需注意导入 torch_npu 及设备转移细节，合理预期性能表现。