当前位置：首页 > news >正文

iree 上运行qwen2.5 05b运行效率

news 2025/9/27 7:59:05

#!/usr/bin/env python3
"""
IREE + Qwen2.5-0.5B-Instruct 性能测试脚本
修复参数数量不匹配问题，并添加保存iree-run-module所需npy文件功能
"""
import time
import iree.runtime as rt
from transformers import AutoTokenizer
import numpy as np# ---------- 配置项 ----------
MODEL_PATH = "./model-dir"  # 分词器路径
IREE_MODEL_FILE = "qwen2_5_05b.vmfb"  # IREE编译的模型文件
PROMPT = "请介绍一下人工智能的发展历程。"  # 测试用提示
MAX_NEW_TOKENS = 100  # 最大生成token数
WARMUP_ITERATIONS = 2  # 预热次数
# 根据模型结构确定的Qwen模型参数
NUM_LAYERS = 24        # 层数
HIDDEN_SIZE = 1024     # 隐藏层大小
NUM_HEADS = 2          # 注意力头数
HEAD_DIM = 64          # 头维度
# 计算出past_key_value的总参数数量（除了input_ids、attention_mask、position_ids外的参数数量）
PAST_KEY_VALUE_PARAM_COUNT = NUM_LAYERS * 2 * 4  # 每个past_key_value有4个维度（batch_size, num_heads, seq_len, head_dim）
TOTAL_EXPECTED_ARGS = 3 + PAST_KEY_VALUE_PARAM_COUNT  # 加上前面3个常规参数
# ---------------------------def main():# ---------- 1. 加载分词器与处理输入 ----------print("加载分词器...")tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=True)if tokenizer.pad_token is None:tokenizer.pad_token = tokenizer.eos_tokeninputs = tokenizer(PROMPT, return_tensors="np")input_ids = inputs.input_ids.astype(np.int64)print(input_ids.tolist())seq_length = input_ids.shape[1]print(f"提示词长度: {seq_length} tokens")# 保存input_ids.npy文件np.save("input_ids.npy", input_ids)# 生成并保存attention_mask.npy文件attention_mask = np.ones((1, seq_length), dtype=np.int64)np.save("attention_mask.npy", attention_mask)print(attention_mask.tolist())# 生成并保存position_ids.npy文件position_ids = np.arange(seq_length, dtype=np.int64)[None, :]np.save("position_ids.npy", position_ids)print(position_ids.tolist())# 生成并保存past_key和past_value相关的.npy文件（初始全0）for i in range(NUM_LAYERS * 2):past_key_name = f"past_key_{i}.npy"past_value_name = f"past_value_{i}.npy"past_key = np.zeros((1, NUM_HEADS, 0, HEAD_DIM), dtype=np.float32)past_value = np.zeros((1, NUM_HEADS, 0, HEAD_DIM), dtype=np.float32)print(past_value.tolist())np.save(past_key_name, past_key)np.save(past_value_name, past_value)# ---------- 2. 初始化IREE运行时 ----------print("初始化IREE运行时...")config = rt.Config("local-sync")ctx = rt.SystemContext(config=config)with open(IREE_MODEL_FILE, "rb") as f:vmfb_data = f.read()ctx.add_vm_module(rt.VmModule.from_flatbuffer(config.vm_instance, vmfb_data, warn_if_copy=False))invoke = ctx.modules.module["main_graph"]# ---------- 3. 构造初始输入参数 ----------def create_args(input_ids):"""创建模型所需的输入参数列表，确保参数数量正确"""seq_len = input_ids.shape[1]args = [input_ids,                                   # 0: input_idsnp.ones((1, seq_len), dtype=np.int64),       # 1: attention_masknp.arange(seq_len, dtype=np.int64)[None, :], # 2: position_ids]# 添加past_key和past_value缓存（根据模型实际预期形状，明确使用HEAD_DIM）for _ in range(NUM_LAYERS * 2):# 缓存形状: [batch_size, num_heads, seq_len, head_dim]args.append(np.zeros((1, NUM_HEADS, 0, HEAD_DIM), dtype=np.float32))return args# 创建初始参数initial_args = create_args(input_ids)# ---------- 4. 预热模型 ----------print(f"预热模型 {WARMUP_ITERATIONS} 次...")for _ in range(WARMUP_ITERATIONS):warmup_args = create_args(input_ids)invoke(*warmup_args)print("预热完成")# ---------- 5. 测试首token延迟 ----------print("\n测试首token延迟...")first_args = create_args(input_ids)start_time = time.perf_counter()outputs = invoke(*first_args)first_logits = np.asarray(outputs[0])first_token_id = int(np.argmax(first_logits[0, -1, :]))end_time = time.perf_counter()first_token_time = (end_time - start_time) * 1000print(f"首token延迟: {first_token_time:.2f} ms")print(f"首token: {tokenizer.decode([first_token_id])}")# 保存首token生成后的输出结果（示例，可根据需要进一步细化保存内容，比如单独保存logits等）np.save("first_token_output.npy", outputs)# ---------- 6. 测试连续生成速度 ----------print(f"\n测试生成最多{MAX_NEW_TOKENS}个token的速度...")current_ids = input_ids.copy()generated_tokens = [first_token_id]current_outputs = outputstotal_start_time = time.perf_counter()for i in range(1, MAX_NEW_TOKENS):seq_len = current_ids.shape[1] + 1current_ids = np.append(current_ids, [[generated_tokens[-1]]], axis=1)new_args = create_args(current_ids)[:TOTAL_EXPECTED_ARGS]  # 确保参数数量正确，截取期望数量的参数# 更新动态变化的参数（如input_ids、attention_mask、position_ids等）new_args[0] = current_idsnew_args[1] = np.ones((1, seq_len), dtype=np.int64)new_args[2] = np.arange(seq_len, dtype=np.int64)[None, :]current_outputs = invoke(*new_args)logits = np.asarray(current_outputs[0])next_token_id = int(np.argmax(logits[0, -1, :]))if next_token_id == tokenizer.eos_token_id:print(f"遇到结束标记，提前终止生成")breakgenerated_tokens.append(next_token_id)# 保存每次生成token后的结果（示例，可按需调整）np.save(f"output_{i}.npy", current_outputs)if i % 10 == 0:print(f"已生成 {i} 个token...")total_end_time = time.perf_counter()total_time = total_end_time - total_start_timetotal_tokens = len(generated_tokens) - 1# 计算性能指标print(f"\n性能统计:")print(f"首token延迟: {first_token_time:.2f} ms")print(f"生成总token数（含首token）: {len(generated_tokens)}")print(f"生成耗时（不含首token）: {total_time:.4f} s")if total_time > 0 and total_tokens > 0:tokens_per_second = total_tokens / total_timeprint(f"平均生成速度（不含首token）: {tokens_per_second:.2f} tokens/s")# ---------- 7. 输出生成结果 ----------print(f"\n生成内容:")full_ids = np.append(input_ids, generated_tokens).astype(np.int64)print(tokenizer.decode(full_ids, skip_special_tokens=True))if __name__ == "__main__":main()

python run_qwen2_5_05b.py

加载分词器...

input_ids: [[14880, 109432, 104455, 103949, 103168, 1773]]

提示词长度: 6 tokens

初始化IREE运行时...

预热模型 2 次...

预热完成

测试首token延迟...

首token延迟: 472.10 ms

首token:

测试生成最多100个token的速度...

已生成 10 个token...