github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Added Automagic v3

Browse Source
This commit is contained in:
Jaret Burkett
2026-06-07 12:06:42 -06:00
parent 82190b41e6
commit 1dc9a797cf
3 changed files with 378 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
toolkit/optimizer.py
View File

@@ -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

331
toolkit/optimizers/automagic3.py Normal file
View File

@@ -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)

73
toolkit/optimizers/test_optimizers.py
View File

@@ -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__":