diff --git a/toolkit/optimizers/automagic3.py b/toolkit/optimizers/automagic3.py index 70810b2..93de889 100644 --- a/toolkit/optimizers/automagic3.py +++ b/toolkit/optimizers/automagic3.py @@ -19,10 +19,17 @@ class Automagic3(torch.optim.Optimizer): 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 + 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. + 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) @@ -81,6 +88,7 @@ class Automagic3(torch.optim.Optimizer): eps: float = 1e-30, clip_threshold: float = 1.0, weight_decay: float = 0.0, + fused: bool = True, ): if lr > 1e-3: print(f"Warning! Start lr {lr} is very high; forcing to 1e-6.") @@ -97,14 +105,27 @@ class Automagic3(torch.optim.Optimizer): ) super().__init__(params, defaults) + self.fused = fused self._hook_handles = [] for group in self.param_groups: for p in group["params"]: - if p.requires_grad: + 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:,}") @@ -151,6 +172,35 @@ class Automagic3(torch.optim.Optimizer): 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 @@ -193,6 +243,17 @@ class Automagic3(torch.optim.Optimizer): 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 @@ -231,6 +292,11 @@ class Automagic3(torch.optim.Optimizer): # 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"]) # Per-row sign agreement vs. the previous step: the plain fraction of # elements that kept their sign, every element counting equally @@ -283,12 +349,7 @@ class Automagic3(torch.optim.Optimizer): 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)) + self._stochastic_copy_(p, new_p_fp32) p.grad = None @@ -300,6 +361,22 @@ class Automagic3(torch.optim.Optimizer): 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) return loss def get_learning_rates(self) -> List[float]: diff --git a/toolkit/optimizers/test_optimizers.py b/toolkit/optimizers/test_optimizers.py index 8a36f37..9c48cb3 100644 --- a/toolkit/optimizers/test_optimizers.py +++ b/toolkit/optimizers/test_optimizers.py @@ -8,7 +8,7 @@ Compares speed (ms/step) and peak VRAM across: - Adafactor - Automagic v1 - Automagic v2 (fused-backward) - - Automagic v3 (fused-backward) + - Automagic v3 (fused-backward and traditional/unfused) - Prodigy """ import contextlib @@ -186,7 +186,8 @@ def main(): ("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)), + ("Automagic v3 fused", lambda p: Automagic3(p, lr=1e-4, fused=True)), + ("Automagic v3 unfused", lambda p: Automagic3(p, lr=1e-4, fused=False)), ("Prodigy", lambda p: Prodigy(p, lr=1.0, eps=1e-6)), ]