From aeca7fe404231abd857b22b83a6813431a11a47e Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Tue, 5 May 2026 09:09:07 -0600
Subject: [PATCH] Add Automagic v2 optimizer. It uses significantly less vram
 and is much more efficient.

---
 toolkit/optimizer.py              |   3 +
 toolkit/optimizers/automagic2.py  | 212 ++++++++++++++++++++++++++++++
 ui/src/app/jobs/new/SimpleJob.tsx |   1 +
 version.py                        |   2 +-
 4 files changed, 217 insertions(+), 1 deletion(-)
 create mode 100644 toolkit/optimizers/automagic2.py

diff --git a/toolkit/optimizer.py b/toolkit/optimizer.py
index 355512e..0357f79 100644
--- a/toolkit/optimizer.py
+++ b/toolkit/optimizer.py
@@ -97,6 +97,9 @@ def get_optimizer(
     elif lower_type == 'automagic':
         from toolkit.optimizers.automagic import Automagic
         optimizer = Automagic(params, lr=float(learning_rate), **optimizer_params)
+    elif lower_type == 'automagic2':
+        from toolkit.optimizers.automagic2 import Automagic2
+        optimizer = Automagic2(params, lr=float(learning_rate), **optimizer_params)
     else:
         raise ValueError(f'Unknown optimizer type {optimizer_type}')
     return optimizer
diff --git a/toolkit/optimizers/automagic2.py b/toolkit/optimizers/automagic2.py
new file mode 100644
index 0000000..903313b
--- /dev/null
+++ b/toolkit/optimizers/automagic2.py
@@ -0,0 +1,212 @@
+from typing import List
+import torch
+
+
+class Automagic2(torch.optim.Optimizer):
+    """
+    Automagic v2.
+
+    A single scalar learning rate is kept per parameter (e.g. one lr for the
+    full weight matrix of a Linear layer rather than one per element). The lr
+    is nudged up when the per-element update direction stays consistent with
+    the previous step and nudged down when it flips, 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). Stochastic rounding is applied only when
+    writing back to a bf16 parameter.
+    """
+
+    def __init__(
+        self,
+        params,
+        lr: float = 1e-6,
+        min_lr: float = 1e-7,
+        max_lr: float = 1e-3,
+        lr_bump: float = 1e-6,
+        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=lr_bump,
+            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)
+
+    def _init_state(self, p: torch.Tensor, group: dict) -> None:
+        state = self.state[p]
+        state["step"] = 0
+        state["lr"] = torch.full(
+            (), 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("Automagic2 does not support sparse gradients.")
+        if grad.dtype != torch.float32:
+            grad = grad.to(torch.float32)
+
+        beta2 = group["beta2"]
+        eps = group["eps"]
+        sq = grad * grad + eps
+
+        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), alpha=1.0 - beta2)
+                col.mul_(beta2).add_(sq.mean(dim=-2), 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), alpha=1.0 - beta2)
+                col.mul_(beta2).add_(sq.mean(dim=-2), 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.rsqrt().mul_(grad)
+
+        update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+
+        # Per-element sign agreement collapsed to a single bump decision.
+        # Kept on-device as a 0-D tensor to avoid a CPU<->GPU sync in the hot path.
+        cur_polarity = update > 0
+        last_polarity = state["last_polarity"]
+        agreement = (cur_polarity == last_polarity).to(torch.float32).mean()
+        state["last_polarity"] = cur_polarity
+
+        lr_t = state["lr"]
+        if state["step"] > 0:
+            direction = (agreement >= 0.5).to(lr_t.dtype) * 2.0 - 1.0
+            lr_t.add_(direction, alpha=group["lr_bump"]).clamp_(
+                min=group["min_lr"], max=group["max_lr"]
+            )
+        state["step"] += 1
+
+        update.mul_(lr_t)
+        if group["weight_decay"] != 0.0:
+            p_fp32 = p if p.dtype == torch.float32 else p.to(torch.float32)
+            update.addcmul_(p_fp32, lr_t, value=group["weight_decay"])
+
+        if p.dtype == torch.bfloat16:
+            # Stochastic rounding fp32 -> bf16: add random noise into the lower
+            # 16 mantissa bits, then truncate. Done in place on new_p_fp32 so
+            # we don't allocate a separate int32 work buffer.
+            new_p_fp32 = p.to(torch.float32).sub_(update)
+            as_int = new_p_fp32.view(torch.int32)
+            as_int.add_(torch.randint_like(as_int, 1 << 16)).bitwise_and_(-65536)
+            p.copy_(new_p_fp32)
+        else:
+            p.add_(update.to(p.dtype), alpha=-1.0)
+
+        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"])
+                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_bump (default 1e-6) isn't rounded away on bf16 weights.
+        super().load_state_dict(state_dict)
+        for group in self.param_groups:
+            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)
diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx
index 4442c81..05ff846 100644
--- a/ui/src/app/jobs/new/SimpleJob.tsx
+++ b/ui/src/app/jobs/new/SimpleJob.tsx
@@ -580,6 +580,7 @@ export default function SimpleJob({
                     { value: 'adamw', label: 'AdamW' },
                     { value: 'adamw8bit', label: 'AdamW8Bit' },
                     { value: 'automagic', label: 'Automagic' },
+                    { value: 'automagic2', label: 'Automagic v2' },
                     { value: 'prodigyopt', label: 'Prodigy' },
                     { value: 'prodigy8bit', label: 'Prodigy8Bit' },
                   ]}
diff --git a/version.py b/version.py
index 7608758..73e5053 100644
--- a/version.py
+++ b/version.py
@@ -1 +1 @@
-VERSION = "0.9.6"
+VERSION = "0.9.7"