File size: 1,769 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 | import time
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# 1. 路径定义
MODEL_36_LAYERS = "./Qwen3-4B-Thinking-2507" # 原始 36 层 Base
MODEL_32_LAYERS = "./QiMing-Polaris-Qwen3-4B-Final-Accelerated" # 终极 32 层硬剪枝模型
prompt = "What is the 'Burden-based Training' method?"
def test_speed(model_path, name):
print(f"\n正在加载模型 [{name}]...")
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
)
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# 预热
with torch.no_grad():
_ = model.generate(**inputs, max_new_tokens=10)
# 强制生成 200 个 Token (不提前结束，精确测量 TPS)
torch.cuda.synchronize()
start_time = time.time()
with torch.no_grad():
outputs = model.generate(
**inputs,
min_new_tokens=200,
max_new_tokens=200,
do_sample=False, # 贪婪采样，排除随机性
)
torch.cuda.synchronize()
end_time = time.time()
latency = end_time - start_time
tps = 200 / latency
print(
f"⚡ [{name}] 生成 200 Tokens 耗时: {latency:.4f} 秒 | 吞吐率: {tps:.2f} Tokens/s"
)
# 释放显存
del model
torch.cuda.empty_cache()
return tps
# 运行对比
tps_36 = test_speed(MODEL_36_LAYERS, "原始 36 层模型")
tps_32 = test_speed(MODEL_32_LAYERS, "剪枝 32 层终极模型")
speedup = ((tps_32 - tps_36) / tps_36) * 100
print("\n" + "=" * 50)
print(f"🏆 物理剪枝硬加速结果: 每秒生成速度提升了 +{speedup:.2f}% !")
print("=" * 50)
|