(9/n) Support 2D Parallelism - Remaining Checkpoint Logic #19888

awaelchli · 2024-05-22T01:15:09Z

What does this PR do?

Implements the remaining distributed checkpoint saving and loading logic to the ModelParallelStrategy for Trainer.
The tests were adopted from the existing FSDP strategy tests.

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

cc @Borda @awaelchli @carmocca @justusschock

github-actions · 2024-05-22T11:14:47Z

⚡ 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/model_parallel.py, src/lightning/pytorch/strategies/model_parallel.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/strategies/test_model_parallel.py, tests/tests_pytorch/strategies/test_model_parallel_integration.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/model_parallel.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/strategies/test_model_parallel.py, tests/tests_pytorch/strategies/test_model_parallel_integration.py, src/lightning/fabric/strategies/model_parallel.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/model_parallel.py, src/lightning/pytorch/strategies/model_parallel.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/model_parallel.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/model_parallel.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/model_parallel.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 src/lightning/fabric/strategies/model_parallel.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/model_parallel.py, src/lightning/pytorch/strategies/model_parallel.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/model_parallel.py, src/lightning/pytorch/strategies/model_parallel.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.

src/lightning/fabric/strategies/model_parallel.py

tests/tests_pytorch/strategies/test_fsdp.py

codecov · 2024-05-22T11:31:37Z

Codecov Report

Attention: Patch coverage is 82.75862% with 5 lines in your changes missing coverage. Please review.

Project coverage is 59%. Comparing base (987c2c4) to head (9ad7c2a).
Report is 121 commits behind head on master.

❗ There is a different number of reports uploaded between BASE (987c2c4) and HEAD (9ad7c2a). Click for more details.

HEAD has 64 uploads less than BASE

Flag BASE (987c2c4) HEAD (9ad7c2a)

python3.10 19 16

cpu 65 48

lightning 39 32

pytest 45 28

examples 9 0

app 9 0

lightning_app 2 0

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #19888     +/-   ##
=========================================
- Coverage      84%      59%    -25%     
=========================================
  Files         426      421      -5     
  Lines       35233    35149     -84     
=========================================
- Hits        29506    20745   -8761     
- Misses       5727    14404   +8677

lantiga

LGTM!

src/lightning/pytorch/strategies/model_parallel.py

Co-authored-by: Luca Antiga <[email protected]>

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

awaelchli added feature Is an improvement or enhancement pl Generic label for PyTorch Lightning package labels May 22, 2024

github-actions bot added the fabric lightning.fabric.Fabric label May 22, 2024

awaelchli added this to the 2.3 milestone May 22, 2024

awaelchli added the checkpointing Related to checkpointing label May 22, 2024

awaelchli marked this pull request as ready for review May 22, 2024 11:14

awaelchli requested review from carmocca, justusschock, Borda and williamFalcon as code owners May 22, 2024 11:14

awaelchli commented May 22, 2024

View reviewed changes

src/lightning/fabric/strategies/model_parallel.py Outdated Show resolved Hide resolved

awaelchli commented May 22, 2024

View reviewed changes

tests/tests_pytorch/strategies/test_fsdp.py Show resolved Hide resolved

checkpoint saving and loading

f558f07

awaelchli force-pushed the feature/pl-model-parallel-checkpoints branch from 33f47ef to f558f07 Compare May 22, 2024 11:21

awaelchli added 2 commits May 22, 2024 08:02

test

c8ae9e4

test

0ed924d

carmocca approved these changes May 22, 2024

View reviewed changes

awaelchli requested a review from lantiga May 22, 2024 18:00

lantiga approved these changes May 22, 2024

View reviewed changes

src/lightning/pytorch/strategies/model_parallel.py Outdated Show resolved Hide resolved

mergify bot added the ready PRs ready to be merged label May 22, 2024

awaelchli and others added 2 commits May 22, 2024 17:40

Update src/lightning/pytorch/strategies/model_parallel.py

3768328

Co-authored-by: Luca Antiga <[email protected]>

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

9ad7c2a

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

awaelchli merged commit 414c863 into master May 22, 2024
117 of 118 checks passed

awaelchli deleted the feature/pl-model-parallel-checkpoints branch May 22, 2024 22:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

(9/n) Support 2D Parallelism - Remaining Checkpoint Logic #19888

(9/n) Support 2D Parallelism - Remaining Checkpoint Logic #19888

awaelchli commented May 22, 2024 •

edited by github-actions bot

Loading

github-actions bot commented May 22, 2024 •

edited

Loading

codecov bot commented May 22, 2024 •

edited

Loading

lantiga left a comment

(9/n) Support 2D Parallelism - Remaining Checkpoint Logic #19888

(9/n) Support 2D Parallelism - Remaining Checkpoint Logic #19888

Conversation

awaelchli commented May 22, 2024 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented May 22, 2024 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented May 22, 2024 • edited Loading

Codecov Report

lantiga left a comment

Choose a reason for hiding this comment

awaelchli commented May 22, 2024 •

edited by github-actions bot

Loading

github-actions bot commented May 22, 2024 •

edited

Loading

codecov bot commented May 22, 2024 •

edited

Loading