Lossless stop/resume for fit(): rolling full-state last.ckpt + resume_from, init_from warm start - #17
Merged
Merged
Conversation
…ferCheckpoint) Restores backbone + head; optimizer state is not carried, so expect a brief transient. Enables continuing a pretrain whose stopping criterion changed (e.g. weighted-val monitoring). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Log the per-event median (and IQR) of the CURTAIN validation loss alongside the mean; the median is far more stable epoch-to-epoch than the outlier-dominated mean. Adds the callback and configs wiring it in. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ume_from enable_checkpointing was hard-off, so a stopped pretrain lost optimizer, scheduler and early-stop state and could only warm-start from weights. save_state=<dir> keeps a rolling full Lightning last.ckpt (save_top_k=0; the exported artifact stays the TransferCheckpoint) and resume_from=<ckpt> resumes it losslessly. TransferCheckpoint gains state_dict/load_state_dict so its export threshold survives the resume.
…ain_end Lightning 2.6 writes last.ckpt in on_validation_end only when a top-k save happened in the same step; with save_top_k=0 that never fires, so the file appeared only at on_train_end -- worthless for crash/timeout recovery. monitor=None + save_top_k=1 makes each validation epoch a save.
save_state becomes an optional override instead of an opt-in: every run now leaves a resumable full-state checkpoint beside its transfer checkpoint, so any crash/timeout/stop can continue via resume_from.
Lightning replays on_train_epoch_end when resuming a validation-end checkpoint, without validation metrics; EarlyStopping's default train-epoch-end check then raises on the missing val_loss_epoch.
…lint) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015wesmn95UL96xPBx3tXKQA
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015wesmn95UL96xPBx3tXKQA
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015wesmn95UL96xPBx3tXKQA
…tate continue Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015wesmn95UL96xPBx3tXKQA
… configurable Logs val_loss_p<q> per configured percentile (default 25/50/75; the former val_loss_median is now val_loss_p50). The callback configs also pointed at the pre-rename spine.pretext module path, which does not exist on main. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015wesmn95UL96xPBx3tXKQA
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Carries the training-framework improvements from
dev/unstabletomain(six cherry-picked commits, original messages kept). No pretext-task or backbone additions — those stay ondev/unstablefor a separate PR.What's included
init_fromwarm start —fit()can initialize the full pretext model (backbone + head) from a prior TransferCheckpoint; weights only, fresh optimizer.last.ckpt(<out stem>_state/by default, override withsave_state=), andresume_from=<last.ckpt>restores weights, optimizer moments, LR-scheduler state, early-stop counters, and the epoch/step loop.TransferCheckpointis now a stateful callback so its export threshold survives a resume. Mutually exclusive withinit_from.save_top_k=0Lightning writeslast.ckptonly aton_train_end(useless for crash/timeout recovery) → the checkpoint callback usesmonitor=None, save_top_k=1so every validation epoch refreshes it; and resuming a validation-end checkpoint replayson_train_epoch_endwithout validation metrics, which makes EarlyStopping's default train-epoch-end check raise → early stopping is evaluated at validation end.CurtainValLossPercentilescallback + configs — per-event validation-loss percentiles (val_loss_p<q>, configurable, default 25/50/75).Why
Multi-day pretrains have to cross SLURM walltime limits and maintenance reservations. With these changes a run is segment-chained losslessly: a TIMEOUT/kill resumes exactly where it stopped instead of falling back to best-so-far weights or a weights-only restart.
Validation
Verified with a kill-and-resume smoke (checkpoint written mid-run, job killed, resumed run logs "Restoring states…" and continues to the correct epoch), then used in production: two multi-segment pretraining chains (2 and 4 segments) ran to their epoch ceiling across walltime limits and a cluster maintenance window with bit-continuous optimizer/scheduler state.
🤖 Generated with Claude Code
https://claude.ai/code/session_015wesmn95UL96xPBx3tXKQA