diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 | ||||
| -rw-r--r-- | ep_run/muon.py | 85 |
2 files changed, 86 insertions, 3 deletions
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) |