diff --git a/toolkit/optimizers/automagic3.py b/toolkit/optimizers/automagic3.py index b99f2eb..906c72d 100644 --- a/toolkit/optimizers/automagic3.py +++ b/toolkit/optimizers/automagic3.py @@ -36,10 +36,9 @@ class Automagic3(torch.optim.Optimizer): whose update is exactly zero (dead/masked grads, low-precision underflow) carry no direction and abstain from the vote, so a pool of frozen elements can't quietly bias a row's lr upward. Noisy and clean layers each find their - own operating point automatically; the lr - neither collapses to min_lr nor runs away to max_lr. ``lr_bump_rate`` only - sets how fast it gets there, not where it lands. lr is clamped to - [min_lr, max_lr]. + own operating point automatically; the lr neither collapses to min_lr nor + runs away to max_lr. ``lr_bump_rate`` only sets how fast it gets there, not + where it lands. lr is clamped to [min_lr, max_lr]. With ``fused=True`` (default) the step is fused into the backward pass via ``register_post_accumulate_grad_hook``: each parameter is updated and its @@ -57,6 +56,45 @@ class Automagic3(torch.optim.Optimizer): the state is lower precision). Updates to low-precision (e.g. bf16/fp16) parameters are applied in fp32 and stochastically rounded on write-back. + Parameters + ---------- + lr : float + Starting learning rate for every row. The controller adapts away from + this, so it is a launch point, not a tuned target -- a low value just + lets the lr ramp up on its own (there is no warmup). Values above 1e-3 + are rejected and forced back to 1e-6. + min_lr, max_lr : float + Hard clamps on every per-row lr. ``max_lr`` doubles as the per-step + trust region: since the per-element update is capped at + ``clip_threshold`` (~1), a weight moves at most ~``max_lr`` per step, so + keep it modest -- a high ceiling lets the hottest rows take destabilising + steps (and, when fused, the trainer's grad clip can't catch them). + lr_bump_rate : float + Fractional, log-space size of each lr nudge (~10% at 0.1). Sets how fast + the lr moves, NOT where it settles (the flip dynamics fix that); a full + up-nudge and a full down-nudge cancel exactly. + beta2 : float + EMA decay for the second moment, as in Adam/Adafactor. + eps : float + Floor added to the second moment before the rsqrt, to avoid div-by-zero. + clip_threshold : float + Trust region on the update: its RMS is scaled to <= this, then every + element is clamped to +/- this, so no single weight takes an outsized + step. + weight_decay : float + Decoupled (AdamW-style) weight decay; 0 disables it. + lr_smoothing_steps : int + How many steps of the flip signal to EMA-average before nudging the lr + (>=1, default 3). Higher = smoother/slower lr, lower = twitchier/faster; + it does not change where the lr lands. Held as an EMA, so it costs O(1) + state per row regardless of the value. + fused : bool + If True (default), each param is updated inside the backward pass the + moment its grad is ready -- low peak VRAM, but it bypasses the trainer's + grad clipping / nan-skip and cannot be combined with multi-backward + gradient accumulation. If False, a normal ``.step()``-time update, with + low-precision grads accumulated using stochastic rounding. + Improvements over v2 -------------------- 1. Per-row learning rate (was a single scalar per parameter tensor). @@ -95,10 +133,12 @@ class Automagic3(torch.optim.Optimizer): 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, and the lr - scale and parameter update are fused into one ``addcmul_``. Plain English: - each step issues fewer GPU passes over the weights, so it runs faster - (notably in bf16/fp16) without changing the result. + row/col vectors instead of the full gradient-square tensor; the lr scale + and parameter update are fused into one ``addcmul_``; and the per-element + agree/flip vote is a single int8 sign-product (``cur_sign * prev_sign``) + instead of several boolean masks plus float casts. 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__( @@ -165,10 +205,17 @@ class Automagic3(torch.optim.Optimizer): @staticmethod def _rms(t: torch.Tensor) -> torch.Tensor: + # Root-mean-square of a tensor; used to size the trust-region clip. return t.norm(2) / (t.numel() ** 0.5) @staticmethod def _approx_sq_grad(row: torch.Tensor, col: torch.Tensor) -> torch.Tensor: + # Adafactor's factored second moment (inherited from v2). Rather than + # store a full RxC tensor of running grad^2, only its per-row and + # per-col means are kept; this rebuilds the rank-1 approximation of + # 1/sqrt(v) -- the per-element update scale -- as the outer product + # rsqrt(row / mean(row)) (x) rsqrt(col). That is the standard HF + # Adafactor reconstruction and is what keeps optimizer state small. r = (row / row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1) c = col.unsqueeze(-2).rsqrt() return torch.mul(r, c) @@ -297,6 +344,13 @@ class Automagic3(torch.optim.Optimizer): # which saves a full-size kernel pass. sq = grad * grad + # Second moment: a beta2-EMA of grad^2, then update = grad / sqrt(v), + # exactly as Adam/Adafactor (this magnitude-normalises the step; only the + # *sign* of the result drives the lr controller further down). For >=2D + # params v is Adafactor-factored into row/col means (small state, see + # _approx_sq_grad); 1D params (biases, norms) keep the full per-element + # second moment. State lives in p.dtype; when that is low precision the + # math is done in an fp32 copy and written back. if p.dim() >= 2: row_state = state["exp_avg_sq_row"] col_state = state["exp_avg_sq_col"] @@ -348,6 +402,10 @@ class Automagic3(torch.optim.Optimizer): cur_sign = update.sign().to(torch.int8) prev_sign = state["prev_sign"] + # dims: the within-row axes to reduce the per-element vote over (so each + # output channel gets one nudge). lr_b: the per-row lr reshaped to + # broadcast across the full param for the weight update below. For 1D + # params there is no row to reduce over (one lr per element). lr_t = state["lr"] if p.dim() >= 2: dims = tuple(range(1, p.dim())) @@ -438,6 +496,14 @@ class Automagic3(torch.optim.Optimizer): return loss def get_learning_rates(self) -> List[float]: + # Reporting helper: one representative lr per param group. Uses the + # arithmetic mean of the per-row lrs because it is magnitude-weighted -- + # the rows actually taking sizeable steps dominate, so the number reads + # like the scalar lr you'd set in another trainer. (A geometric mean or + # median instead gets dragged toward min_lr by frozen rows and reads + # misleadingly low.) A handful of rows riding at max_lr can lift this + # average even while the typical row is flat; that's expected and bounded + # by the max_lr clamp, not a runaway. out = [] for group in self.param_groups: lrs = [