332 lines
15 KiB
Python
332 lines
15 KiB
Python
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)
|