| 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() |
|
|
