File size: 1,727 Bytes
76bbe95 | 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 | """
Day-1 throughput test: forward+backward tokens/sec for a preset on this machine.
Use it to size the token budget before any real run.
python bench.py --config hoard_small --seq_len 1024 --batch 16 --iters 5
"""
import argparse, time
import mlx.core as mx
import mlx.nn as nn
import configs
from model import HOARD
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--config", default="hoard_small")
ap.add_argument("--seq_len", type=int, default=1024)
ap.add_argument("--batch", type=int, default=16)
ap.add_argument("--iters", type=int, default=5)
ap.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"])
a = ap.parse_args()
cfg = configs.get(a.config)
model = HOARD(cfg)
if a.dtype == "bfloat16":
model.set_dtype(mx.bfloat16)
c = model.param_counts()
print(f"{a.config}: total {c['total']/1e6:.1f}M, active/token/loop {c['active_per_token_per_loop']/1e6:.1f}M, loops {cfg.n_loops}")
def loss_fn(model, x):
logits = model(x[:, :-1])
return nn.losses.cross_entropy(logits.astype(mx.float32), x[:, 1:], reduction="mean")
grad_fn = nn.value_and_grad(model, loss_fn)
x = mx.random.randint(0, cfg.vocab_size, (a.batch, a.seq_len + 1))
# warmup
l, g = grad_fn(model, x); mx.eval(l, g)
t0 = time.time()
for _ in range(a.iters):
l, g = grad_fn(model, x); mx.eval(l, g)
dt = (time.time() - t0) / a.iters
toks = a.batch * a.seq_len
print(f"fwd+bwd: {dt*1000:.0f} ms/step, {toks/dt:,.0f} tok/s "
f"-> 1B tokens ~ {1e9/(toks/dt)/3600:.1f} h on this machine")
print(f"peak memory: {mx.get_peak_memory()/2**30:.2f} GB")
if __name__ == "__main__":
main()
|