File size: 3,620 Bytes
1836a26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import time
import torch
import gc
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed
# 固定随机种子，保证采样逻辑尽量可比
set_seed(3407)
# ==================== 对决模型路径配置 ====================
MODEL_A_PATH = "./Qwen3-4B-Thinking-2507" # 1. 原始 36 层完整模型
MODEL_B_PATH = (
"./QiMing-Polaris-Qwen3-4B-Final-Accelerated" # 2. 终极 32 层硬剪枝加速版模型
)
test_prompt = "What is the 'Burden-based Training' method?"
# 统一生成参数
gen_kwargs = {
"max_new_tokens": 256,
"temperature": 0.7,
"top_p": 0.8,
"do_sample": True,
}
def test_model(model_path, model_name):
print(f"\n==========================================================")
print(f"🚀 正在测试模型: [{model_name}]")
print(f"📁 模型路径 : {model_path}")
print(f"==========================================================")
# 1. 加载
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
model_path, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True
)
# 2. 格式化 Prompt
messages = [{"role": "user", "content": test_prompt}]
formatted_input = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(formatted_input, return_tensors="pt").to("cuda")
# 3. GPU 预热
print("🔥 正在预热 GPU...")
with torch.no_grad():
_ = model.generate(**inputs, max_new_tokens=10)
# 4. 正式计时生成
print("⚡ 开始生成...")
torch.cuda.synchronize()
start_time = time.time()
with torch.no_grad():
outputs = model.generate(
**inputs, pad_token_id=tokenizer.eos_token_id, **gen_kwargs
)
torch.cuda.synchronize()
latency = time.time() - start_time
# 5. 解码与数据统计
generated_tokens = outputs[0][inputs.input_ids.shape[1] :]
response_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
num_tokens = len(generated_tokens)
tps = num_tokens / latency if latency > 0 else 0
print(f"\n💬 [{model_name}] 回答:\n{response_text}")
print("-" * 58)
print(
f"⏱️ 耗时: {latency:.4f} 秒 | 生成 Token 数: {num_tokens} | 吞吐率: {tps:.2f} Tokens/s"
)
# 6. 清理 GPU 显存
del model, tokenizer
gc.collect()
torch.cuda.empty_cache()
return {
"name": model_name,
"latency": latency,
"num_tokens": num_tokens,
"tps": tps,
}
print("==========================================================")
print("⚔️ 启动终极对决: 原始 36 层模型 VS 终极 32 层加速模型")
print("==========================================================")
# 依次测试两个成品模型
res_A = test_model(MODEL_A_PATH, "原始 36 层模型 (Base)")
res_B = test_model(MODEL_B_PATH, "终极 32 层加速模型 (Final-Accelerated)")
# 终极对决结算表
print("\n" + "=" * 60)
print("🏆 最终成品模型对决结算表 (Final Models Duel)")
print("=" * 60)
print(
f"1. {res_A['name']:<32} : 耗时 {res_A['latency']:.4f}s | 生成 {res_A['num_tokens']} Tokens | 吞吐率 {res_A['tps']:.2f} Tokens/s"
)
print(
f"2. {res_B['name']:<32} : 耗时 {res_B['latency']:.4f}s | 生成 {res_B['num_tokens']} Tokens | 吞吐率 {res_B['tps']:.2f} Tokens/s"
)
tps_diff = ((res_B["tps"] - res_A["tps"]) / res_A["tps"]) * 100
print("-" * 60)
print(f"🚀 每秒生成吞吐率 (TPS) 提升幅度 : {tps_diff:+.2f}%")
print("=" * 60)
|