Fabric checkpointing 2/n: DeepSpeed implementation #16452

awaelchli · 2023-01-20T12:17:56Z

What does this PR do?

Implement checkpoint saving and loading methods on the DeepSpeedStrategy in Fabric.
Follow up to #16434
Part of #14816

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

I made sure I had fun coding 🙃

cc @Borda @awaelchli @carmocca @justusschock

github-actions · 2023-01-20T14:01:28Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, pytorch, 3.8, 1.11)	success	✅
pl-cpu (macOS-11, pytorch, 3.9, 1.12)	success	✅
pl-cpu (macOS-11, pytorch, 3.10, 1.13)	success	✅
pl-cpu (macOS-11, pytorch, 3.8, 1.10, oldest)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 1.10)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.9, 1.11)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.10, 1.12)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.10, 1.13)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.7, 1.10, oldest)	success	✅
pl-cpu (windows-2022, pytorch, 3.9, 1.11)	success	✅
pl-cpu (windows-2022, pytorch, 3.10, 1.12)	success	✅
pl-cpu (windows-2022, pytorch, 3.10, 1.13)	success	✅
pl-cpu (windows-2022, pytorch, 3.7, 1.10, oldest)	success	✅
pl-cpu (slow, macOS-11, pytorch, 3.7, 1.11)	success	✅
pl-cpu (slow, ubuntu-20.04, pytorch, 3.7, 1.11)	success	✅
pl-cpu (slow, windows-2022, pytorch, 3.7, 1.11)	success	✅
pl-cpu (macOS-11, lightning, 3.8, 1.13)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.13)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py.

🟢 pytorch_lightning: Azure HPU

Check ID	Status
pytorch-lightning (HPUs)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py.

🟢 pytorch_lightning: Azure IPU

Check ID	Status
pytorch-lightning (IPUs)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, fabric, 3.8, 1.11)	success	✅
fabric-cpu (macOS-11, fabric, 3.9, 1.12)	success	✅
fabric-cpu (macOS-11, fabric, 3.10, 1.13)	success	✅
fabric-cpu (macOS-11, fabric, 3.7, 1.10, oldest)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 1.10)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.9, 1.11)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.10, 1.12)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.10, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.7, 1.10, oldest)	success	✅
fabric-cpu (windows-2022, fabric, 3.9, 1.11)	success	✅
fabric-cpu (windows-2022, fabric, 3.10, 1.12)	success	✅
fabric-cpu (windows-2022, fabric, 3.10, 1.13)	success	✅
fabric-cpu (windows-2022, fabric, 3.7, 1.10, oldest)	success	✅
fabric-cpu (macOS-11, lightning, 3.8, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.13)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py, tests/tests_fabric/strategies/test_deepspeed.py, tests/tests_fabric/strategies/test_deepspeed_integration.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py, tests/tests_fabric/strategies/test_deepspeed.py, tests/tests_fabric/strategies/test_deepspeed_integration.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py.

🟢 install

Check ID	Status
install-pkg (ubuntu-22.04, app, 3.7)	success	✅
install-pkg (ubuntu-22.04, app, 3.10)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.7)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.10)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.7)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.10)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.7)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.10)	success	✅
install-pkg (ubuntu-22.04, notset, 3.7)	success	✅
install-pkg (ubuntu-22.04, notset, 3.10)	success	✅
install-pkg (macOS-12, app, 3.7)	success	✅
install-pkg (macOS-12, app, 3.10)	success	✅
install-pkg (macOS-12, fabric, 3.7)	success	✅
install-pkg (macOS-12, fabric, 3.10)	success	✅
install-pkg (macOS-12, pytorch, 3.7)	success	✅
install-pkg (macOS-12, pytorch, 3.10)	success	✅
install-pkg (macOS-12, lightning, 3.7)	success	✅
install-pkg (macOS-12, lightning, 3.10)	success	✅
install-pkg (macOS-12, notset, 3.7)	success	✅
install-pkg (macOS-12, notset, 3.10)	success	✅
install-pkg (windows-2022, app, 3.7)	success	✅
install-pkg (windows-2022, app, 3.10)	success	✅
install-pkg (windows-2022, fabric, 3.7)	success	✅
install-pkg (windows-2022, fabric, 3.10)	success	✅
install-pkg (windows-2022, pytorch, 3.7)	success	✅
install-pkg (windows-2022, pytorch, 3.10)	success	✅
install-pkg (windows-2022, lightning, 3.7)	success	✅
install-pkg (windows-2022, lightning, 3.10)	success	✅
install-pkg (windows-2022, notset, 3.7)	success	✅
install-pkg (windows-2022, notset, 3.10)	success	✅

