Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

(2/n) Support 2D Parallelism - Distributed Checkpoints #19852

Merged
merged 6 commits into from May 15, 2024
Merged

(2/n) Support 2D Parallelism - Distributed Checkpoints #19852

merged 6 commits into from May 15, 2024

Conversation

awaelchli
Copy link
Member

@awaelchli awaelchli commented May 7, 2024 •
edited by github-actions bot

What does this PR do?

Follow up to #19846

Adds the checkpointing logic to support all features we have in Fabric. This uses the new torch.distributed.checkpoint APIs from PyTorch 2.3+.

The caveat is that loading non-distributed checkpoints into a distributed model is now more difficult. I have not yet been able to resolve all edge cases around loading optimizer state, so for this iteration of the PR it is not yet supported to load optimizer states.

📚 Documentation preview 📚: https://pytorch-lightning--19852.org.readthedocs.build/en/19852/

cc @Borda @awaelchli @carmocca @justusschock

@awaelchli awaelchli added feature Is an improvement or enhancement checkpointing Related to checkpointing labels May 7, 2024
@github-actions github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package and removed feature Is an improvement or enhancement checkpointing Related to checkpointing labels May 7, 2024
@awaelchli awaelchli added this to the 2.3 milestone May 7, 2024
@awaelchli awaelchli added feature Is an improvement or enhancement checkpointing Related to checkpointing labels May 13, 2024
@awaelchli awaelchli marked this pull request as ready for review May 13, 2024 12:20
@github-actions GitHub Actions
Copy link
Contributor

github-actions bot commented May 13, 2024 •
edited

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow
Check ID Status
pl-cpu (macOS-11, lightning, 3.8, 2.0, oldest) success ✅
pl-cpu (macOS-11, lightning, 3.10, 2.0) success ✅
pl-cpu (macOS-11, lightning, 3.10, 2.1) success ✅
pl-cpu (macOS-11, lightning, 3.10, 2.2) success ✅
pl-cpu (macOS-14, lightning, 3.10, 2.3) success ✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 2.0, oldest) success ✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.0) success ✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.1) success ✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.2) success ✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.3) success ✅
pl-cpu (windows-2022, lightning, 3.8, 2.0, oldest) success ✅
pl-cpu (windows-2022, lightning, 3.10, 2.0) success ✅
pl-cpu (windows-2022, lightning, 3.10, 2.1) success ✅
pl-cpu (windows-2022, lightning, 3.10, 2.2) success ✅
pl-cpu (windows-2022, lightning, 3.10, 2.3) success ✅
pl-cpu (macOS-11, pytorch, 3.8, 2.0) success ✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 2.0) success ✅
pl-cpu (windows-2022, pytorch, 3.8, 2.0) success ✅
pl-cpu (macOS-12, pytorch, 3.11, 2.0) success ✅
pl-cpu (macOS-12, pytorch, 3.11, 2.1) success ✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.0) success ✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.1) success ✅
pl-cpu (windows-2022, pytorch, 3.11, 2.0) success ✅
pl-cpu (windows-2022, pytorch, 3.11, 2.1) success ✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py, src/lightning/pytorch/strategies/fsdp.py.

🟢 pytorch_lightning: Azure GPU
Check ID Status
pytorch-lightning (GPUs) (testing Lightning | latest) success ✅
pytorch-lightning (GPUs) (testing PyTorch | latest) success ✅

These checks are required after the changes to src/lightning/pytorch/strategies/fsdp.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py.

🟢 pytorch_lightning: Benchmarks
Check ID Status
lightning.Benchmarks success ✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py, src/lightning/pytorch/strategies/fsdp.py.

🟢 fabric: Docs
Check ID Status
docs-make (fabric, doctest) success ✅
docs-make (fabric, html) success ✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py.

🟢 pytorch_lightning: Docs
Check ID Status
docs-make (pytorch, doctest) success ✅
docs-make (pytorch, html) success ✅

These checks are required after the changes to src/lightning/pytorch/strategies/fsdp.py.

