Lightning-AI / Lightning-AI/pytorch-lightning

enable loading `universal checkpointing` checkpoint in `DeepSpeedStrategy`

Open
#20,065 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

feature help wanted strategy: deepspeed
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.