From 1dc9a797cf4c849f776d75f67e4fb939050aae0d Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 7 Jun 2026 12:06:42 -0600 Subject: [PATCH] Added Automagic v3 --- toolkit/optimizer.py | 3 + toolkit/optimizers/automagic3.py | 331 ++++++++++++++++++++++++++ toolkit/optimizers/test_optimizers.py | 73 +++--- 3 files changed, 378 insertions(+), 29 deletions(-) create mode 100644 toolkit/optimizers/automagic3.py diff --git a/toolkit/optimizer.py b/toolkit/optimizer.py index 0357f79..433106b 100644 --- a/toolkit/optimizer.py +++ b/toolkit/optimizer.py @@ -100,6 +100,9 @@ def get_optimizer( elif lower_type == 'automagic2': from toolkit.optimizers.automagic2 import Automagic2 optimizer = Automagic2(params, lr=float(learning_rate), **optimizer_params) + elif lower_type == 'automagic3': + from toolkit.optimizers.automagic3 import Automagic3 + optimizer = Automagic3(params, lr=float(learning_rate), **optimizer_params) else: raise ValueError(f'Unknown optimizer type {optimizer_type}') return optimizer diff --git a/toolkit/optimizers/automagic3.py b/toolkit/optimizers/automagic3.py new file mode 100644 index 0000000..70810b2 --- /dev/null +++ b/toolkit/optimizers/automagic3.py @@ -0,0 +1,331 @@ +from typing import List +import torch + + +class Automagic3(torch.optim.Optimizer): + """ + Automagic v3. + + A learning rate is kept per row of each parameter: one lr per output + channel for >=2D weights (e.g. one lr per output neuron of a Linear layer) + and one lr per element for 1D weights (biases, norms). Each step the lr is + nudged by how strongly the per-element update direction agrees with the + previous step. Agreement is the plain fraction of elements that kept their + sign (every element counts equally, regardless of update magnitude) and is + used as a proportional signal in [-1, 1] that scales a multiplicative + (geometric) bump ``lr *= exp(signal * lr_bump_rate)``, so each step is a fixed + fractional move that behaves uniformly across the whole [min_lr, max_lr] + range and a full up bump is exactly undone by a full down bump. The lr + drifts smoothly, self-stabilises near random agreement, and is clamped to + [min_lr, max_lr]. + + The optimizer step is fused into the backward pass via + ``register_post_accumulate_grad_hook``: each parameter is updated and its + grad freed as soon as autograd finishes accumulating into it. ``.step()`` + therefore does no real work and peak VRAM stays low. + + Second-moment EMA state is stored in ``p.dtype`` (math runs in fp32 when + the state is lower precision). Updates to low-precision (e.g. bf16/fp16) + parameters are applied in fp32 and stochastically rounded on write-back. + + Improvements over v2 + -------------------- + 1. Per-row learning rate (was a single scalar per parameter tensor). + v2 kept one lr for an entire weight matrix; v3 keeps one per output + channel (per element for 1D params). Plain English: different neurons in + the same layer can now learn at different speeds instead of being forced + to share one rate, so a layer where some rows have converged and others + have not is handled gracefully. + + 2. Proportional lr control (was a hard threshold flip). v2 bumped the lr up + or down by a fixed amount depending on whether agreement crossed a + threshold, which jitters when agreement hovers near the boundary. v3 + scales the bump by how strongly the directions agree + (``2*agreement - 1`` in [-1, 1]). Plain English: the lr nudges gently + when the signal is weak and firmly when it is strong, and parks itself + instead of oscillating when gradients are basically noise. + + 3. Multiplicative (geometric) lr bump (was additive). v2 added/subtracted a + fixed absolute amount, so the same bump was a huge relative jump near + min_lr and a negligible one near max_lr. v3 multiplies by + ``exp(signal * lr_bump_rate)`` -- a fixed *percentage* step. Plain + English: the lr moves at the same relative pace whether it is tiny or + large, traverses its whole range in a predictable number of steps, and a + full up bump is exactly cancelled by a full down bump (no drift). The + knob was renamed ``lr_bump`` -> ``lr_bump_rate`` to signal the change. + + 4. Stochastic rounding for fp16, not just bf16. v2 only rounded bf16 + write-backs and let fp16 fall back to round-to-nearest, silently + discarding updates smaller than an fp16 ULP. v3 stochastically rounds + both (fast bit-trick for bf16/fp16, generic fallback for other low + precisions). Plain English: fp16 training no longer throws away small + weight updates, so it actually keeps learning instead of stalling. + + 5. Faster hot path, identical math. eps is folded into the small reduced + row/col vectors instead of the full gradient-square tensor; the lr scale, + weight decay and parameter update are fused into one ``addcmul_``; and the + sign-agreement is summed straight off the bool mask with no full-size + float cast. Plain English: each step issues fewer GPU passes over the + weights, so it runs faster (notably in bf16/fp16) without changing the + result. + """ + + def __init__( + self, + params, + lr: float = 1e-6, + min_lr: float = 1e-8, + max_lr: float = 1e-2, + lr_bump_rate: float = 0.1, # fractional/log step per bump (~10%); see step logic + beta2: float = 0.999, + eps: float = 1e-30, + clip_threshold: float = 1.0, + weight_decay: float = 0.0, + ): + if lr > 1e-3: + print(f"Warning! Start lr {lr} is very high; forcing to 1e-6.") + lr = 1e-6 + defaults = dict( + lr=lr, + min_lr=min_lr, + max_lr=max_lr, + lr_bump_rate=lr_bump_rate, + beta2=beta2, + eps=eps, + clip_threshold=clip_threshold, + weight_decay=weight_decay, + ) + super().__init__(params, defaults) + + self._hook_handles = [] + for group in self.param_groups: + for p in group["params"]: + if p.requires_grad: + handle = p.register_post_accumulate_grad_hook( + self._make_backward_hook(group) + ) + self._hook_handles.append(handle) + + total = sum(p.numel() for g in self.param_groups for p in g["params"]) + print(f"Total training paramiters: {total:,}") + + # ------------------------------------------------------------------ utils + + @staticmethod + def _rms(t: torch.Tensor) -> torch.Tensor: + return t.norm(2) / (t.numel() ** 0.5) + + @staticmethod + def _approx_sq_grad(row: torch.Tensor, col: torch.Tensor) -> torch.Tensor: + r = (row / row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1) + c = col.unsqueeze(-2).rsqrt() + return torch.mul(r, c) + + @staticmethod + def _sr_truncate(v_fp32: torch.Tensor, drop_bits: int) -> torch.Tensor: + # Fast in-place stochastic rounding for a low-precision float that is a + # mantissa truncation of fp32: add uniform noise into the dropped low + # mantissa bits of the fp32 bit pattern, then zero them, so the + # subsequent narrowing cast is exact and rounds up with probability + # equal to the truncated fractional part. bf16 drops 16 bits (it is the + # high half of fp32); fp16 drops 13 bits (23 - 10 mantissa) and is exact + # within its normal exponent range -- values past fp16's overflow / + # subnormal limits are rounded at fp32 granularity, which trained + # weights effectively never reach. + as_int = v_fp32.view(torch.int32) + as_int.add_(torch.randint_like(as_int, 1 << drop_bits)) + as_int.bitwise_and_(-(1 << drop_bits)) + return v_fp32 + + @staticmethod + def _stochastic_round(v: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + # Generic fp32 -> low-precision stochastic rounding for dtypes that are + # not a mantissa truncation of fp32 (the bf16/fp16 fast path in + # _sr_truncate does not apply). Adds uniform noise of +/- half a target + # ULP and rounds to nearest, so P(round up) equals the fractional + # distance to the next representable value -> unbiased in expectation. + # The ULP at |v| is 2**floor(log2|v|) * eps(dtype). + finfo = torch.finfo(dtype) + absv = v.abs().clamp_(min=finfo.tiny) + ulp = torch.exp2(torch.floor(torch.log2(absv))).mul_(finfo.eps) + noise = torch.rand_like(v).sub_(0.5).mul_(ulp) + return v.add_(noise).to(dtype) + + def _init_state(self, p: torch.Tensor, group: dict) -> None: + state = self.state[p] + state["step"] = 0 + # Per-row lr: one entry per output channel for >=2D params, per element + # for 1D params, a scalar for 0D params. + lr_shape = (p.shape[0],) if p.dim() >= 2 else p.shape + state["lr"] = torch.full( + lr_shape, float(group["lr"]), dtype=torch.float32, device=p.device + ) + state["last_polarity"] = torch.zeros(p.shape, dtype=torch.bool, device=p.device) + if p.dim() >= 2: + state["exp_avg_sq_row"] = torch.zeros( + p.shape[:-1], dtype=p.dtype, device=p.device + ) + state["exp_avg_sq_col"] = torch.zeros( + p.shape[:-2] + p.shape[-1:], dtype=p.dtype, device=p.device + ) + else: + state["exp_avg_sq"] = torch.zeros(p.shape, dtype=p.dtype, device=p.device) + + def _make_backward_hook(self, group): + def _hook(p: torch.Tensor): + self._update_param(p, group) + + return _hook + + # -------------------------------------------------------------- per-param + + @torch.no_grad() + def _update_param(self, p: torch.Tensor, group: dict) -> None: + if p.grad is None: + return + state = self.state[p] + if len(state) == 0: + self._init_state(p, group) + + grad = p.grad + if grad.is_sparse: + raise RuntimeError("Automagic3 does not support sparse gradients.") + if grad.dtype != torch.float32: + grad = grad.to(torch.float32) + + beta2 = group["beta2"] + eps = group["eps"] + # eps is folded into the reduced row/col (or rsqrt) instead of being + # added to the full-size sq tensor: mean(g^2 + eps) == mean(g^2) + eps, + # which saves a full-size kernel pass. + sq = grad * grad + + if p.dim() >= 2: + row_state = state["exp_avg_sq_row"] + col_state = state["exp_avg_sq_col"] + if row_state.dtype == torch.float32: + row, col = row_state, col_state + row.mul_(beta2).add_(sq.mean(dim=-1).add_(eps), alpha=1.0 - beta2) + col.mul_(beta2).add_(sq.mean(dim=-2).add_(eps), alpha=1.0 - beta2) + else: + row = row_state.to(torch.float32) + col = col_state.to(torch.float32) + row.mul_(beta2).add_(sq.mean(dim=-1).add_(eps), alpha=1.0 - beta2) + col.mul_(beta2).add_(sq.mean(dim=-2).add_(eps), alpha=1.0 - beta2) + row_state.copy_(row.to(row_state.dtype)) + col_state.copy_(col.to(col_state.dtype)) + update = self._approx_sq_grad(row, col).mul_(grad) + else: + v_state = state["exp_avg_sq"] + if v_state.dtype == torch.float32: + v = v_state + v.mul_(beta2).add_(sq, alpha=1.0 - beta2) + else: + v = v_state.to(torch.float32) + v.mul_(beta2).add_(sq, alpha=1.0 - beta2) + v_state.copy_(v.to(v_state.dtype)) + update = v.add(eps).rsqrt().mul_(grad) + + # Update-RMS clip (trust region): scale so the update RMS never exceeds + # clip_threshold. No bias-correction warmup -- LoRA runs are short and a + # slow ramp wastes steps; for a soft start the user can set a low start + # lr and let the lr bump up on its own. + update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0)) + + # Per-row sign agreement vs. the previous step: the plain fraction of + # elements that kept their sign, every element counting equally + # regardless of update magnitude. For >=2D params this is reduced to one + # value per output channel; for 1D it is per element. The result drives + # a proportional lr bump in [-1, 1]. The match is summed straight off the + # bool tensor (no full-size float cast) into the per-row reduction. + cur_polarity = update > 0 + eqb = cur_polarity == state["last_polarity"] + state["last_polarity"] = cur_polarity + + lr_t = state["lr"] + if p.dim() >= 2: + dims = tuple(range(1, p.dim())) + agreement = eqb.sum(dim=dims, dtype=torch.float32).div_( + eqb.shape[1:].numel() + ) + lr_b = lr_t.view(lr_t.shape[0], *([1] * (p.dim() - 1))) + else: + agreement = eqb.to(torch.float32) + lr_b = lr_t + + if state["step"] > 0: + direction = agreement.mul_(2.0).sub_(1.0) + # Multiplicative (geometric) bump: lr *= exp(direction * lr_bump_rate). + # A full up step multiplies lr by exp(lr_bump_rate), a full down step by + # its reciprocal, so up and down are symmetric in log space (a full + # up bump is exactly undone by a full down bump). lr_bump_rate is thus a + # fractional rate (~lr_bump_rate per step for small values), giving a + # uniform relative move at every scale across [min_lr, max_lr] + # instead of a fixed absolute amount. + lr_t.mul_(torch.exp(direction.mul_(group["lr_bump_rate"]))).clamp_( + min=group["min_lr"], max=group["max_lr"] + ) + state["step"] += 1 + + wd = group["weight_decay"] + + if p.dtype == torch.float32: + # Decoupled weight decay folded in (update += wd*p), then a single + # fused p -= lr_b * update. + if wd != 0.0: + update.add_(p, alpha=wd) + p.addcmul_(update, lr_b, value=-1.0) + else: + # Low precision: apply the update in fp32 then stochastically round + # back, so tiny updates aren't lost to round-to-nearest. Single + # bf16/fp16 -> fp32 conversion shared by weight decay and rounding. + new_p_fp32 = p.to(torch.float32) + if wd != 0.0: + update.add_(new_p_fp32, alpha=wd) + new_p_fp32.addcmul_(update, lr_b, value=-1.0) + if p.dtype == torch.bfloat16: + p.copy_(self._sr_truncate(new_p_fp32, 16)) + elif p.dtype == torch.float16: + p.copy_(self._sr_truncate(new_p_fp32, 13)) + else: + p.copy_(self._stochastic_round(new_p_fp32, p.dtype)) + + p.grad = None + + # ----------------------------------------------------------- optimizer API + + @torch.no_grad() + def step(self, closure=None): + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + return loss + + def get_learning_rates(self) -> List[float]: + out = [] + for group in self.param_groups: + lrs = [ + float(self.state[p]["lr"].mean()) + for p in group["params"] + if p in self.state and "lr" in self.state[p] + ] + out.append(sum(lrs) / len(lrs) if lrs else float(group["lr"])) + return out + + def get_avg_learning_rate(self) -> float: + lrs = self.get_learning_rates() + return sum(lrs) / len(lrs) if lrs else float(self.defaults["lr"]) + + def load_state_dict(self, state_dict): + # Parent casts every fp state tensor to param.dtype; force lr back to fp32 + # so subsequent lr bumps aren't rounded away on bf16 weights. + super().load_state_dict(state_dict) + # Constructor args always win over whatever was saved in the checkpoint. + for group in self.param_groups: + for k, v in self.defaults.items(): + group[k] = v + for p in group["params"]: + st = self.state.get(p) + if st is not None and isinstance(st.get("lr"), torch.Tensor): + st["lr"] = st["lr"].to(torch.float32) diff --git a/toolkit/optimizers/test_optimizers.py b/toolkit/optimizers/test_optimizers.py index 7d81e15..8a36f37 100644 --- a/toolkit/optimizers/test_optimizers.py +++ b/toolkit/optimizers/test_optimizers.py @@ -1,12 +1,14 @@ """ -Optimizer benchmark on a ~500M parameter fp32 transformer. +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 (only optimizer using fused-backward) + - Automagic v2 (fused-backward) + - Automagic v3 (fused-backward) - Prodigy """ import contextlib @@ -74,7 +76,7 @@ class Transformer(nn.Module): # ---- benchmark ----------------------------------------------------------- DEVICE = "cuda" -DTYPE = torch.float32 +DTYPES = [torch.float32, torch.float16, torch.bfloat16] D_MODEL = 1024 N_HEADS = 16 N_LAYERS = 40 @@ -85,21 +87,21 @@ WARMUP = 3 ITERS = 10 -def build_model(): +def build_model(dtype): torch.manual_seed(0) - return Transformer(D_MODEL, N_HEADS, N_LAYERS, D_FF).to(DEVICE, dtype=DTYPE) + return Transformer(D_MODEL, N_HEADS, N_LAYERS, D_FF).to(DEVICE, dtype=dtype) -def benchmark(results: list, label: str, opt_factory): +def benchmark(results: list, label: str, opt_factory, dtype): gc.collect() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() - model = build_model() + 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) + x = torch.randn(BATCH, SEQ, D_MODEL, device=DEVICE, dtype=dtype) print(f" running {label}...", flush=True) try: @@ -119,21 +121,31 @@ def benchmark(results: list, label: str, opt_factory): 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}) + 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): +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"], "OOM", "-"]) + 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"]) @@ -144,6 +156,7 @@ def print_table(results: list): line_top = "─" * (sum(widths) + 3 * (len(widths) - 1)) print() + print(f" dtype: {dtype_name}") print(line_top) print(fmt(headers)) print(line_top) @@ -153,35 +166,37 @@ def print_table(results: list): def main(): - n_params = sum(p.numel() for p in build_model().parameters()) - dtype_name = str(DTYPE).replace("torch.", "") + 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 ({dtype_name})") + print(f" {n_params/1e6:.1f}M params") print(f"Step: batch={BATCH}, seq={SEQ}") print(f"Timing: {WARMUP} warmup + {ITERS} timed iters") - print() + 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.adafactor import Adafactor from prodigyopt import Prodigy import bitsandbytes as bnb - results: list = [] - benchmark(results, "AdamW", - lambda p: torch.optim.AdamW(p, lr=1e-4, eps=1e-6, foreach=False, fused=False)) - benchmark(results, "AdamW8bit", - lambda p: bnb.optim.AdamW8bit(p, lr=1e-4, eps=1e-6)) - benchmark(results, "Adafactor", - lambda p: Adafactor(p, lr=1e-4, scale_parameter=False, relative_step=False, warmup_init=False)) - benchmark(results, "Automagic v1", - lambda p: Automagic(p, lr=1e-4)) - benchmark(results, "Automagic v2", - lambda p: Automagic2(p, lr=1e-4)) - benchmark(results, "Prodigy", - lambda p: Prodigy(p, lr=1.0, eps=1e-6)) + optimizers = [ + ("AdamW", lambda p: torch.optim.AdamW(p, lr=1e-4, eps=1e-6, foreach=False, fused=False)), + ("AdamW8bit", lambda p: bnb.optim.AdamW8bit(p, lr=1e-4, eps=1e-6)), + ("Adafactor", lambda p: Adafactor(p, lr=1e-4, scale_parameter=False, relative_step=False, warmup_init=False)), + ("Automagic v1", lambda p: Automagic(p, lr=1e-4)), + ("Automagic v2", lambda p: Automagic2(p, lr=1e-4)), + ("Automagic v3", lambda p: Automagic3(p, lr=1e-4)), + ("Prodigy", lambda p: Prodigy(p, lr=1.0, eps=1e-6)), + ] - print_table(results) + 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__":