🟢 lightning_fabric: CPU workflow
Check ID Status
fabric-cpu (macOS-11, lightning, 3.8, 2.0, oldest) success ✅
fabric-cpu (macOS-11, lightning, 3.10, 2.0) success ✅
fabric-cpu (macOS-11, lightning, 3.11, 2.1) success ✅
fabric-cpu (macOS-11, lightning, 3.11, 2.2) success ✅
fabric-cpu (macOS-14, lightning, 3.10, 2.3) success ✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 2.0, oldest) success ✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 2.0) success ✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.1) success ✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.2) success ✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.3) success ✅
fabric-cpu (windows-2022, lightning, 3.8, 2.0, oldest) success ✅
fabric-cpu (windows-2022, lightning, 3.10, 2.0) success ✅
fabric-cpu (windows-2022, lightning, 3.11, 2.1) success ✅
fabric-cpu (windows-2022, lightning, 3.11, 2.2) success ✅
fabric-cpu (windows-2022, lightning, 3.11, 2.3) success ✅
fabric-cpu (macOS-11, fabric, 3.8, 2.0) success ✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 2.0) success ✅
fabric-cpu (windows-2022, fabric, 3.8, 2.0) success ✅
fabric-cpu (macOS-12, fabric, 3.11, 2.0) success ✅
fabric-cpu (macOS-12, fabric, 3.11, 2.1) success ✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.0) success ✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.1) success ✅
fabric-cpu (windows-2022, fabric, 3.11, 2.0) success ✅
fabric-cpu (windows-2022, fabric, 3.11, 2.1) success ✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py, tests/tests_fabric/strategies/test_fsdp.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/strategies/test_model_parallel.py, tests/tests_fabric/strategies/test_model_parallel_integration.py.

🟢 lightning_fabric: Azure GPU
Check ID Status
lightning-fabric (GPUs) (testing Fabric | latest) success ✅
lightning-fabric (GPUs) (testing Lightning | latest) success ✅

These checks are required after the changes to examples/fabric/tensor_parallel/parallelism.py, examples/fabric/tensor_parallel/train.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py, tests/tests_fabric/strategies/test_fsdp.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/strategies/test_model_parallel.py, tests/tests_fabric/strategies/test_model_parallel_integration.py.

🟢 mypy
Check ID Status
mypy success ✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py, src/lightning/pytorch/strategies/fsdp.py.

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

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/model_parallel.py, src/lightning/fabric/utilities/imports.py, src/lightning/pytorch/strategies/fsdp.py.

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.

@codecov Codecov
Copy link

codecov bot commented May 13, 2024 •
edited

Codecov Report

Attention: Patch coverage is 75.42373% with 29 lines in your changes are missing coverage. Please review.

Project coverage is 59%. Comparing base (90d04b5) to head (6948957).

Additional details and impacted files 
@@            Coverage Diff            @@
##           master   #19852     +/-   ##
=========================================
- Coverage      84%      59%    -25%     
=========================================
  Files         425      420      -5     
  Lines       34917    34912      -5     
=========================================
- Hits        29292    20525   -8767     
- Misses       5625    14387   +8762

carmocca
carmocca approved these changes May 13, 2024
View reviewed changes
Copy link
Member

@carmocca carmocca left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. It's nice to be able to ignore the non 2.3 code paths

src/lightning/fabric/strategies/model_parallel.py Outdated Show resolved Hide resolved
src/lightning/fabric/strategies/model_parallel.py Outdated Show resolved Hide resolved
src/lightning/fabric/strategies/model_parallel.py Outdated Show resolved Hide resolved
@carmocca carmocca mentioned this pull request May 13, 2024
Closed
@awaelchli awaelchli mentioned this pull request May 15, 2024
Merged
@lantiga lantiga merged commit 9455871 into master May 15, 2024
117 of 118 checks passed
@lantiga lantiga deleted the examples/tp-ckpt branch May 15, 2024 12:19
@lantiga
Copy link
Collaborator

lantiga commented May 15, 2024

Looks great!

@mergify mergify bot added the ready PRs ready to be merged label May 15, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@lantiga lantiga lantiga approved these changes

@carmocca carmocca carmocca approved these changes

@justusschock justusschock Awaiting requested review from justusschock justusschock is a code owner

@williamFalcon williamFalcon Awaiting requested review from williamFalcon williamFalcon is a code owner

@Borda Borda Awaiting requested review from Borda Borda is a code owner

@tchaton tchaton Awaiting requested review from tchaton tchaton is a code owner

Assignees
No one assigned
Labels
checkpointing Related to checkpointing fabric lightning.fabric.Fabric feature Is an improvement or enhancement pl Generic label for PyTorch Lightning package ready PRs ready to be merged
Projects
None yet
Milestone
2.3
Development

Successfully merging this pull request may close these issues.

None yet
3 participants
@awaelchli @lantiga @carmocca