From 907588538ead4d005b6cf4c0fa26ac2450fa5d7e Mon Sep 17 00:00:00 2001 From: Yuren Hao Date: Mon, 13 Jul 2026 12:01:46 -0500 Subject: fw72m crown run launched (NCCL DDP prod), exact resume (opt state), RESULT 17, shuffle+prep scripts Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn --- ep_run/casc_eq_train.py | 8 ++++++++ 1 file changed, 8 insertions(+) (limited to 'ep_run/casc_eq_train.py') diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 05273df..573f030 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -412,6 +412,13 @@ def evaluate(nb=6): tot += F.cross_entropy(readout(z).reshape(-1, vocab), y.reshape(-1)).item() return tot / nb +if args.resume and _ck.get('opt') is not None: + try: + opt.load_state_dict(_ck['opt']) + print('[resume] optimizer state restored (exact chunked-resume)', flush=True) + except Exception as e: + print(f'[resume] optimizer state NOT restored ({e}) — cold optimizer', flush=True) + if DDP: # belt & suspenders on top of identical init seeds: rank0's params are law with torch.no_grad(): for p in all_params: @@ -509,6 +516,7 @@ for step in range(start_step, args.steps + 1): torch.save({'tok': tok.state_dict(), 'pos': pos.state_dict(), 'blocks': blocks.state_dict(), 'wout': (W_out.detach().cpu() if args.untie else None), 'lnf': (ln_f.state_dict() if not isinstance(ln_f, nn.Identity) else None), + 'opt': opt.state_dict(), # full optimizer state -> exact resume for chunked HPC jobs 'step': step, 'val': best, 'config': vars(args)}, Path('runs') / f'{args.tag}_s{step}.pt') if RANK == 0: print(f'[{args.tag}] DONE best val CE {best:.4f}', flush=True) -- cgit v1.2.3