These checks are required after the changes to src/lightning_fabric/strategies/deepspeed.py.

🟢 link-check

Check ID	Status
markdown-link-check	success	✅

These checks are required after the changes to src/lightning_fabric/CHANGELOG.md.

Thank you for your contribution! 💜

Note
This comment is automatically generated and updates for 60 minutes every 180 seconds. If you have any other questions, contact carmocca for help.

src/lightning_fabric/strategies/deepspeed.py

Co-authored-by: Carlos Mocholí <[email protected]>

for more information, see https://pre-commit.ci

Co-authored-by: Carlos Mocholí <[email protected]>

docs/source-pytorch/fabric/guide/checkpoint.rst

pick

f4a1c2d

github-actions bot added the pl Generic label for PyTorch Lightning package label Jan 20, 2023

squashs all commits

a3c97c9

awaelchli force-pushed the lite/checkpoint-deepspeed branch from b9df0dd to a3c97c9 Compare January 20, 2023 12:23

github-actions bot removed the pl Generic label for PyTorch Lightning package label Jan 20, 2023

awaelchli added 7 commits January 20, 2023 13:34

engine validation

0e6a3ca

update

8c81413

save testing

2694556

testing

ef5ad47

test

6637691

format doc

f77d84a

changelog updaste

339ae46

awaelchli added checkpointing Related to checkpointing strategy: deepspeed fabric lightning.fabric.Fabric labels Jan 20, 2023

awaelchli added this to the 2.0 milestone Jan 20, 2023

awaelchli self-assigned this Jan 20, 2023

awaelchli marked this pull request as ready for review January 20, 2023 14:00

awaelchli requested review from carmocca and justusschock as code owners January 20, 2023 14:00

Merge branch 'master' into lite/checkpoint-deepspeed

53ed69f

awaelchli requested review from Borda and akihironitta as code owners January 20, 2023 21:25

awaelchli force-pushed the lite/checkpoint-deepspeed branch from ea7f93f to 53ed69f Compare January 20, 2023 21:36

mergify bot added the has conflicts label Jan 22, 2023

Merge branch 'master' into lite/checkpoint-deepspeed

87f2e4e

mergify bot removed the has conflicts label Jan 23, 2023

justusschock approved these changes Jan 23, 2023

View reviewed changes

carmocca reviewed Jan 24, 2023

View reviewed changes

awaelchli and others added 2 commits January 24, 2023 07:43

Update src/lightning_fabric/strategies/deepspeed.py

4b4c95d

Co-authored-by: Carlos Mocholí <[email protected]>

Typo

e2e60a8

carmocca requested review from edenlightning and lantiga as code owners January 24, 2023 13:43

github-actions bot added the pl Generic label for PyTorch Lightning package label Jan 24, 2023

awaelchli and others added 5 commits January 24, 2023 15:25

warning

46b925c

test

639e2d3

[pre-commit.ci] auto fixes from pre-commit.com hooks

44a0884

for more information, see https://pre-commit.ci

Merge branch 'master' into lite/checkpoint-deepspeed

267697b

Update src/lightning_fabric/strategies/deepspeed.py

b8fe273

Co-authored-by: Carlos Mocholí <[email protected]>

awaelchli commented Jan 24, 2023

View reviewed changes

docs/source-pytorch/fabric/guide/checkpoint.rst Outdated Show resolved Hide resolved

Update docs/source-pytorch/fabric/guide/checkpoint.rst

65191bc

github-actions bot removed the pl Generic label for PyTorch Lightning package label Jan 24, 2023

carmocca approved these changes Jan 24, 2023

View reviewed changes

carmocca merged commit 7603dd0 into master Jan 24, 2023

carmocca deleted the lite/checkpoint-deepspeed branch January 24, 2023 17:53

mergify bot added the ready PRs ready to be merged label Jan 24, 2023

awaelchli mentioned this pull request Jan 24, 2023

Fabric checkpointing 3/n: Implement missing get_module_state_dict for strategies #16487

Merged

11 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fabric checkpointing 2/n: DeepSpeed implementation #16452

Fabric checkpointing 2/n: DeepSpeed implementation #16452

awaelchli commented Jan 20, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jan 20, 2023 •

edited

Loading

Fabric checkpointing 2/n: DeepSpeed implementation #16452

Fabric checkpointing 2/n: DeepSpeed implementation #16452

Conversation

awaelchli commented Jan 20, 2023 • edited by github-actions bot Loading

What does this PR do?

Before submitting

PR review

Did you have fun?

github-actions bot commented Jan 20, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

awaelchli commented Jan 20, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jan 20, 2023 •

edited

Loading