Lightning-AI/pytorch-lightning

enable loading `universal checkpointing` checkpoint in `DeepSpeedStrategy`

Open

#20,065 opened on 2024年7月9日

GitHub で見る
 (1 comment) (0 reactions) (0 assignees)Python (26,687 stars) (3,233 forks)batch import
featurehelp wantedstrategy: deepspeed

説明

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

コントリビューターガイド

技術スタック
python
領域
backendmachine learning
Issue 種別
feature
難度新しい貢献者にとっての実装難度です。1 は非常に小さな変更、5 は専門的な作業です。
5
推定時間経験ある貢献者が調査、実装、テスト、pull request 準備にかけるおおよその時間範囲です。
over 1 week
活動状況Issue が今どれだけ取り組みやすいかを示します。新しい、活発、古い、ブロック中、またはメンテナー入力待ちなどです。
fresh
明確さ期待される変更、受け入れ条件、次の手順がどれだけ明確かを示します。
mostly clear
前提条件
PyTorchDeepSpeedcheckpointing basics
初心者向け度初回貢献者にどれだけ取り組みやすいかを 1-100 で推定したスコアです。
10
調査方針
The issue requests adding universal checkpointing support to DeepSpeedStrategy. The contributor should study the DeepSpeed universal checkpointing tutorial at https://www.deepspeed.ai/tutorials/universal checkpointing/. The implementation likely involves modifying the DeepSpeedStrategy class in pytorch lightning's source code (e.g., lightning fabric/strategies/deepspeed.py). The maintainer comments from @borda and @awaelchli may provide guidance. The work includes adding a parameter to enable universal checkpoint and adapting the save/load methods.