diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 19 |
1 files changed, 15 insertions, 4 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index f6ae71b..74d6a3b 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -31,6 +31,9 @@ ap.add_argument('--beta_fixed', action='store_true') # disable sig^2 schedul ap.add_argument('--cosine', action='store_true') # warmup then cosine decay to lr_min_ratio*lr over --steps (long runs) ap.add_argument('--lr_min_ratio', type=float, default=0.1) ap.add_argument('--qk_norm', action='store_true') # RMS-norm q,k per head before scores (OLMo2-style; bounds logits, analog-friendly) +ap.add_argument('--final_ln', action='store_true') # final LayerNorm before readout (standard GPT; bounds sig_tok growth -> keeps beta/estimator healthy on long runs) +ap.add_argument('--resume', default='') # path to a ckpt (tok/pos/blocks) to continue from; step taken from ckpt +ap.add_argument('--sig0', type=float, default=-1.0) # override SIG0 (beta-schedule ref); needed on resume to restore original beta regime ap.add_argument('--dtop_every', type=int, default=1) # 1 = exact (DEFAULT, BP-parity); 2 = fast mode (~20% cheaper, ~4% CE tax at high lr) ap.add_argument('--gate_every', type=int, default=200) # in-training cos(EP,BP) telemetry; <=0 = fully BP-free (no bp_gate at all) ap.add_argument('--gate_govern', action='store_true') # let gate cos adjust K/bscale (default: observe-only => training control is BP-free) @@ -94,8 +97,15 @@ if args.compile: print(f'[compile] disabled ({e})', flush=True) mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1) 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 []) +ln_f = nn.LayerNorm(args.C).to(dev) if args.final_ln else nn.Identity() +readout = (lambda z: ln_f(z) @ W_out.t()) if args.untie else (lambda z: ln_f(z) @ tok.weight.t()) +all_params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + list(ln_f.parameters()) + ([W_out] if args.untie else []) +start_step = 0 +if args.resume: + _ck = torch.load(args.resume, map_location=dev, weights_only=False) + tok.load_state_dict(_ck['tok']); pos.load_state_dict(_ck['pos']); blocks.load_state_dict(_ck['blocks']) + start_step = int(_ck.get('step', 0)) + print(f'[resume] loaded {args.resume} at step {start_step}', flush=True) if args.opt == 'muon': from muon import build_hybrid opt, sched = build_hybrid(blocks, all_params, args.lr, args.muon_lr, args.warmup) @@ -189,7 +199,7 @@ def ep_step(x, y): GOV['sig'] = tok_sigma() GOV['step'] = GOV.get('step', 0) + 1 sig = GOV['sig'] - if SIG0 is None: SIG0 = sig + if SIG0 is None: SIG0 = args.sig0 if args.sig0 > 0 else sig beta_t = args.beta * GOV['bscale'] * (SIG0 * SIG0) / max(sig * sig, 1e-9) if args.beta_fixed: beta_t = args.beta * GOV['bscale'] if args.beta_floor > 0.0: beta_t = max(beta_t, args.beta_floor) @@ -253,7 +263,8 @@ print(f'[{args.tag}] cascade-EP(EQUILIBRIUM/fb) L{args.L} C{args.C} T{args.T} be f'K={args.K} geta={args.geta} | {n/1e6:.2f}M | {dev}', flush=True) best, t0 = 1e9, time.time() skips = 0 -for step in range(args.steps + 1): +for _ in range(start_step): sched.step() # advance LR schedule to the resumed step +for step in range(start_step, args.steps + 1): x, y = get_batch('train') ce, beta_t, rounds, ok = ep_step(x, y) if not ok: skips += 1 |