(5/n) Support 2D Parallelism in Lightning Trainer #19878

awaelchli · 2024-05-16T20:32:52Z

What does this PR do?

Ports the functionality added in #19846 to the Trainer. The same tests are adopted and rewritten for the Trainer semantics. Some tests were also taken from the Trainer FSDP Strategy test files.

To keep the PRs minimal, I'm not including checkpointing here and will submit it in a follow up PR.
A concrete example of how to use the strategy is in #19879. In summary:

# 1. Define your LightningModule normally
class MyParallelModel(LightningModule):

	# 2. Add this hook
    def configure_model(self):
        # 3. Add parallelization here. You can access `self.device_mesh`.
        tp_mesh = self.device_mesh["tensor_parallel"]
        dp_mesh = self.device_mesh["data_parallel"]
		
        plan = {"layer.w1": ColwiseParallel())
        parallelize_module(model, tp_mesh, plan)


# 4. Select the `ModelParallelStrategy` in the Trainer.
from lightning.pytorch.strategies import ModelParallelStrategy

strategy = ModelParallelStrategy()
trainer = Trainer(strategy=strategy) 

# Defaults are:
strategy = ModelParallelStrategy(
	tensor_parallel_size="auto",   # number of GPUs in the machine
    data_parallel_size="auto",  # number of machines in the cluster
)

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

cc @Borda @carmocca @justusschock @awaelchli

github-actions · 2024-05-17T09:54:13Z

⚡ 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/core/module.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/model_parallel.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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/core/module.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/model_parallel.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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/core/module.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/model_parallel.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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/core/module.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/model_parallel.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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, 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 src/lightning/fabric/strategies/model_parallel.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/model_parallel.py, src/lightning/pytorch/core/module.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/model_parallel.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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/core/module.py, src/lightning/pytorch/strategies/__init__.py, src/lightning/pytorch/strategies/model_parallel.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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 · 2024-05-17T14:04:48Z

Codecov Report

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

Project coverage is 59%. Comparing base (1d0c6aa) to head (8eebff9).

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #19878     +/-   ##
=========================================
- Coverage      84%      59%    -25%     
=========================================
  Files         425      421      -4     
  Lines       35028    35139    +111     
=========================================
- Hits        29319    20714   -8605     
- Misses       5709    14425   +8716

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

lantiga

This looks great

src/lightning/pytorch/strategies/model_parallel.py

awaelchli changed the base branch from master to feature/tp-full-optim-load May 16, 2024 20:33

github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels May 16, 2024

awaelchli mentioned this pull request May 16, 2024

(6/n) Support 2D Parallelism - Trainer example #19879

Merged

awaelchli added this to the 2.3 milestone May 17, 2024

awaelchli added the feature Is an improvement or enhancement label May 17, 2024

awaelchli marked this pull request as ready for review May 17, 2024 09:53

awaelchli requested review from carmocca, justusschock, williamFalcon, tchaton and Borda as code owners May 17, 2024 09:53

awaelchli force-pushed the feature/tp-pl-strategy branch from c4d3f75 to 2585763 Compare May 17, 2024 13:45

awaelchli changed the base branch from feature/tp-full-optim-load to master May 17, 2024 13:46

ModelParallelStrategy for Lightning Trainer

d806b64

awaelchli force-pushed the feature/tp-pl-strategy branch from 2585763 to d806b64 Compare May 17, 2024 13:48

awaelchli added 2 commits May 17, 2024 15:52

mypy

3a54bf7

import fix

76e206d

awaelchli and others added 4 commits May 17, 2024 16:17

fix torchscript errors

a7937b6

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

939d5e9

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

fix docs issue

5e44dc1

fix test execution

9b34c99

justusschock approved these changes May 17, 2024

View reviewed changes

lantiga reviewed May 17, 2024

View reviewed changes

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

lantiga approved these changes May 17, 2024

View reviewed changes

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

lantiga added 2 commits May 17, 2024 18:15

Merge branch 'master' into feature/tp-pl-strategy

38f9317

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

8eebff9

lantiga merged commit 32e2418 into master May 17, 2024
116 of 117 checks passed

lantiga deleted the feature/tp-pl-strategy branch May 17, 2024 23:03

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

(5/n) Support 2D Parallelism in Lightning Trainer #19878

(5/n) Support 2D Parallelism in Lightning Trainer #19878

awaelchli commented May 16, 2024 •

edited by github-actions bot

github-actions bot commented May 17, 2024 •

edited

codecov bot commented May 17, 2024 •

edited

lantiga left a comment

(5/n) Support 2D Parallelism in Lightning Trainer #19878

(5/n) Support 2D Parallelism in Lightning Trainer #19878

Conversation

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

What does this PR do?

github-actions bot commented May 17, 2024 • edited

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented May 17, 2024 • edited

Codecov Report

lantiga left a comment

Choose a reason for hiding this comment

awaelchli commented May 16, 2024 •

edited by github-actions bot

github-actions bot commented May 17, 2024 •

edited

codecov bot commented May 17, 2024 •

edited