""" Optimizer benchmark on a ~500M parameter transformer, run once per dtype (float32, float16, bfloat16). Compares speed (ms/step) and peak VRAM across: - AdamW (torch, unfused — traditional Python loop) - AdamW8bit (bitsandbytes) - Adafactor - Automagic v1 - Automagic v2 (fused-backward) - Automagic v3 (fused-backward and traditional/unfused) - Automagic AdamW (fused-backward and traditional/unfused) - AdamConvRot (AdamW with ConvRot-quantized moments; 8, 4 and 2 bit, traditional/unfused and fused-backward) - Prodigy Pass optimizer keys as CLI args to run a subset (all when none passed), e.g.: python test_optimizers.py adamw8bit automagic3 automagicEXPERIMENT """ import contextlib import gc import io import os import sys import time import torch import torch.nn as nn import torch.nn.functional as F # Allow running this file directly: `python test_optimizers.py` without setting PYTHONPATH. # Toolkit imports happen inside main() so they pick this up. _REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) if _REPO_ROOT not in sys.path: sys.path.insert(0, _REPO_ROOT) # ---- model --------------------------------------------------------------- class TransformerBlock(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int): super().__init__() self.n_heads = n_heads self.d_head = d_model // n_heads self.ln1 = nn.LayerNorm(d_model) self.q = nn.Linear(d_model, d_model, bias=False) self.k = nn.Linear(d_model, d_model, bias=False) self.v = nn.Linear(d_model, d_model, bias=False) self.o = nn.Linear(d_model, d_model, bias=False) self.ln2 = nn.LayerNorm(d_model) self.ffn_up = nn.Linear(d_model, d_ff, bias=False) self.ffn_down = nn.Linear(d_ff, d_model, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: B, S, D = x.shape h = self.ln1(x) q = self.q(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2) k = self.k(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2) v = self.v(h).view(B, S, self.n_heads, self.d_head).transpose(1, 2) a = F.scaled_dot_product_attention(q, k, v, is_causal=True) a = a.transpose(1, 2).contiguous().view(B, S, D) x = x + self.o(a) h = self.ln2(x) x = x + self.ffn_down(F.gelu(self.ffn_up(h))) return x class Transformer(nn.Module): def __init__(self, d_model=1024, n_heads=16, n_layers=40, d_ff=4096): super().__init__() self.blocks = nn.ModuleList([ TransformerBlock(d_model, n_heads, d_ff) for _ in range(n_layers) ]) self.norm = nn.LayerNorm(d_model) def forward(self, x: torch.Tensor) -> torch.Tensor: for b in self.blocks: x = b(x) return self.norm(x) # ---- benchmark ----------------------------------------------------------- DEVICE = "cuda" DTYPES = [torch.float32, torch.float16, torch.bfloat16] D_MODEL = 1024 N_HEADS = 16 N_LAYERS = 40 D_FF = 4096 BATCH = 1 SEQ = 128 WARMUP = 3 ITERS = 10 def build_model(dtype): torch.manual_seed(0) return Transformer(D_MODEL, N_HEADS, N_LAYERS, D_FF).to(DEVICE, dtype=dtype) def benchmark(results: list, label: str, opt_factory, dtype): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() model = build_model(dtype) # Some optimizers print on construction; mute that so the final table is clean. with contextlib.redirect_stdout(io.StringIO()): opt = opt_factory(model.parameters()) x = torch.randn(BATCH, SEQ, D_MODEL, device=DEVICE, dtype=dtype) print(f" running {label}...", flush=True) try: for _ in range(WARMUP): opt.zero_grad(set_to_none=True) model(x).sum().backward() opt.step() torch.cuda.synchronize() t0 = time.perf_counter() for _ in range(ITERS): opt.zero_grad(set_to_none=True) model(x).sum().backward() opt.step() torch.cuda.synchronize() dt = (time.perf_counter() - t0) / ITERS * 1000 peak = torch.cuda.max_memory_allocated() / 1024**3 results.append({"label": label, "ms": dt, "peak": peak, "ok": True}) except torch.cuda.OutOfMemoryError: results.append({"label": label, "ms": float("inf"), "peak": float("inf"), "ok": False, "note": "OOM"}) except Exception as e: # An optimizer may not support a given dtype; record it and keep going. print(f" {label} failed: {type(e).__name__}: {e}", flush=True) results.append({"label": label, "ms": float("inf"), "peak": float("inf"), "ok": False, "note": "ERR"}) finally: # Fused-backward optimizers register post-accumulate-grad hooks on the # params; without removing them the hook closures keep this model alive # past `del`, inflating the *next* fused optimizer's peak. Unregister so # each run reports its true isolated peak regardless of run order. for h in getattr(opt, "_hook_handles", []): h.remove() del opt, model gc.collect() torch.cuda.empty_cache() def print_table(results: list, dtype_name: str): results = sorted(results, key=lambda r: r["peak"]) headers = ["#", "Optimizer", "Peak VRAM", "Time/step"] rows = [] for i, r in enumerate(results, 1): if not r["ok"]: rows.append([str(i), r["label"], r.get("note", "OOM"), "-"]) continue rows.append([str(i), r["label"], f"{r['peak']:.2f} GB", f"{r['ms']:.1f} ms"]) widths = [max(len(str(row[c])) for row in [headers] + rows) for c in range(len(headers))] def fmt(row, sep=" │ "): return sep.join(s.ljust(widths[c]) if c == 1 else s.rjust(widths[c]) for c, s in enumerate(row)) line_top = "─" * (sum(widths) + 3 * (len(widths) - 1)) print() print(f" dtype: {dtype_name}") print(line_top) print(fmt(headers)) print(line_top) for row in rows: print(fmt(row)) print(line_top) def main(): n_params = sum(p.numel() for p in build_model(torch.float32).parameters()) print(f"Model: {N_LAYERS} blocks × d_model={D_MODEL} × d_ff={D_FF}") print(f" {n_params/1e6:.1f}M params") print(f"Step: batch={BATCH}, seq={SEQ}") print(f"Timing: {WARMUP} warmup + {ITERS} timed iters") print(f"Rounds: {', '.join(str(d).replace('torch.', '') for d in DTYPES)}") from toolkit.optimizers.automagic import Automagic from toolkit.optimizers.automagic2 import Automagic2 from toolkit.optimizers.automagic3 import Automagic3 from toolkit.optimizers.automagicEXPERIMENT import AutomagicEXPERIMENT from toolkit.optimizers.adamconvrot import AdamConvRot from toolkit.optimizers.adafactor import Adafactor from prodigyopt import Prodigy import bitsandbytes as bnb optimizers = [ ("adamw", "AdamW", lambda p: torch.optim.AdamW(p, lr=1e-4, eps=1e-6, foreach=False, fused=False)), ("adamw8bit", "AdamW8bit", lambda p: bnb.optim.AdamW8bit(p, lr=1e-4, eps=1e-6)), ("adafactor", "Adafactor", lambda p: Adafactor(p, lr=1e-4, scale_parameter=False, relative_step=False, warmup_init=False)), ("automagic", "Automagic v1", lambda p: Automagic(p, lr=1e-4)), ("automagic2", "Automagic v2", lambda p: Automagic2(p, lr=1e-4)), ("automagic3", "Automagic v3 fused", lambda p: Automagic3(p, lr=1e-4, fused=True)), ("automagic3", "Automagic v3 unfused", lambda p: Automagic3(p, lr=1e-4, fused=False)), ("automagicEXPERIMENT", "AutomagicEXPERIMENT", lambda p: AutomagicEXPERIMENT(p, lr=1e-4, fused=True)), ("automagicEXPERIMENT", "AutomagicEXPERIMENT", lambda p: AutomagicEXPERIMENT(p, lr=1e-4, fused=False)), ("adamconvrot", "AdamConvRot 8bit", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrot8")), ("adamconvrot", "AdamConvRot 4bit", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint4")), ("adamconvrot", "AdamConvRot 2bit", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint2")), ("adamconvrot", "AdamConvRot 8bit fused", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrot8", fused=True)), ("adamconvrot", "AdamConvRot 4bit fused", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint4", fused=True)), ("adamconvrot", "AdamConvRot 2bit fused", lambda p: AdamConvRot(p, lr=1e-4, eps=1e-6, qtype="convrotint2", fused=True)), ("prodigy", "Prodigy", lambda p: Prodigy(p, lr=1.0, eps=1e-6)), ] selected = [a.lower() for a in sys.argv[1:]] if selected: valid = {key.lower() for key, _, _ in optimizers} unknown = [a for a in selected if a not in valid] if unknown: print(f"Unknown optimizer keys: {', '.join(unknown)}") print(f"Valid keys: {', '.join(sorted(valid))}") sys.exit(1) optimizers = [o for o in optimizers if o[0].lower() in selected] for dtype in DTYPES: dtype_name = str(dtype).replace("torch.", "") print(f"\n=== Round: {dtype_name} ===") results: list = [] for _, label, factory in optimizers: benchmark(results, label, factory, dtype) print_table(results, dtype_name) if __name__ == "__main__": main()