Upload benchmark_throughput.py with huggingface_hub

benchmark_throughput.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+import torch.nn.functional as F
+
+from audio_dit import AudioDiT
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--batch-sizes", type=int, nargs="+", default=[64, 128, 256])
+    ap.add_argument("--x-res", type=int, default=384)
+    ap.add_argument("--y-res", type=int, default=256)
+    ap.add_argument("--warmup", type=int, default=10)
+    ap.add_argument("--iters", type=int, default=40)
+    args = ap.parse_args()
+    dev = "cuda"
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    model = AudioDiT(x_res=args.x_res, y_res=args.y_res).to(dev)
+    opt = torch.optim.AdamW(model.parameters(), lr=2e-4)
+    print(f"[bench] {model.num_params()/1e6:.2f}M params, grid {model.grid_h}x{model.grid_w} "
+          f"({model.grid_h*model.grid_w} tokens) at x_res={args.x_res} y_res={args.y_res}", flush=True)
+
+    for B in args.batch_sizes:
+        try:
+            torch.cuda.reset_peak_memory_stats()
+            x1 = torch.randn(B, 1, args.y_res, args.x_res, device=dev)
+            seq = torch.randn(B, 32, 768, device=dev)
+            pool = torch.randn(B, 512, device=dev)
+
+            def step():
+                x0 = torch.randn_like(x1)
+                t = torch.sigmoid(torch.randn(B, device=dev))
+                tb = t.view(-1, 1, 1, 1)
+                xt = (1 - tb) * x0 + tb * x1
+                target = x1 - x0
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    v = model(xt, t, seq, pool)
+                    loss = F.mse_loss(v.float(), target)
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                opt.step()
+
+            for _ in range(args.warmup):
+                step()
+            torch.cuda.synchronize()
+            t0 = time.time()
+            for _ in range(args.iters):
+                step()
+            torch.cuda.synchronize()
+            el = time.time() - t0
+            sps = args.iters / el
+            peak = torch.cuda.max_memory_allocated() / 2**30
+            print(f"[bench] batch={B:<5} {sps:6.2f} steps/s  {sps*B:8.1f} samples/s  "
+                  f"peak {peak:5.1f} GiB  90k-step ETA {90000/sps/3600:5.2f} h", flush=True)
+        except torch.cuda.OutOfMemoryError:
+            print(f"[bench] batch={B:<5} OOM", flush=True)
+            torch.cuda.empty_cache()
+
+if __name__ == "__main__":
+    main()