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
|
"""Throughput vs probe_chunk for ZBP steps (informs the H200 config)."""
import os, sys, time, argparse
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
import torch, torch.nn as nn
from zbp_scaling.model import ScalingLM
from zbp_scaling.zbp import ZBPConfig, zbp_blocks
from zbp_scaling.zbp.autograd import seed_probes
p = argparse.ArgumentParser()
p.add_argument("--device", default="cuda:3")
p.add_argument("--d", type=int, default=512); p.add_argument("--layers", type=int, default=8)
p.add_argument("--vocab", type=int, default=8192); p.add_argument("--seq", type=int, default=1024)
p.add_argument("--bs", type=int, default=8)
a = p.parse_args()
dev = torch.device(a.device); seed_probes(0, dev)
x = torch.randint(0, a.vocab, (a.bs, a.seq), device=dev); y = torch.randint(0, a.vocab, (a.bs, a.seq), device=dev)
def bench(mode, n, chunk, iters=4):
cfg = ZBPConfig(mode=mode, n_probes=n, eps=0.1, probe_chunk=chunk)
m = ScalingLM(a.vocab, a.d, a.layers, 8, a.seq, cfg=cfg).to(dev)
opt = torch.optim.AdamW(m.parameters(), lr=1e-4)
for _ in range(2):
loss = nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()
torch.cuda.synchronize(dev); t = time.time()
for _ in range(iters):
loss = nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()); opt.zero_grad(); loss.backward(); opt.step()
torch.cuda.synchronize(dev)
dt = (time.time() - t) / iters
print(f"{mode:3s} n={n:3d} chunk={chunk:3d}: {dt*1000:7.0f} ms/step {a.bs*a.seq/dt/1000:7.1f} ktok/s peakmem {torch.cuda.max_memory_allocated(dev)/2**30:.1f} GB")
torch.cuda.reset_peak_memory_stats(dev)
return dt
t_bp = bench("bp", 0, 0)
for n in (16, 64):
for chunk in (8, 16, 32, 64):
if chunk > 2 * n: continue
dt = bench("cd", n, chunk)
print(f" -> multiplier vs BP: {dt/t_bp:.1f}x")
|
