562 lines
29 KiB
Python
562 lines
29 KiB
Python
"""
|
|
NOTE: This is experimental and under active development; expect breaking changes and bugs. Feedback welcome.
|
|
"""
|
|
|
|
from typing import List
|
|
import torch
|
|
|
|
|
|
class Automagic3(torch.optim.Optimizer):
|
|
"""
|
|
Automagic v3.
|
|
|
|
A single learning rate is kept per parameter tensor (one lr per weight
|
|
matrix / layer). Each step the lr is nudged by whether the per-element update
|
|
direction *flipped* vs the previous step (RProp-style edge-of-stability
|
|
control), aggregated over the whole tensor.
|
|
|
|
A sign flip means the step jumped past the local minimum (overshoot) -- the
|
|
one event whose frequency genuinely rises with the lr, so it provides a true
|
|
restoring force. Each element votes: agree with last step -> +``lr_bump_rate``;
|
|
flip -> -``lr_bump_rate`` (symmetric). The votes are reduced to ONE value for
|
|
the whole tensor and EMA-smoothed over ~``lr_smoothing_steps`` steps, then
|
|
applied multiplicatively: ``lr *= exp(nudge)``. Reducing over the whole tensor
|
|
(rather than per row/channel) is deliberate: per-channel lrs let coupled
|
|
channels fight -- one drives its lr up while a neighbour drives its down to
|
|
compensate -- and split to opposite extremes, which visibly wrecks the model;
|
|
one lr per tensor makes opposing channels cancel into a single vote.
|
|
|
|
The flip equilibrium is flip fraction == 0.5, the only rate that is both the
|
|
pure-noise point and the edge of stability: a tensor descending cleanly flips
|
|
less than half the time -> its lr grows; once it overshoots it flips more than
|
|
half -> its lr shrinks; pure noise flips ~half -> the lr holds. Symmetry is
|
|
load-bearing (any asymmetry drags a noise-dominated tensor to zero). Elements
|
|
with an exactly-zero update (dead/masked grads, low-precision underflow) carry
|
|
no direction and abstain from the vote.
|
|
|
|
On top of the per-tensor flip control, every tensor's lr is gently pulled
|
|
toward the GLOBAL average lr each step (``lr_pull``, geometric / log-space:
|
|
``lr *= (avg_lr / lr) ** lr_pull``). This mean-reversion is the restoring
|
|
force that replaces hard min/max clamps: layers may settle at their own level,
|
|
but cannot drift apart to opposite extremes (one frozen near zero, one running
|
|
away) -- the failure mode that destroys full finetunes. The target is the
|
|
emergent average across layers, not a fixed number, so it stays fully
|
|
automatic; there are no min/max lr bounds. ``lr_bump_rate`` sets how fast the
|
|
lr moves, not where it lands.
|
|
|
|
With ``fused=True`` (default) the 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. Note this bypasses the
|
|
trainer's grad clipping / nan-skip (they run after backward) and is not
|
|
compatible with multi-backward gradient accumulation.
|
|
|
|
With ``fused=False`` it behaves like a traditional optimizer: grads
|
|
accumulate across backward passes and the update happens in ``.step()``.
|
|
Low-precision (bf16/fp16) grads are accumulated with stochastic rounding so
|
|
small per-micro-batch grads aren't lost; fp32 grads accumulate normally.
|
|
|
|
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.
|
|
|
|
Parameters
|
|
----------
|
|
lr : float
|
|
Starting learning rate for every layer. 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. There are no min/max lr clamps;
|
|
the mean-reversion pull keeps the spread bounded instead.
|
|
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.
|
|
lr_pull : float
|
|
Per-step strength of the mean-reversion pulling each layer's lr toward
|
|
the global average (log space; 0 disables). This is the restoring force
|
|
that replaces min/max clamps and stops layers drifting to opposite
|
|
extremes. Small (default 0.05) lets layers keep their own level while
|
|
bounding the spread; larger forces all layers toward a common lr.
|
|
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 layer 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-layer lr with mean-reversion (v2 had one static lr per tensor and no
|
|
coupling between layers). v3 still keeps one lr per tensor, but it is
|
|
adaptive (driven by the flip controller) and every layer's lr is pulled
|
|
toward the global average (``lr_pull``). Plain English: each layer finds
|
|
its own learning rate automatically, but no layer can run away or freeze
|
|
relative to the others -- which is what used to split a full finetune into
|
|
over-cooked and dead layers and destroy it. (An earlier v3 used a separate
|
|
lr per output channel; coupled channels fought and split to opposite
|
|
extremes, so it was reduced back to one lr per tensor.)
|
|
|
|
2. Overshoot-based (RProp-style) lr control with a real equilibrium. v2
|
|
bumped the lr from raw direction agreement, which has no upper fixed point
|
|
-- a parameter that is simply still descending keeps agreeing at any lr,
|
|
so the lr ratchets up and eventually runs away on long runs. v3 drives the
|
|
lr from sign *flips* (overshoot) instead, nudging up on agree and down on
|
|
flip symmetrically; the equilibrium is a flip fraction of 0.5, which is
|
|
both the noise point and the edge of stability. Plain English: the lr
|
|
speeds up while a layer is making clean progress, backs off the moment it
|
|
starts overshooting, and simply holds when the gradient is pure noise --
|
|
so it neither climbs without bound on long runs nor collapses to nothing
|
|
on a fresh, noisy LoRA.
|
|
|
|
3. Multiplicative (geometric) lr bump (was additive). v2 added/subtracted a
|
|
fixed absolute amount, so the same bump was a huge relative jump when the
|
|
lr was tiny and a negligible one when it was large. 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
|
|
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__(
|
|
self,
|
|
params,
|
|
lr: float = 1e-6,
|
|
lr_bump_rate: float = 0.1, # fractional/log step per bump (~10%); see step logic
|
|
lr_pull: float = 0.025, # per-step pull of each layer's lr toward the global avg (log space)
|
|
beta2: float = 0.999,
|
|
eps: float = 1e-30,
|
|
clip_threshold: float = 1.0,
|
|
weight_decay: float = 0.0,
|
|
lr_smoothing_steps: int = 3, # lr-nudge EMA smoothing horizon, in steps (min 1)
|
|
fused: bool = True,
|
|
):
|
|
if lr > 1e-3:
|
|
print(f"Warning! Start lr {lr} is very high; forcing to 1e-6.")
|
|
lr = 1e-6
|
|
# The lr nudge is EMA-smoothed over ~this many steps; at least 1.
|
|
lr_smoothing_steps = max(1, int(lr_smoothing_steps))
|
|
defaults = dict(
|
|
lr=lr,
|
|
lr_bump_rate=lr_bump_rate,
|
|
lr_pull=max(0.0, float(lr_pull)),
|
|
beta2=beta2,
|
|
eps=eps,
|
|
clip_threshold=clip_threshold,
|
|
weight_decay=weight_decay,
|
|
lr_smoothing_steps=lr_smoothing_steps,
|
|
# EMA decay for the per-layer lr nudge, derived from the smoothing
|
|
# horizon (n steps -> beta = n/(n+1)).
|
|
dir_beta=lr_smoothing_steps / (lr_smoothing_steps + 1.0),
|
|
)
|
|
super().__init__(params, defaults)
|
|
|
|
self.fused = fused
|
|
# Global geometric-mean lr across all layers; the mean-reversion pull
|
|
# targets this. Seeded at the start lr (all layers equal) and refreshed
|
|
# each .step() from the current per-layer lrs.
|
|
self._avg_lr = float(lr)
|
|
self._hook_handles = []
|
|
for group in self.param_groups:
|
|
for p in group["params"]:
|
|
if not p.requires_grad:
|
|
continue
|
|
if self.fused:
|
|
# Fused: update each param the moment its grad is ready.
|
|
handle = p.register_post_accumulate_grad_hook(
|
|
self._make_backward_hook(group)
|
|
)
|
|
self._hook_handles.append(handle)
|
|
elif p.dtype != torch.float32:
|
|
# Non-fused: the actual update happens in .step(); here we
|
|
# only stochastically accumulate low-precision grads across
|
|
# micro-batches so repeated round-to-nearest doesn't drop
|
|
# small grads (fp32 grads accumulate losslessly on their own).
|
|
handle = p.register_post_accumulate_grad_hook(
|
|
self._make_accum_hook()
|
|
)
|
|
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:
|
|
# 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)
|
|
|
|
@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)
|
|
|
|
@classmethod
|
|
def _stochastic_copy_(cls, dst: torch.Tensor, src_fp32: torch.Tensor) -> None:
|
|
# Stochastically round the fp32 ``src`` into the low-precision ``dst`` in
|
|
# place. Uses the fast mantissa-truncation path for bf16/fp16 and the
|
|
# generic method otherwise. ``src_fp32`` may be mutated (caller owns it).
|
|
if dst.dtype == torch.bfloat16:
|
|
dst.copy_(cls._sr_truncate(src_fp32, 16))
|
|
elif dst.dtype == torch.float16:
|
|
dst.copy_(cls._sr_truncate(src_fp32, 13))
|
|
else:
|
|
dst.copy_(cls._stochastic_round(src_fp32, dst.dtype))
|
|
|
|
def _make_accum_hook(self):
|
|
# Non-fused grad accumulation for low-precision params: accumulate the
|
|
# running sum in fp32 then stochastically round it back into the
|
|
# low-precision ``_accum_grad`` buffer, so small per-micro-batch grads
|
|
# are not lost to repeated round-to-nearest. .step() consumes the buffer.
|
|
def _hook(p: torch.Tensor):
|
|
if p.grad is None:
|
|
return
|
|
if hasattr(p, "_accum_grad"):
|
|
acc = p._accum_grad.to(torch.float32).add_(p.grad.to(torch.float32))
|
|
self._stochastic_copy_(p._accum_grad, acc)
|
|
else:
|
|
p._accum_grad = p.grad.clone()
|
|
p.grad = None
|
|
|
|
return _hook
|
|
|
|
def _init_state(self, p: torch.Tensor, group: dict) -> None:
|
|
state = self.state[p]
|
|
state["step"] = 0
|
|
# ONE lr per parameter tensor (scalar), not per row/channel. Per-channel
|
|
# lrs let coupled channels fight -- one drives up while another drives
|
|
# down to compensate -- so they split to the rails and wreck the model.
|
|
# A single lr per tensor averages the vote over all elements, so opposing
|
|
# channels cancel instead of diverging.
|
|
state["lr"] = torch.tensor(
|
|
float(group["lr"]), dtype=torch.float32, device=p.device
|
|
)
|
|
# Previous update-sign snapshot (int8 {-1, 0, +1}, full param shape); the
|
|
# current sign is compared against it to detect per-element flips. Set on
|
|
# the first step.
|
|
state["prev_sign"] = None
|
|
# EMA of the (scalar) log lr-nudge, smoothing the flip signal over time.
|
|
state["dir_ema"] = torch.zeros((), dtype=torch.float32, 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)
|
|
|
|
# This step is fused into backward, so the trainer's grad clipping and
|
|
# nan/inf-skip run too late to protect us -- the weights are already
|
|
# updated here. A single non-finite gradient would poison the
|
|
# second-moment EMA (NaN*beta2 + ... stays NaN forever) and corrupt the
|
|
# weights, which surfaces as the model "randomly" blowing up. Neutralise
|
|
# non-finite grads in place (we own this fp32 grad) so those elements
|
|
# contribute nothing this step instead of destroying state. Large but
|
|
# finite grads are left alone -- the second-moment normalisation already
|
|
# bounds their effect.
|
|
grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
|
|
|
|
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
|
|
|
|
# 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"]
|
|
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))
|
|
# The RMS clip only bounds the aggregate, so a single outlier element can
|
|
# still survive at ~sqrt(numel)*clip_threshold and hit one weight hard,
|
|
# distorting the model. Cap each element to clip_threshold (a true
|
|
# max-norm trust region) so no single weight can take an outsized step.
|
|
update.clamp_(-group["clip_threshold"], group["clip_threshold"])
|
|
|
|
# RProp-style edge-of-stability lr control. The signal is whether each
|
|
# element's update direction *flipped* vs the previous step, not how
|
|
# steady it has been: a flip means we stepped past the local minimum
|
|
# (overshoot), the one event whose frequency actually rises with the lr,
|
|
# so it gives a true restoring force. Steadiness does not -- a parameter
|
|
# descending monotonically agrees with itself at any non-overshooting lr,
|
|
# which is why a consistency-vs-noise-floor signal has no upper
|
|
# equilibrium and runs away on long tunes.
|
|
# Trinary sign {-1, 0, +1}: zero updates (dead/masked grads, flat
|
|
# activation regions, low-precision underflow) are kept distinct from
|
|
# negatives rather than bucketed with them by a bare ``> 0``.
|
|
cur_sign = update.sign().to(torch.int8)
|
|
prev_sign = state["prev_sign"]
|
|
lr_t = state["lr"] # scalar (one lr for the whole tensor)
|
|
|
|
if prev_sign is not None:
|
|
# Per-element vote via the sign product. With signs in {-1, 0, +1},
|
|
# cur_sign * prev_sign is +1 when the direction held (agree), -1 when
|
|
# it flipped (overshoot), and 0 whenever either step's update was zero
|
|
# -- so a zero update automatically ABSTAINS (contributes nothing and
|
|
# isn't counted), no separate masking needed.
|
|
#
|
|
# Reduced over the WHOLE tensor (not per row) this is
|
|
# bump*(1 - 2*flip_fraction): the lr grows while the layer mostly
|
|
# holds its direction, shrinks once it mostly flips, and holds at the
|
|
# flip_fraction == 0.5 point. Aggregating over all elements means
|
|
# channels pushing opposite ways cancel into one vote instead of
|
|
# splitting to opposite lr extremes.
|
|
bump = group["lr_bump_rate"]
|
|
prod = cur_sign * prev_sign # int8 {-1, 0, +1} per element
|
|
# Reduce directly on the int8 product (sum -> int64, no full-size
|
|
# float cast; count_nonzero -> valid votes) -- two reductions instead
|
|
# of casting the whole tensor to float twice.
|
|
num = prod.sum()
|
|
den = prod.count_nonzero().clamp_(min=1)
|
|
log_dir = num.float().div_(den.float()).mul_(bump)
|
|
# EMA-smooth the nudge so a single noisy step doesn't swing the lr,
|
|
# then apply it multiplicatively (a fixed fractional move).
|
|
ema = state["dir_ema"]
|
|
beta = group["dir_beta"]
|
|
ema.mul_(beta).add_(log_dir, alpha=1.0 - beta)
|
|
lr_t.mul_(torch.exp(ema))
|
|
# Mean-reversion: pull this layer's lr toward the global average lr
|
|
# (geometric, in log space) -- lr *= (avg/lr)**lr_pull. This is the
|
|
# restoring force that replaces the hard min/max rails: layers can
|
|
# still settle at their own level, but can't drift apart to opposite
|
|
# extremes (one frozen, one runaway) and wreck the model. The pull is
|
|
# toward an emergent average, not a fixed target, so it stays
|
|
# automatic. self._avg_lr is refreshed once per .step().
|
|
pull = group["lr_pull"]
|
|
if pull > 0.0 and self._avg_lr > 0.0:
|
|
lr_t.mul_(lr_t.reciprocal().mul_(self._avg_lr).pow_(pull))
|
|
|
|
state["prev_sign"] = cur_sign
|
|
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 * update (lr is a scalar, broadcasts).
|
|
if wd != 0.0:
|
|
update.add_(p, alpha=wd)
|
|
p.addcmul_(update, lr_t, 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_t, value=-1.0)
|
|
self._stochastic_copy_(p, new_p_fp32)
|
|
|
|
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()
|
|
# Fused mode already updated every param in the backward hook; nothing
|
|
# left to do. Non-fused mode does the real work here.
|
|
if not self.fused:
|
|
for group in self.param_groups:
|
|
for p in group["params"]:
|
|
if not p.requires_grad:
|
|
continue
|
|
# Low-precision grads were stochastically accumulated into
|
|
# _accum_grad; hand it back as the grad to update from.
|
|
accum = getattr(p, "_accum_grad", None)
|
|
if accum is not None:
|
|
p.grad = accum
|
|
del p._accum_grad
|
|
if p.grad is None:
|
|
continue
|
|
self._update_param(p, group)
|
|
self._refresh_avg_lr()
|
|
return loss
|
|
|
|
def _all_lrs(self) -> list:
|
|
# The per-layer lr scalars (0-d tensors), gathered without any device
|
|
# sync. Callers stack these and reduce in one op so there is a single
|
|
# GPU->CPU sync instead of one per layer.
|
|
return [
|
|
st["lr"]
|
|
for group in self.param_groups
|
|
for p in group["params"]
|
|
if (st := self.state.get(p)) is not None and "lr" in st
|
|
]
|
|
|
|
def _refresh_avg_lr(self) -> None:
|
|
# Global geometric-mean lr across all layers, the mean-reversion target.
|
|
# Geometric (mean of log) because the lr is controlled/pulled in log
|
|
# space. One stacked reduction + one sync, refreshed once per step.
|
|
lrs = self._all_lrs()
|
|
if lrs:
|
|
self._avg_lr = float(torch.stack(lrs).log_().mean().exp_())
|
|
|
|
def get_learning_rates(self) -> List[float]:
|
|
# Reporting helper: one representative lr per param group -- the mean of
|
|
# the per-layer (per-tensor) lrs. Stacked reduction -> one sync per group.
|
|
out = []
|
|
for group in self.param_groups:
|
|
lrs = [
|
|
self.state[p]["lr"]
|
|
for p in group["params"]
|
|
if p in self.state and "lr" in self.state[p]
|
|
]
|
|
out.append(float(torch.stack(lrs).mean()) 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)
|
|
# prev_sign / dir_ema are transient; rebuild them after load
|
|
# rather than persisting a sign tensor and an fp32 EMA.
|
|
if st is not None and "prev_sign" in st:
|
|
st["prev_sign"] = None
|
|
if st is not None and isinstance(st.get("dir_ema"), torch.Tensor):
|
|
st["dir_ema"] = torch.zeros_like(st["dir_ema"], dtype=torch.float32)
|
|
# Rebuild the global average lr from the restored per-layer lrs.
|
|
self._refresh_avg_lr()
|