1 files changed, 64 insertions, 0 deletions
diff --git a/ep_run/muon.py b/ep_run/muon.py
new file mode 100644
index 0000000..5ac6fc6
--- /dev/null
+++ b/ep_run/muon.py
@@ -0,0 +1,64 @@
+"""Muon optimizer (Newton-Schulz orthogonalized momentum) + hybrid helpers.
+Convention: Muon on 2D hidden matrices, AdamW on everything else (emb/pos/LN/bias)."""
+import torch
+
+
+def newton_schulz(G, steps=5, eps=1e-7):
+    """approximate polar factor of G via the quintic NS iteration (Keller Jordan coefficients)."""
+    a, b, c = 3.4445, -4.7750, 2.0315
+    X = G / (G.norm() + eps)
+    transposed = X.size(0) > X.size(1)
+    if transposed: X = X.T
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * (A @ A)
+        X = a * X + B @ X
+    return X.T if transposed else X
+
+
+class Muon(torch.optim.Optimizer):
+    def __init__(self, params, lr=0.02, momentum=0.95, ns_steps=5, nesterov=True):
+        super().__init__(params, dict(lr=lr, momentum=momentum, ns_steps=ns_steps, nesterov=nesterov))
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None: continue
+                g = p.grad
+                st = self.state[p]
+                if 'mom' not in st: st['mom'] = torch.zeros_like(g)
+                buf = st['mom']
+                buf.mul_(group['momentum']).add_(g)
+                u = g.add(buf, alpha=group['momentum']) if group['nesterov'] else buf
+                if u.ndim == 2:
+                    u = newton_schulz(u, group['ns_steps'])
+                    u = u * max(1.0, u.size(0) / u.size(1)) ** 0.5   # rms-matched scaling
+                p.add_(u, alpha=-group['lr'])
+
+
+class MultiOpt:
+    """duck-typed bundle of optimizers (step/zero_grad API-compatible)."""
+    def __init__(self, opts): self.optimizers = opts
+    def step(self):
+        for o in self.optimizers: o.step()
+    def zero_grad(self, set_to_none=True):
+        for o in self.optimizers: o.zero_grad(set_to_none=set_to_none)
+
+
+class MultiSched:
+    def __init__(self, scheds): self.scheds = scheds
+    def step(self):
+        for s in self.scheds: s.step()
+
+
+def build_hybrid(blocks, other_params, lr_adamw, lr_muon, warmup):
+    """Muon(2D block matrices) + AdamW(everything else), with linear-warmup scheds for both."""
+    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]
+    om = Muon(mats, lr=lr_muon)
+    oa = torch.optim.AdamW(rest, lr=lr_adamw, weight_decay=1e-4)
+    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)