Lightning-AI / Lightning-AI/pytorch-lightning
HPC Resubmit resume on most recent epoch checkpoint
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
## 🚀 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the SLURMEnvironment plugin and its checkpoint pause/resume process, then trace how default_root_dir and job IDs are used after resubmission. Done should support resuming killed and resubmitted SLURM training jobs from the most recent normal end-of-epoch checkpoint, with the requested customization or option clearly defined.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- hpc
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 30/100