diff options
| author | Yuren Hao <yurenh2@illinois.edu> | 2026-08-05 06:48:06 -0500 |
|---|---|---|
| committer | Yuren Hao <yurenh2@illinois.edu> | 2026-08-05 06:48:06 -0500 |
| commit | d0d1849e16127d4248bf5ef6c0bb1d6686aa9924 (patch) | |
| tree | c65736712fd122d8520ac6f41af24549ad42df75 | |
| parent | d5158e64cad8a008794b0e61270c3b18821b41b7 (diff) | |
第二代原语优化器: SignLine/ConsensusLion/DitherLion 实装+本地筛选电池发射
R4后设计原则=只追不翻腾的目标。SignLine核心洞察: 正行/列缩放在sign内是无操作, 作为脉冲幅度
在sign外才有效=因子化自适应的模拟原生形态; ConsensusLion=双漏率电容+符合门的噪声否决;
DitherLion=免费噪声做期望软sign。KFAC-lite(追平稳的活动协方差)列为下一轮。
9臂与stage A同协议本地筛选中; 发射纪律教训追加: 发射与等待必须分Bash调用(超时杀树会扫掉setsid)。
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
| -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) |
