diff options
| -rw-r--r-- | docs/hardware/EQOLION_DESIGN.md | 10 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 | ||||
| -rw-r--r-- | ep_run/muon.py | 85 |
3 files changed, 96 insertions, 3 deletions
diff --git a/docs/hardware/EQOLION_DESIGN.md b/docs/hardware/EQOLION_DESIGN.md index 7ba1577..a44fccb 100644 --- a/docs/hardware/EQOLION_DESIGN.md +++ b/docs/hardware/EQOLION_DESIGN.md @@ -78,3 +78,13 @@ sign(M) vs sign(polar(M)) 一致率**中位 0.700（0.647-0.816）**，谱偏斜 (b) Occam 臂升为主候选——行/列 AGC 可能拿走谱均衡里"有用的部分"而不碰病态尾部； (c) 限定：本轨迹取自训练最早期（步 100-600, LR 峰值区, 漂移最快的最难工况）——**晚期 ckpt 重 dump 后再判死刑**。 + +--- +## 第二代设计（2026-08-05，R4 后重排：只追翻腾慢或不翻腾的目标） +| 名字 | 规则 | 原语构成 | 针对的洞 | +|---|---|---|---| +| **SignLine** | a_i·b_j·sign(m_ij)，a/b 来自行/列 RMS 的漏积分（^-α/4，均值归一） | 漏积分×(d1+d2) + 比较器 + **每线驱动幅度**（BoM 已有的 per-line DAC） | 关键洞察：正的行/列缩放放进 sign 里是无操作，只有作为**脉冲幅度**放在 sign 外面才起作用 = 因子化自适应的模拟原生形态 | +| **ConsensusLion** | 快慢两个动量的 sign 一致才写，不一致处不写 | 两个不同漏率的电容 + 比较器 + 符合门 | 用免费的多时间尺度做噪声否决；整个优化器=电容加逻辑 | +| **DitherLion** | sign(m + τ·noise·rms(m)) | 基底噪声（免费）+ 比较器 | 硬 sign 变期望意义的软 sign，小分量跨步携带比例信息 | +| KFAC-lite（下一轮） | u = m·Q，Q 用反 Hebbian 局部规则追踪**输入活动协方差**的白化 | 搭车探针 + 外积写 + 漏衰减 | **数据统计是平稳的**——R4 杀掉的是追极因子（高翻腾），活动协方差不翻腾，追踪前提在这里成立 | +测试协议 = 与 stage A 完全同网格（C128/2500 步/3 档 LR），本地双卡 ~45 分钟出结果，与农场标准组同表对比。 diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 4a436d1..6a6865d 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', 'sgdm', 'lion', 'olion', 'adafactor'], default='adamw') +ap.add_argument('--opt', choices=['adamw', 'muon', 'sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion'], 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) @@ -344,7 +344,7 @@ 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 in ('sgdm', 'lion', 'olion', 'adafactor'): +if args.opt in ('sgdm', 'lion', 'olion', 'adafactor', 'signline', 'conslion', 'ditherlion'): 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, diff --git a/ep_run/muon.py b/ep_run/muon.py index 1466ee6..e6c2692 100644 --- a/ep_run/muon.py +++ b/ep_run/muon.py @@ -173,7 +173,10 @@ def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ 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]() + 'adafactor': lambda: Adafactor2D(mats, lr=lm, wd=wd), + 'signline': lambda: SignLine(mats, lr=lm, wd=wd), + 'conslion': lambda: ConsensusLion(mats, lr=lm, wd=wd), + 'ditherlion': lambda: DitherLion(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: @@ -185,3 +188,83 @@ def build_alt(opt_name, blocks, other_params, lr, warmup, total_steps=0, lr_min_ 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) + + +class SignLine(torch.optim.Optimizer): + """sign(momentum) with per-row/per-column pulse amplitudes from leaky RMS line statistics. + Insight: positive row/col scaling INSIDE a sign is a no-op; applied OUTSIDE as amplitudes + (a_i b_j sign(m_ij)) it is the analog-native form of factored adaptivity: per-line DACs set + drive amplitude, comparators give the sign. Amplitudes normalized to unit mean so lr keeps + its scale; alpha in [0,1] interpolates flat-Lion (0) -> full line-adaptive (1).""" + def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), beta_line=0.999, alpha=1.0, wd=0.0): + super().__init__(params, dict(lr=lr, betas=betas, bl=beta_line, alpha=alpha, 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) + st['r'] = torch.ones(p.shape[0], device=p.device) + st['c'] = torch.ones(p.shape[1], device=p.device) + m, r, c = st['m'], st['r'], st['c'] + u = b1 * m + (1 - b1) * p.grad + r.mul_(g_['bl']).add_(u.float().pow(2).mean(1), alpha=1 - g_['bl']) + c.mul_(g_['bl']).add_(u.float().pow(2).mean(0), alpha=1 - g_['bl']) + a = (r / r.mean()).clamp_min(1e-12).pow(-0.25 * g_['alpha']) + b = (c / c.mean()).clamp_min(1e-12).pow(-0.25 * g_['alpha']) + amp = torch.outer(a, b).to(p.dtype) + amp = amp / amp.mean() + if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd']) + p.add_(u.sign_() * amp, alpha=-g_['lr']) + m.mul_(b2).add_(p.grad, alpha=1 - b2) + + +class ConsensusLion(torch.optim.Optimizer): + """two-timescale sign consensus: update only where sign(fast momentum) == sign(slow momentum). + Two leaky integrators with different leaks + one comparator + coincidence gate — the entire + optimizer is capacitors and logic. Where they disagree, write nothing (noise veto).""" + def __init__(self, params, lr=3e-4, beta_fast=0.9, beta_slow=0.99, wd=0.0): + super().__init__(params, dict(lr=lr, bf=beta_fast, bs=beta_slow, 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 + st = self.state[p] + if 'mf' not in st: + st['mf'] = torch.zeros_like(p); st['ms'] = torch.zeros_like(p) + mf, ms = st['mf'], st['ms'] + mf.mul_(g_['bf']).add_(p.grad, alpha=1 - g_['bf']) + ms.mul_(g_['bs']).add_(p.grad, alpha=1 - g_['bs']) + sf, ss = mf.sign(), ms.sign() + u = sf * (sf == ss) + if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd']) + p.add_(u, alpha=-g_['lr']) + + +class DitherLion(torch.optim.Optimizer): + """Lion with dithered sign: sign(m + tau*noise*rms(m)). Free substrate noise turns the hard + sign into an unbiased soft-sign in expectation, letting small entries carry proportional + information across steps. tau=0 recovers Lion.""" + def __init__(self, params, lr=3e-4, betas=(0.9, 0.99), tau=0.5, wd=0.0): + super().__init__(params, dict(lr=lr, betas=betas, tau=tau, 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 + d = torch.randn_like(u) * (g_['tau'] * u.float().pow(2).mean().sqrt().to(u.dtype)) + if g_['wd'] > 0: p.mul_(1 - g_['lr'] * g_['wd']) + p.add_((u + d).sign_(), alpha=-g_['lr']) + m.mul_(b2).add_(p.grad, alpha=1 - b2) |