File size: 3,888 Bytes
71dabc7 | 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 | #!/usr/bin/env python3
"""吞吐基准：测不同 batch × 序列长度 × 精度下的前向吞吐。
输出 samples/s、tokens/s、延迟 p50 与显存峰值，用于选训练 batch 和部署配置。
输入为随机 token（测的是纯算力上限）；真实短文本的收益还要叠加动态 padding，
见 README「吞吐优化」一节。
用法:
python scripts/benchmark.py # seq=128/512 × batch=8..128, bf16
python scripts/benchmark.py --fp32 # fp32 对照
python scripts/benchmark.py --compile # 额外测 torch.compile 版本
"""
import argparse
import time
import torch
from transformers import AutoModelForSequenceClassification
def bench_once(model, x, attn, iters):
for _ in range(3):
with torch.inference_mode():
model(input_ids=x, attention_mask=attn)
torch.cuda.synchronize()
lat = []
t0 = time.perf_counter()
for _ in range(iters):
t1 = time.perf_counter()
with torch.inference_mode():
model(input_ids=x, attention_mask=attn)
torch.cuda.synchronize()
lat.append((time.perf_counter() - t1) * 1000)
total = time.perf_counter() - t0
return total, lat
def run(model, args, dtype_name, results):
for seq in args.seqs:
for batch in args.batches:
x = torch.randint(1000, 200000, (batch, seq), device="cuda")
attn = torch.ones_like(x)
torch.cuda.reset_peak_memory_stats()
try:
total, lat = bench_once(model, x, attn, args.iters)
except torch.cuda.OutOfMemoryError:
torch.cuda.empty_cache()
print(f"[{dtype_name}] seq={seq:>4} batch={batch:>4} | OOM，跳过")
continue
row = {
"dtype": dtype_name, "seq": seq, "batch": batch,
"samples_per_s": batch * args.iters / total,
"tokens_per_s": batch * args.iters * seq / total,
"p50_ms": sorted(lat)[len(lat) // 2],
"vram_gb": torch.cuda.max_memory_allocated() / 2**30,
}
results.append(row)
print(f"[{row['dtype']}] seq={seq:>4} batch={batch:>4} | "
f"{row['samples_per_s']:>9.1f} 样本/s | "
f"{row['tokens_per_s'] / 1e6:>6.2f}M tok/s | "
f"p50={row['p50_ms']:>7.1f}ms | VRAM={row['vram_gb']:.2f}GB")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", default="models/mmbert-ad")
ap.add_argument("--seqs", type=int, nargs="+", default=[128, 512])
ap.add_argument("--batches", type=int, nargs="+", default=[8, 32, 64, 128])
ap.add_argument("--iters", type=int, default=20)
ap.add_argument("--fp32", action="store_true", help="只测 fp32")
ap.add_argument("--compile", action="store_true", help="追加测 torch.compile")
args = ap.parse_args()
results = []
dtypes = ["fp32"] if args.fp32 else ["bf16"] + (["fp32"] if False else [])
for dt in dtypes:
model = AutoModelForSequenceClassification.from_pretrained(
args.model, attn_implementation="sdpa",
dtype=torch.float32 if dt == "fp32" else torch.bfloat16).eval().cuda()
run(model, args, dt, results)
if args.compile and dt == "bf16":
compiled = torch.compile(model)
print("（torch.compile：每个新 shape 都要重编译，首次会慢）")
run(compiled, args, dt + "-compile", results)
del model
torch.cuda.empty_cache()
if results:
best = max(results, key=lambda r: r["samples_per_s"])
print(f"\n最佳吞吐配置: {best['dtype']} seq={best['seq']} batch={best['batch']} "
f"-> {best['samples_per_s']:.0f} 样本/s ({best['tokens_per_s'] / 1e6:.2f}M tok/s)")
if __name__ == "__main__":
main()
|