昇腾 NPU 实战指南：部署与推理 CodeLlama

注：device_map="auto" 会自动把模型分配到可用的 NPU 上，同时 FP16 精度可以在保证计算精度的前提下降低显存使用。

我们接下来可以来验证一下模型是否加载成功：

import torch_npu
from transformers import AutoTokenizer, AutoModelForCausalLM

# 模型名称
model_name = "code-llama-7b-instruct"
# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 加载模型到 NPU
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)
# 验证模型是否在 NPU 上
device = next(model.parameters()).device
print(f"模型已加载，当前设备：{device}")
# 测试一次简单推理
prompt = "def fibonacci(n):"
inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(**inputs, max_new_tokens=10)
# 输出生成结果
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("生成结果:", generated_text)

从运行结果我们可以得知当前模型已经加载成功了。

5. 基础推理演示

在完成模型加载后，带大家开始进行 CodeLlama 的基础推理实验，从简单代码生成到多输入批量推理，展示 NPU 的实战效果。

先尝试最基础的单条 prompt 生成：

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "code-llama-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

# 简单 prompt
prompt = "def fibonacci(n):"
# Tokenize 并移动到 NPU
inputs = tokenizer(prompt, return_tensors="pt").to(next(model.parameters()).device)
# 执行推理
outputs = model.generate(**inputs, max_new_tokens=50)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("生成结果:\n", generated_text)

批量推理： 在实际场景中，我们经常需要同时处理多条请求，尝试批量推理：

prompts = [
    "def factorial(n):",
    "def quicksort(arr):",
    "def gcd(a, b):"
]
# 批量 token 化
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(next(model.parameters()).device)
# 批量生成
outputs = model.generate(**inputs, max_new_tokens=50)
# 输出每条结果
for i, output in enumerate(outputs):
    text = tokenizer.decode(output, skip_special_tokens=True)
    print(f"\nPrompt {i+1} 生成结果:\n{text}")

控制生成风格和长度： 还可以通过调整生成参数优化结果，例如：

outputs = model.generate(
    **inputs,
    max_new_tokens=100,  # 最大生成长度
    temperature=0.7,     # 控制生成随机性
    top_p=0.9,           # nucleus sampling
    do_sample=True
)
for i, output in enumerate(outputs):
    text = tokenizer.decode(output, skip_special_tokens=True)
    print(f"\nPrompt {i+1} 生成结果 (控制风格):\n{text}")

多轮推理： 在实际开发中，模型经常用于多轮交互，例如根据不断变化的需求生成或优化代码。接下来写一个代码案例展示如何实现多轮对话式推理：

# 初始对话：生成一个 Python 函数计算平方根
conversation = ["# 请写一个 Python 函数计算平方根"]
for i in range(2):
    # 将对话内容 token 化并移动到 NPU
    inputs = tokenizer(conversation, return_tensors="pt", padding=True).to(next(model.parameters()).device)
    # 执行生成
    outputs = model.generate(**inputs, max_new_tokens=50)
    # 解码生成结果
    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"\n第 {i+1} 轮生成:\n{reply}")
    # 将模型输出追加到对话中，用于下一轮生成
    conversation.append(reply)

6. 性能测试

在完成模型加载和部署后，对 CodeLlama 的推理性能进行了评估。测试中使用了一个包含多个 Python 函数的 prompt 列表，通过批量输入测量模型生成结果的延迟和吞吐量。每次生成的最大 token 数设置为 50，模拟常见的代码生成场景。

示例代码：

import time
from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "code-llama-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)

prompts = [
    "def fibonacci(n):",
    "def factorial(n):",
    "def quicksort(arr):",
    "def gcd(a, b):",
    "def is_prime(n):"
]

# 批量 token 化并移动到 NPU
inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(next(model.parameters()).device)
num_trials = 5
total_time = 0.0

for i in range(num_trials):
    start_time = time.time()
    outputs = model.generate(**inputs, max_new_tokens=50)
    end_time = time.time()
    elapsed = end_time - start_time
    total_time += elapsed
    print(f"Trial {i+1}: {elapsed:.3f} s")

avg_time = total_time / num_trials
num_tokens = sum(len(tokenizer.decode(output, skip_special_tokens=True)) for output in outputs)
throughput = num_tokens / avg_time

print(f"\n平均延迟：{avg_time:.3f} s")
print(f"吞吐量：{throughput:.1f} 字符/s")

在执行过程中，每次生成的延迟都被记录下来，并计算了多次试验的平均延迟。通过解码生成结果，可以统计生成字符数量，从而得到吞吐量。

从实际测试结果来看，CodeLlama 在昇腾 NPU 上的单次生成延迟稳定在几十毫秒级别，平均延迟仅 0.110 秒；同时还能达到 4727.2 字符/秒的较高吞吐量，足以支撑多 prompt 并行处理的场景需求。

7. 总结

对于希望快速体验和操作昇腾 NPU 的开发者来说，云端 Notebook 提供了即开即用的环境，无需依赖本地硬件或复杂配置。在昇腾环境中部署 CodeLlama 并不复杂，同时也有一些需要注意的环节。整体来看，昇腾已经能够比较稳定地承载主流代码生成模型的推理任务，显著降低了上手门槛，让开发者能够更多地专注于模型本身，而非环境搭建。