diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_bp_train.py | 4 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 27 |
2 files changed, 23 insertions, 8 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 8e039ec..5194c45 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -14,6 +14,7 @@ ap.add_argument('--lr', type=float, default=3e-4); ap.add_argument('--warmup', t ap.add_argument('--seed', type=int, default=0) ap.add_argument('--save_every', type=int, default=500); ap.add_argument('--log', type=int, default=200) ap.add_argument('--wandb', default=''); ap.add_argument('--wandb_run', default='') +ap.add_argument('--tok_init', type=float, default=0.0) # >0: init tok/pos std (GPT-standard 0.02) args = ap.parse_args() torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -40,6 +41,9 @@ class Block(nn.Module): tok = nn.Embedding(vocab, args.C).to(dev) pos = nn.Embedding(args.T, args.C).to(dev) +if args.tok_init > 0: + with torch.no_grad(): + tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init) blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev) mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1) params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 96b7ab3..547312e 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -19,6 +19,9 @@ ap.add_argument('--geta', type=float, default=1.0) # fb mixing (1.0 = undam ap.add_argument('--save_every', type=int, default=1000); ap.add_argument('--log', type=int, default=100) ap.add_argument('--wandb', default=''); 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('--tok_init', type=float, default=0.0) # >0: init tok/pos with this std (GPT-standard 0.02) ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry args = ap.parse_args() torch.manual_seed(args.seed) @@ -46,10 +49,14 @@ class Block(nn.Module): tok = nn.Embedding(vocab, args.C).to(dev) pos = nn.Embedding(args.T, args.C).to(dev) +if args.tok_init > 0: + with torch.no_grad(): + tok.weight.normal_(0, args.tok_init); pos.weight.normal_(0, args.tok_init) blocks = nn.ModuleList([Block(args.C, args.H) for _ in range(args.L)]).to(dev) mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1) -readout = lambda z: z @ tok.weight.t() -all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) +W_out = nn.Parameter(torch.randn(vocab, args.C, device=dev) * 0.02) if args.untie else None +readout = (lambda z: z @ W_out.t()) if args.untie else (lambda z: z @ tok.weight.t()) +all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + ([W_out] if args.untie else []) opt = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=1e-4) sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(args.warmup, 1))) NBT = args.B * args.T @@ -65,7 +72,7 @@ def free_states(x): @torch.no_grad() def tok_sigma(iters=8): """top singular value of tok.weight (power iteration on the raw matrix).""" - W = tok.weight + W = W_out if args.untie else tok.weight v = torch.randn(W.shape[1], device=dev); v /= v.norm() sig = 1.0 for _ in range(iters): @@ -104,13 +111,14 @@ def dFdtheta(zs, x, y, beta): return [g if g is not None else None for g in gs] SIG0 = None -GOV = {'K': None, 'bscale': 1.0, 'gema': None} +GOV = {'K': None, 'bscale': 1.0, 'gema': None, 'drift': 0.0, 'gn': 0.0, 'sig': 0.0} def ep_step(x, y): """single-sided EP with a QUALITY-GOVERNED estimator: beta_t = beta0*bscale*sig0^2/sig^2, K = GOV['K'] fb rounds; guard = finiteness + drift + grad-norm sanity only.""" global SIG0 if GOV['K'] is None: GOV['K'] = args.K sig = tok_sigma() + GOV['sig'] = sig if SIG0 is None: SIG0 = sig beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9) z0, zs_free = free_states(x) @@ -119,9 +127,10 @@ def ep_step(x, y): with torch.no_grad(): drift = sum(float((a - b).norm()) for a, b in zip(zp, zs_free)) / max( sum(float(b.norm()) for b in zs_free), 1e-9) - if (not math.isfinite(drift)) or drift > 0.5: + if (not math.isfinite(drift)) or (drift > 0.5 and not args.noguard): for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False + GOV['drift'] = drift prev = tok(x) + pos(torch.arange(args.T, device=dev))[None] E = 0.0 for z, b in zip(zp, blocks): E = E + 0.5 * ((z - b(prev, mask)) ** 2).sum(); prev = z @@ -132,10 +141,11 @@ def ep_step(x, y): if g is not None: gn += float((g ** 2).sum()) gn = gn ** 0.5 if GOV['gema'] is None: GOV['gema'] = gn - if not math.isfinite(gn) or gn > 8 * GOV['gema']: + GOV['gema'] = 0.99 * GOV['gema'] + 0.01 * gn # EMA always updates (frozen-ref bugfix) + GOV['gn'] = gn + if not math.isfinite(gn) or (gn > 8 * GOV['gema'] and not args.noguard): for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False - GOV['gema'] = 0.99 * GOV['gema'] + 0.01 * gn for p, g in zip(all_params, gs): p.grad = g return free_ce, beta_t, GOV['K'], True @@ -193,7 +203,8 @@ for step in range(args.steps + 1): val = evaluate(); best = min(best, val) gtag = '' if math.isnan(gcos) else f' cos={gcos:.4f}' print(f'step {step:5d}/{args.steps} | train {ce:.4f} val {val:.4f} (best {best:.4f}) ' - f'| beta={beta_t:.2e} K={rounds} skips={skips}{gtag} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True) + f'| beta={beta_t:.2e} K={rounds} skips={skips}{gtag} ' + f'drift={GOV["drift"]:.3f} gn={GOV["gn"]:.2e} sig={GOV["sig"]:.1f} | {step/max(time.time()-t0,1e-9):.3f} it/s', flush=True) if wb is not None: try: wb.log({'train_ce': ce, 'val_ce': val, 'best': best, 'beta_t': beta_t, 'rounds': rounds, 'skips': skips, 'gate_cos': (None if math.isnan(gcos) else gcos)}, step=step) |