| |
| """吞吐基准：测不同 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() |
|
|
