昇腾 NPU 部署 Llama 大模型实战指南

2. 核心代码

注意设备转移方式，字典类型不能直接调用 .npu() 方法。

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)

start = time.time()
outputs = model.generate(**inputs, max_new_tokens=50)
end = time.time()

text = tokenizer.decode(outputs[0])
print(f"生成文本：{text}")
print(f"吞吐量：{50/(end-start):.2f} tokens/s")

性能测试

为获得准确数据，需加入预热环节并同步 NPU 状态。

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 = 10

def load_model_and_tokenizer(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"显存占用：{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)
    
    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)
        torch.npu.synchronize()
        end_time = time.time()
        latencies.append(end_time - start_time)
    
    avg_latency = sum(latencies) / len(latencies)
    throughput = max_new_tokens / avg_latency
    return {"平均延迟 (秒)": round(avg_latency, 2), "平均吞吐量 (tokens/秒)": round(throughput, 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}
    ]
    
    for case in test_cases:
        result = benchmark(case["输入"], tokenizer, model, case["生成长度"])
        print(f"{case['场景']}：{result['平均吞吐量 (tokens/秒)']} tokens/秒")

实测结果显示，三种场景下的吞吐量约为 16-17 tokens/s，性能稳定。

常见问题与解决方案

1. torch.npu 找不到：报错 AttributeError: module 'torch' has no attribute 'npu'。需在代码开头添加 import torch_npu。
2. tokenizer.npu() 不存在：字典类型无此方法。应使用 .to('npu:0') 转移张量。
3. Llama 下载需要权限：官方仓库受限。改用开源镜像版本 NousResearch/Llama-2-7b-hf。
4. 网络超时：国内访问 HuggingFace 不稳定。设置 HF_ENDPOINT 或使用 ModelScope 镜像。

性能优化建议

1. MindSpeed-LLM 框架：使用昇腾官方框架进行深度优化。
2. INT8 量化：使用 BitsAndBytesConfig 加载 INT8 精度模型，减少显存占用并提升速度。
3. 批处理推理：增加 batch size，显著提升吞吐量。

总结

昇腾 NPU 适合对供应链自主可控有要求的政企项目，以及预算有限但需跑大模型的团队。云上测试可大幅降低验证成本。部署时需特别注意环境配置细节，如导入 torch_npu 及设备转移方式。虽然当前性能略低于顶级 NVIDIA GPU，但在性价比和自主可控方面具有优势。