Lightning-AI / Lightning-AI/pytorch-lightning
enable loading `universal checkpointing` checkpoint in `DeepSpeedStrategy`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
After I trained a model in some numbers of gpus, say, 8 gpus for a while, It's difficult to load the checkpoint to 16 gpus with optimizer and model states unchanged. The deepspeed has developed the universal checkpointing strategy to solve this problem, but I didn't see the `pytorch-lightning` has this feature.
### Pitch
I want the `pytorch-lightning` could support this feature
### Alternatives
try to add `universal_checkpoint` as a param of `DeepSpeedStrategy` and modify the class refering to `https://www.deepspeed.ai/tutorials/universal-checkpointing/`
### Additional context
_No response_
cc @borda @awaelchli
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 DeepSpeedStrategy entry point and the DeepSpeed universal checkpointing tutorial linked in the issue. Trace how model and optimizer checkpoints are currently saved and loaded, then define done as loading a checkpoint created with one GPU count onto a different GPU count while preserving those states.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 32/100