| """Fair per-step cost comparison of the five optimizers. |
| |
| The wall-clock totals from the training runs are not comparable: they were |
| produced on a shared GPU, so a run that happened to start while the device was |
| busy looks slower than one that had it to itself, regardless of the update |
| rule. This benchmark instead times every optimizer back-to-back on the same |
| model, same batch shape and same device, reporting the cost of the optimizer |
| *step* separately from the forward/backward pass that all of them share. |
| |
| python -m src.part2.benchmark |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import time |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from ..part1.model import Transformer |
| from ..utils import RUNS, save_json |
| from .optimizers import CATEGORY, DEFAULT_HPARAMS, OPTIMIZERS |
| from .optimizers.sophia import SophiaG |
| from .train import build_optimizer, model_config |
|
|
|
|
| def time_optimizer(name, args, device, warmup=5): |
| torch.manual_seed(0) |
| cfg = model_config(args.vocab_size, args.n_ctx) |
| model = Transformer(cfg).to(device) |
| opt = build_optimizer(name, model, dict(DEFAULT_HPARAMS[name])) |
| is_sophia = isinstance(opt, SophiaG) |
| tok_per_step = args.batch_size * args.n_ctx |
|
|
| x = torch.randint(0, args.vocab_size, (args.batch_size, args.n_ctx), |
| device=device) |
| y = torch.randint(0, args.vocab_size, (args.batch_size, args.n_ctx), |
| device=device) |
|
|
| fwd_bwd, step_t = 0.0, 0.0 |
| for i in range(warmup + args.steps): |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| with torch.autocast("cuda", dtype=torch.bfloat16, |
| enabled=device.type == "cuda"): |
| logits, _, _ = model(x) |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)).float(), |
| y.reshape(-1)) |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| t1 = time.perf_counter() |
|
|
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| if is_sophia: |
| opt.step(bs_scale=tok_per_step) |
| else: |
| opt.step() |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| t2 = time.perf_counter() |
|
|
| if i >= warmup: |
| fwd_bwd += t1 - t0 |
| step_t += t2 - t1 |
|
|
| n_params = sum(p.numel() for p in model.parameters()) |
| state_mb = opt.memory_bytes() / 1e6 |
| del model, opt |
| torch.cuda.empty_cache() |
|
|
| fwd_bwd_ms = 1000 * fwd_bwd / args.steps |
| step_ms = 1000 * step_t / args.steps |
| |
| |
| |
| extra = fwd_bwd_ms / args.sophia_k if is_sophia else 0.0 |
| return { |
| "optimizer": name, |
| "category": CATEGORY[name], |
| "fwd_bwd_ms": fwd_bwd_ms, |
| "step_ms": step_ms, |
| "hessian_amortised_ms": extra, |
| "total_ms": fwd_bwd_ms + step_ms + extra, |
| "state_mb": state_mb, |
| "state_multiple": state_mb * 1e6 / (n_params * 4), |
| } |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--steps", type=int, default=30) |
| ap.add_argument("--batch_size", type=int, default=24) |
| ap.add_argument("--n_ctx", type=int, default=512) |
| ap.add_argument("--vocab_size", type=int, default=16000) |
| ap.add_argument("--sophia_k", type=int, default=10, |
| help="curvature refresh interval, to amortise its cost") |
| args = ap.parse_args() |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"[benchmark] {args.steps} timed steps each, batch " |
| f"{args.batch_size}x{args.n_ctx}, on {device}", flush=True) |
| print(f"\n{'optimizer':10s} {'fwd+bwd':>9s} {'opt.step':>9s} " |
| f"{'hess':>7s} {'total':>9s} {'overhead':>9s} {'state':>8s}") |
| out = [] |
| for name in OPTIMIZERS: |
| r = time_optimizer(name, args, device) |
| out.append(r) |
| over = 100 * (r["step_ms"] + r["hessian_amortised_ms"]) / r["total_ms"] |
| print(f"{name:10s} {r['fwd_bwd_ms']:8.1f}ms {r['step_ms']:8.1f}ms " |
| f"{r['hessian_amortised_ms']:6.1f}ms {r['total_ms']:8.1f}ms " |
| f"{over:8.1f}% {r['state_multiple']:7.1f}x", flush=True) |
|
|
| base = next(r for r in out if r["optimizer"] == "adamw") |
| print(f"\nrelative to AdamW (total ms/step):") |
| for r in out: |
| print(f" {r['optimizer']:10s} {r['total_ms']/base['total_ms']:.3f}x") |
| save_json(out, RUNS / "part2" / "benchmark.json") |
| print(f"\nwrote {RUNS/'part2'/'benchmark.json'}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
