From 044973e2a78e0faa1296ac370cefbfc2f6feadd6 Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Wed, 5 Aug 2026 06:11:06 -0500 Subject: =?UTF-8?q?=E4=BC=98=E5=8C=96=E5=99=A8=E4=BB=B7=E7=9B=AE=E8=A1=A8?= =?UTF-8?q?=E5=AE=9E=E8=A3=85:=20Lion/OLion/Adafactor2D/build=5Falt=20+=20?= =?UTF-8?q?--opt=20=E5=85=AD=E9=80=89=E9=A1=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OLion按论文算法框(慢动量+Nesterov混合→NS→逐元素sign→γ=0.2定标+解耦wd; betas取Lion默认 0.9/0.99并注明为假设); Adafactor2D=行/列因子化二阶矩+RMS-1裁剪(07-11审计的"模拟原生Adam"); build_alt与build_hybrid同分割(矩阵走被测规则+其余AdamW), sgdm例外全参数SGD(全局部基线)。 四路CPU冒烟全过(损失下降, gate cos 1.0000)。 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/muon.py | 103 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 103 insertions(+) (limited to 'ep_run/muon.py') 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) -- cgit v1.2.3