Add Automagic V4

toolkit/optimizer.py

@@ -103,6 +103,9 @@ def get_optimizer(
     elif lower_type == 'automagic3':
         from toolkit.optimizers.automagic3 import Automagic3
         optimizer = Automagic3(params, lr=float(learning_rate), **optimizer_params)
+    elif lower_type == 'automagic4':
+        from toolkit.optimizers.automagic4 import Automagic4
+        optimizer = Automagic4(params, lr=float(learning_rate), **optimizer_params)
     elif lower_type == 'automagicexperiment':
         from toolkit.optimizers.automagicEXPERIMENT import AutomagicEXPERIMENT
         optimizer = AutomagicEXPERIMENT(params, lr=float(learning_rate), **optimizer_params)

toolkit/optimizers/automagic4.py

@@ -0,0 +1,259 @@
+"""
+NOTE: This is experimental and under active development; expect breaking changes and bugs. Feedback welcome.
+"""
+
+import math
+import torch
+
+from toolkit.optimizers.automagic3 import Automagic3
+
+
+class Automagic4(Automagic3):
+    """
+    Automagic v4: Automagic3 with sensitivity-balanced updates.
+
+    v3 gives every element an Adam-normalized step of about lr, so how far one
+    step moves the OUTPUT scales with that weight's sensitivity. A single shared
+    lr cannot fit a model whose tensors differ in sensitivity by orders of
+    magnitude (a head reading a huge raw residual, a modulation shared by every
+    layer): at the lr the bulk needs, the sensitive ones overshoot.
+
+    v4 multiplies each element's (clipped) v3 update by
+
+        balance = min(1, reference / sensitivity)
+
+    sensitivity = sqrt(bias-corrected second moment). For a noise-dominated MSE
+    loss (diffusion) gradient power tracks the Gauss-Newton curvature, i.e. how
+    far the output moves per unit change of the weight. It comes from the same
+    row/col (or 1D) state v3 keeps, so it is free and per element: per column
+    for a head whose input dims differ in scale, per row for a modulation.
+    Above the reference the step becomes Newton-like (grad / v), which is
+    invariant to rescaling the weights -- including coupled Q/K-style
+    rescalings, so it cannot seed the fights the shared group lr prevents.
+
+    One-sided on purpose: low gradient power cannot be told apart from a weight
+    that is gated off (a LoRA down while its zero-initialized up is still ~0,
+    dead units), where Newton would take an unbounded step. Those keep v3's
+    normalized step, which is what gets zero-initialized LoRA pairs moving.
+
+    reference: the ``sensitivity_quantile`` (default 0.9) quantile of the
+    per-tensor geometric-mean sensitivity, weighted by element count, from the
+    previous step. High so a large near-dead mass (LoRA downs at init are about
+    half of a LoRA's elements) cannot drag it down and stall everything above
+    it; the outliers it targets are a few percent of elements.
+
+    The first step only measures (second moments and the reference): without a
+    reference the sensitive tensors would take a full v3 step.
+
+    lr controller: re-centered flip-rate vote. v3's neutral point is a 50% sign
+    flip rate, but a weight bouncing on gradient noise at its minimum has
+    lag-1 gradient correlation -effective_lr * curvature / 2 at any lr > 0, so
+    v3 only balances at lr = 0 and anneals once the signal runs out. v4 votes
+    (neutral - observed) flip rate per element, with
+
+        neutral = 0.5 + asin(step_target / 2) / pi
+
+    the flip rate of a stationary weight whose effective_lr * curvature equals
+    ``step_target`` (default 0.1; 2 is the stability edge). The lr then has a
+    fixed point set by step_target, with or without signal: trends flip less
+    and raise it, overshoot flips more and lowers it. Votes are weighted by the
+    balanced |update| (output impact). No gain: near the fixed point the flip
+    rate moves ~0.16 * step_target per e-fold of lr, so the loop gain over the
+    H-step window lag stays well below 1.
+    """
+
+    def __init__(
+        self,
+        params,
+        *args,
+        sensitivity_quantile: float = 0.9,
+        step_target: float = 0.1,
+        **kwargs,
+    ):
+        if not 0.0 < sensitivity_quantile <= 1.0:
+            raise ValueError(
+                f"sensitivity_quantile ({sensitivity_quantile}) must be in (0, 1]"
+            )
+        if not 0.0 < step_target < 2.0:
+            raise ValueError(f"step_target ({step_target}) must be in (0, 2)")
+        self.sensitivity_quantile = float(sensitivity_quantile)
+        self.step_target = float(step_target)
+        self._neutral_flip = 0.5 + math.asin(self.step_target / 2.0) / math.pi
+        super().__init__(params, *args, **kwargs)
+
+    # ------------------------------------------------------------------ utils
+
+    @staticmethod
+    def _ema_(
+        state_t: torch.Tensor, sample: torch.Tensor, beta2: float
+    ) -> torch.Tensor:
+        # fp32 beta2-EMA; low-precision state is written back
+        v = state_t if state_t.dtype == torch.float32 else state_t.to(torch.float32)
+        v.mul_(beta2).add_(sample, alpha=1.0 - beta2)
+        if v is not state_t:
+            state_t.copy_(v)
+        return v
+
+    def _rebuild_group_index(self) -> None:
+        super()._rebuild_group_index()
+        n = len(self.param_groups)
+        self._group_ref: list = [None] * n
+        self._group_sens_log: list = [[] for _ in range(n)]
+        self._group_sens_numel: list = [[] for _ in range(n)]
+
+    def _init_state(self, p: torch.Tensor, group: dict) -> None:
+        super()._init_state(p, group)
+        gi = self._param_group_index.get(p)
+        ref = self._group_ref[gi] if gi is not None else None
+        if ref is not None:
+            self.state[p]["sens_ref"] = ref.to(p.device, copy=True)
+
+    # -------------------------------------------------------------- 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("Automagic4 does not support sparse gradients.")
+        if grad.dtype != torch.float32:
+            grad = grad.to(torch.float32)
+        grad.nan_to_num_(nan=0.0, posinf=0.0, neginf=0.0)
+
+        beta2 = group["beta2"]
+        eps = group["eps"]
+        sq = grad * grad
+        if p.dim() >= 2:
+            row = self._ema_(state["exp_avg_sq_row"], sq.mean(dim=-1).add_(eps), beta2)
+            col = self._ema_(state["exp_avg_sq_col"], sq.mean(dim=-2).add_(eps), beta2)
+            del sq
+            # mean log of the factored v = row (x) col / mean(row)
+            log_v = (
+                row.log().mean(-1) - row.mean(-1).log() + col.log().mean(-1)
+            ).mean()
+            inv_sqrt_v = self._approx_sq_grad(row, col)
+        else:
+            v = self._ema_(state["exp_avg_sq"], sq, beta2).add(eps)
+            del sq
+            log_v = v.log().mean()
+            inv_sqrt_v = v.rsqrt_()
+        state["step"] += 1
+        bias_corr = 1.0 - beta2 ** state["step"]
+
+        gi = self._param_group_index.get(p)
+        if gi is not None:
+            self._group_sens_log[gi].append(0.5 * (log_v - math.log(bias_corr)))
+            self._group_sens_numel[gi].append(p.numel())
+
+        ref = state.get("sens_ref")
+        if ref is None:
+            # measurement-only first step (no reference yet)
+            p.grad = None
+            return
+
+        update = grad.mul_(inv_sqrt_v)
+        update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+        update.clamp_(-group["clip_threshold"], group["clip_threshold"])
+        # balance = min(1, ref / sensitivity), sensitivity = sqrt(v / bias_corr)
+        balance = inv_sqrt_v.mul_(ref * math.sqrt(bias_corr)).clamp_(max=1.0)
+        update.mul_(balance)
+        del balance, inv_sqrt_v
+
+        # flip-rate vote (see class docstring), pooled per group, applied in .step()
+        cur_bits = update.gt(0.0)
+        hist = state["sign_history"]
+        idx = state["hist_idx"]
+        H = hist.shape[0]
+        lr_t = state["lr"]
+        hist[idx].copy_(self._pack_bits(cur_bits))
+        state["hist_idx"] = (idx + 1) % H
+        fill = min(H, state["hist_fill"] + 1)
+        state["hist_fill"] = fill
+
+        if fill == H and gi is not None:
+            _, shifts = self._pack_consts(hist.device)
+            chron = torch.roll(hist, -state["hist_idx"], dims=0)
+            bits = (
+                (chron.unsqueeze(-1) >> shifts)
+                .bitwise_and_(1)
+                .view(H, -1)[:, : update.numel()]
+            )
+            flips = (bits[1:] ^ bits[:-1]).sum(0, dtype=torch.int16)
+            w = update.abs().view(-1)
+            den = w.sum()
+            num = den * self._neutral_flip - (w * flips).sum() / (H - 1)
+            if self._group_num[gi] is None:
+                self._group_num[gi] = num
+                self._group_den[gi] = den
+            else:
+                acc = self._group_num[gi]
+                if num.device != acc.device:
+                    num = num.to(acc.device)
+                    den = den.to(acc.device)
+                acc.add_(num)
+                self._group_den[gi].add_(den)
+
+        wd = group["weight_decay"]
+        if p.dtype == torch.float32:
+            if wd != 0.0:
+                update.add_(p, alpha=wd)
+            p.addcmul_(update, lr_t, value=-1.0)
+        else:
+            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
+
+    def _apply_group_votes(self) -> None:
+        self._update_sensitivity_refs()
+        super()._apply_group_votes()
+
+    def _update_sensitivity_refs(self) -> None:
+        # numel-weighted quantile of this step's per-tensor log sensitivities; no GPU sync
+        for gi, group in enumerate(self.param_groups):
+            logs = self._group_sens_log[gi]
+            if not logs:
+                continue
+            dev = logs[0].device
+            vals = torch.stack([t.to(dev) for t in logs])
+            weights = torch.tensor(
+                self._group_sens_numel[gi], dtype=torch.float64, device=dev
+            )
+            order = vals.argsort()
+            cum = weights[order].cumsum(0)
+            pos = torch.searchsorted(cum, cum[-1:] * self.sensitivity_quantile)
+            ref = vals[order][pos.clamp_(max=len(logs) - 1)].squeeze(0).exp_()
+            self._group_ref[gi] = ref
+            for p in group["params"]:
+                st = self.state.get(p)
+                if not st:
+                    continue
+                cur = st.get("sens_ref")
+                if cur is None:
+                    st["sens_ref"] = ref.to(p.device, copy=True)
+                else:
+                    cur.copy_(ref)
+            self._group_sens_log[gi] = []
+            self._group_sens_numel[gi] = []
+
+    def load_state_dict(self, state_dict):
+        super().load_state_dict(state_dict)
+        # parent cast the reference to param dtype; restore fp32 and the group mirror
+        for gi, group in enumerate(self.param_groups):
+            for p in group["params"]:
+                st = self.state.get(p)
+                ref = st.get("sens_ref") if st is not None else None
+                if isinstance(ref, torch.Tensor):
+                    st["sens_ref"] = ref.to(torch.float32)
+                    if self._group_ref[gi] is None:
+                        self._group_ref[gi] = st["sens_ref"].clone()