Lightning-AI/pytorch-lightning

HPC Resubmit resume on most recent epoch checkpoint

Open

#13.773 geöffnet am 20. Juli 2022

Auf GitHub ansehen
 (11 Kommentare) (2 Reaktionen) (0 zugewiesene Personen)Python (3.233 Forks)batch import
checkpointingenvironment: slurmfeaturehelp wanted

Repository-Metriken

Stars
 (26.687 Stars)
PR-Merge-Metriken
 (Durchschn. Merge 9T 15h) (3 gemergte PRs in 30 T)

Beschreibung

🚀 Feature

Allow end-of-epoch checkpoints for resuming killed and resubmitted training jobs in a SLURM environment.

Motivation

Mid-epoch checkpointing does not appear to work with my model, even with fault-tolerant training I still get some weird results. Since I am training on a smaller dataset with a larger number of epochs, it would be really useful for me to be able to resume from the most recent checkpoint I saved using the normal end-of-epoch checkpoints.

Pitch

Instead of forcing users into the checkpoint process defined by the SLURMEnvironment plugin, allowing user's to customize the pause/resume operation would be a useful feature. Maybe add this as an option to the SLURMEnvironment plugin? Since my SLURM job ID is the same after resubmission, the default_root_dir is being set to the same as the previous job so the newest checkpoint should be easy to find.

Alternatives

Just an option for resuming end-of-epoch checkpoints would solve my problem. Allowing hooks for full customization of this function would be the most customizable version but put the most work on the user.

cc @borda @awaelchli @ananthsub @ninginthecloud @rohitgr7 @otaj

Contributor Guide