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/casc_eq_train.py | 9 ++++- ep_run/muon.py | 103 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+), 2 deletions(-) diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 05c1e4e..7ded9d9 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -22,7 +22,7 @@ ap.add_argument('--wandb_run', default='') ap.add_argument('--kmax', type=int, default=8) # adaptive fb rounds cap ap.add_argument('--noguard', action='store_true') # diagnosis: skip only non-finite grads ap.add_argument('--untie', action='store_true') # separate readout matrix (untied from tok) -ap.add_argument('--opt', choices=['adamw', 'muon'], default='adamw') +ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor'], default='adamw') ap.add_argument('--muon_lr', type=float, default=0.02) ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02) ap.add_argument('--compile', action='store_true') # torch.compile each block (free speed where supported) @@ -341,7 +341,12 @@ if args.bf16: if args.untie: with torch.no_grad(): W_out.data = W_out.data.to(torch.bfloat16) print('[bf16] model cast to bfloat16 (E-accum + sigma stay fp32)', flush=True) -if args.opt == 'muon': +if args.opt in ('sgdm', 'lion', 'olion', 'adafactor'): + from muon import build_alt + opt, sched = build_alt(args.opt, blocks, all_params, args.lr, args.warmup, + total_steps=(args.steps if args.cosine else 0), lr_min_ratio=args.lr_min_ratio, + lr_matrix=(args.muon_lr if args.muon_lr != 0.02 else None)) +elif args.opt == 'muon': from muon import build_hybrid opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup, muon_mom=args.muon_mom, adam_b1=args.adam_b1, 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