1 files changed, 103 insertions, 0 deletions
diff --git a/ep_run/muon.py b/ep_run/muon.py
index 0806243..1466ee6 100644
--- a/ep_run/muon.py
+++ b/ep_run/muon.py
@@ -82,3 +82,106 @@ def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup, total_steps=0,
         fn = lambda s: min(1.0, (s + 1) / max(warmup, 1))
     scheds = [torch.optim.lr_scheduler.LambdaLR(om, fn), torch.optim.lr_scheduler.LambdaLR(oa, fn)]
     return MultiOpt([om, oa]), MultiSched(scheds)
+
+
+class Lion(torch.optim.Optimizer):
+    """Lion (Chen et al. 2023): sign of the beta1-mixed momentum; decoupled wd."""
+    def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), wd=0.0):
+        super().__init__(params, dict(lr=lr, betas=betas, wd=wd))
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for g_ in self.param_groups:
+            b1, b2 = g_['betas']
+            for p in g_['params']:
+                if p.grad is None: continue
+                st = self.state[p]
+                if 'm' not in st: st['m'] = torch.zeros_like(p)
+                m = st['m']
+                u = (b1 * m + (1 - b1) * p.grad).sign_()
+                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
+                p.add_(u, alpha=-g_['lr'])
+                m.mul_(b2).add_(p.grad, alpha=1 - b2)
+
+
+class OLion(torch.optim.Optimizer):
+    """OLion (arXiv:2602.01105): Lion-style momentum -> Newton-Schulz -> entrywise sign,
+    RMS alignment gamma (||sign||_F = sqrt(numel) exactly, so D = gamma * sign(Q)), decoupled wd.
+    2D params only; caller routes 1D elsewhere. betas default to Lion's (0.9, 0.99) — the paper's
+    defaults were not in the pages we read; flagged as an assumption in the battery notes."""
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.99), gamma=0.2, ns_steps=5, wd=0.0):
+        super().__init__(params, dict(lr=lr, betas=betas, gamma=gamma, ns_steps=ns_steps, wd=wd))
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for g_ in self.param_groups:
+            b1, b2 = g_['betas']
+            for p in g_['params']:
+                if p.grad is None: continue
+                st = self.state[p]
+                if 'm' not in st: st['m'] = torch.zeros_like(p)
+                m = st['m']
+                m.mul_(b2).add_(p.grad, alpha=1 - b2)          # slow momentum
+                gt = (1 - b1) * p.grad + b1 * m                # Nesterov mix
+                q = newton_schulz(gt, g_['ns_steps'])
+                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
+                p.add_(q.sign_(), alpha=-g_['lr'] * g_['gamma'])
+
+
+class Adafactor2D(torch.optim.Optimizer):
+    """Minimal Shazeer-Stern factored Adam for 2D params: row/col second-moment statistics,
+    RMS-1 update clipping, no relative-step magic (external lr + schedule). The analog-native
+    Adam per the 07-11 BoM audit (row/col stats = AGC channels)."""
+    def __init__(self, params, lr=1e-3, beta2=0.999, eps=1e-30, clip=1.0, wd=0.0):
+        super().__init__(params, dict(lr=lr, beta2=beta2, eps=eps, clip=clip, wd=wd))
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for g_ in self.param_groups:
+            for p in g_['params']:
+                if p.grad is None: continue
+                g = p.grad
+                st = self.state[p]
+                if 'r' not in st:
+                    st['r'] = torch.zeros(g.shape[0], device=g.device)
+                    st['c'] = torch.zeros(g.shape[1], device=g.device)
+                r, c = st['r'], st['c']
+                g2 = g.float().pow(2) + g_['eps']
+                r.mul_(g_['beta2']).add_(g2.mean(1), alpha=1 - g_['beta2'])
+                c.mul_(g_['beta2']).add_(g2.mean(0), alpha=1 - g_['beta2'])
+                v = r[:, None] * c[None, :] / max(float(r.mean()), g_['eps'])
+                u = g / v.sqrt().to(g.dtype)
+                rms = float(u.pow(2).mean().sqrt())
+                if rms > g_['clip']: u = u * (g_['clip'] / rms)
+                if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd'])
+                p.add_(u, alpha=-g_['lr'])
+
+
+def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ratio=0.1,
+              lr_matrix=None, wd=0.0):
+    """Screening-tier builder for the optimizer price list: OPT on block matrices + AdamW on the
+    rest (same split as build_hybrid so arms differ only in the matrix rule). sgdm applies SGD to
+    everything (no split) — the fully-local baseline."""
+    import math as _m
+    mats = [p for p in blocks.parameters() if p.ndim == 2]
+    mat_ids = {id(p) for p in mats}
+    rest = [p for p in other_params if id(p) not in mat_ids]
+    lm = lr_matrix if lr_matrix is not None else lr
+    if opt_name == 'sgdm':
+        om = torch.optim.SGD(mats + rest, lr=lm, momentum=0.95, nesterov=True)
+        opts = [om]
+    else:
+        om = {'lion': lambda: Lion(mats, lr=lm, wd=wd),
+              'olion': lambda: OLion(mats, lr=lm, wd=wd),
+              'adafactor': lambda: Adafactor2D(mats, lr=lm, wd=wd)}[opt_name]()
+        oa = torch.optim.AdamW(rest, lr=lr, weight_decay=1e-4)
+        opts = [om, oa]
+    if total_steps > 0:
+        def fn(s):
+            if s < warmup: return (s + 1) / max(warmup, 1)
+            pr = min(1.0, (s - warmup) / max(1, total_steps - warmup))
+            return lr_min_ratio + 0.5 * (1 - lr_min_ratio) * (1 + _m.cos(_m.pi * pr))
+    else:
+        fn = lambda s: min(1.0, (s + 1) / max(warmup, 1))
+    scheds = [torch.optim.lr_scheduler.LambdaLR(o, fn) for o in opts]
+    return MultiOpt(opts), MultiSched(scheds)