diff options
Diffstat (limited to 'ep_run/casc_eq_train.py')
| -rw-r--r-- | ep_run/casc_eq_train.py | 8 |
1 files changed, 8 insertions, 0 deletions
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) |