File size: 2,461 Bytes
b673e72 | 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 | 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()
|