diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_bp_train.py | 10 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 4 |
2 files changed, 12 insertions, 2 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 7647fae..069d858 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -21,6 +21,7 @@ ap.add_argument('--cosine', action='store_true') # warmup then cosine de 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 args = ap.parse_args() torch.manual_seed(args.seed) dev = 'cuda' if torch.cuda.is_available() else 'cpu' @@ -76,6 +77,12 @@ blocks = nn.ModuleList([Block(args.C, args.H, args.qk_norm) for _ in range(args. mask = torch.triu(torch.full((args.T, args.T), float('-inf'), device=dev), 1) ln_f = nn.LayerNorm(args.C).to(dev) if args.final_ln else nn.Identity() params = list(tok.parameters()) + list(pos.parameters()) + list(blocks.parameters()) + list(ln_f.parameters()) +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, params, args.lr, args.muon_lr, args.warmup) @@ -116,7 +123,8 @@ n = sum(p.numel() for p in params) print(f'[{args.tag}] cascade-BP L{args.L} C{args.C} H{args.H} T{args.T} | {n/1e6:.2f}M params | {dev}', flush=True) best, t0 = 1e9, time.time() outdir = Path('runs'); outdir.mkdir(exist_ok=True) -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') loss = F.cross_entropy(fwd(x).reshape(-1, vocab), y.reshape(-1)) opt.zero_grad(set_to_none=True); loss.backward() diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 74d6a3b..b85ac96 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -211,6 +211,7 @@ def ep_step(x, y): 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 and not args.noguard): + GOV['skd'] = GOV.get('skd', 0) + 1 # drift-guard reject (relaxation non-convergence) for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False GOV['drift'] = drift @@ -226,6 +227,7 @@ def ep_step(x, y): 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): + GOV['skg'] = GOV.get('skg', 0) + 1 # gn-EMA-guard reject (gradient-magnitude spike) for p in all_params: p.grad = None return free_ce, beta_t, GOV['K'], False for p, g in zip(all_params, gs): @@ -287,7 +289,7 @@ for step in range(start_step, 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} ' + f'| beta={beta_t:.2e} K={rounds} skips={skips}(d{GOV.get("skd",0)}/g{GOV.get("skg",0)}){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, |