diff options
| author | Yuren Hao <yurenh2@illinois.edu> | 2026-07-17 00:44:11 -0500 |
|---|---|---|
| committer | Yuren Hao <yurenh2@illinois.edu> | 2026-07-17 00:44:11 -0500 |
| commit | 835d327caa0a86d3b9fdc7b1008d717c5d4c5bed (patch) | |
| tree | 5efbf393a28d646d23a090c8b5ca97993b41d829 /ep_run | |
| parent | 2c32641b59de1ae725d595ce9e3f865d45fabc43 (diff) | |
save final ckpt at last step (fixed-eval protocol); fw72m_ride crown-3 launched (plain+ride governor); ride/BP seed campaign on GPU0
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_bp_train.py | 2 | ||||
| -rw-r--r-- | ep_run/casc_eq_train.py | 2 |
2 files changed, 2 insertions, 2 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 7bafc16..0515ec1 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -238,7 +238,7 @@ for step in range(start_step, args.steps + 1): if wb is not None: try: wb.log({'train_ce': loss.item(), 'val_ce': val, 'best': best}, step=step) except Exception: pass - if step % args.save_every == 0: + if step % args.save_every == 0 or step == args.steps: torch.save({'tok': tok.state_dict(), 'pos': pos.state_dict(), 'blocks': blocks.state_dict(), 'wout': (W_out.detach().cpu() if args.olmo2 else None), 'lnf': (ln_f.state_dict() if not isinstance(ln_f, nn.Identity) else None), diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 6d94ef5..fef44f0 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -708,7 +708,7 @@ for step in range(start_step, args.steps + 1): 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) except Exception: pass - if step % args.save_every == 0 and step > 0 and RANK == 0: + if (step % args.save_every == 0 or step == args.steps) and step > 0 and RANK == 0: 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), |