FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric #20440

lantiga · 2024-11-22T16:13:18Z

What does this PR do?

This PR adds:

examples for FP8 (via torchao), FSDP2 (via DTensors) and torch.compile for both PyTorch Lightning and Fabric, together with the respective READMEs
tests for calling torch.compile from configure_module and parallelize_fn when using ModelParallelStrategy
docs for torch.compile called from configure_module and parallelize_fn in PyTorch Lightning and Fabric

This PR unearthed several issues with integration tests, which are addressed here to get CI green.

cc @Borda @awaelchli @justusschock

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

github-actions · 2024-11-25T23:30:23Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-14, lightning, 3.9, 2.1, oldest)	success	✅
pl-cpu (macOS-14, lightning, 3.10, 2.1)	success	✅
pl-cpu (macOS-14, lightning, 3.11, 2.2.2)	success	✅
pl-cpu (macOS-14, lightning, 3.11, 2.3)	success	✅
pl-cpu (macOS-14, lightning, 3.12, 2.4.1)	success	✅
pl-cpu (macOS-14, lightning, 3.12, 2.5.1)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.9, 2.1, oldest)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.10, 2.1)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.11, 2.2.2)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.11, 2.3)	success	✅
pl-cpu (ubuntu-22.04, lightning, 3.12, 2.4.1)	success	✅
pl-cpu (ubuntu-22.04, lightning, 3.12, 2.5.1)	success	✅
pl-cpu (windows-2022, lightning, 3.9, 2.1, oldest)	success	✅
pl-cpu (windows-2022, lightning, 3.10, 2.1)	success	✅
pl-cpu (windows-2022, lightning, 3.11, 2.2.2)	success	✅
pl-cpu (windows-2022, lightning, 3.11, 2.3)	success	✅
pl-cpu (windows-2022, lightning, 3.12, 2.4.1)	success	✅
pl-cpu (windows-2022, lightning, 3.12, 2.5.1)	success	✅
pl-cpu (macOS-14, pytorch, 3.9, 2.1)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.9, 2.1)	success	✅
pl-cpu (windows-2022, pytorch, 3.9, 2.1)	success	✅
pl-cpu (macOS-14, pytorch, 3.12, 2.5.1)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.12, 2.5.1)	success	✅
pl-cpu (windows-2022, pytorch, 3.12, 2.5.1)	success	✅

These checks are required after the changes to 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 tests/tests_pytorch/strategies/test_model_parallel_integration.py, tests/run_standalone_tests.sh.

🟢 fabric: Docs

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

These checks are required after the changes to docs/source-fabric/advanced/compile.rst.

🟢 pytorch_lightning: Docs

Check ID	Status
docs-make (pytorch, doctest)	success	✅
docs-make (pytorch, html)	success	✅

These checks are required after the changes to docs/source-pytorch/advanced/compile.rst.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-14, lightning, 3.9, 2.1, oldest)	success	✅
fabric-cpu (macOS-14, lightning, 3.10, 2.1)	success	✅
fabric-cpu (macOS-14, lightning, 3.11, 2.2.2)	success	✅
fabric-cpu (macOS-14, lightning, 3.11, 2.3)	success	✅
fabric-cpu (macOS-14, lightning, 3.12, 2.4.1)	success	✅
fabric-cpu (macOS-14, lightning, 3.12, 2.5.1)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.9, 2.1, oldest)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.10, 2.1)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.2.2)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.11, 2.3)	success	✅
fabric-cpu (ubuntu-22.04, lightning, 3.12, 2.4.1)	success	✅
fabric-cpu (ubuntu-22.04, lightning, 3.12, 2.5.1)	success	✅
fabric-cpu (windows-2022, lightning, 3.9, 2.1, oldest)	success	✅
fabric-cpu (windows-2022, lightning, 3.10, 2.1)	success	✅
fabric-cpu (windows-2022, lightning, 3.11, 2.2.2)	success	✅
fabric-cpu (windows-2022, lightning, 3.11, 2.3)	success	✅
fabric-cpu (windows-2022, lightning, 3.12, 2.4.1)	success	✅
fabric-cpu (windows-2022, lightning, 3.12, 2.5.1)	success	✅
fabric-cpu (macOS-14, fabric, 3.9, 2.1)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.9, 2.1)	success	✅
fabric-cpu (windows-2022, fabric, 3.9, 2.1)	success	✅
fabric-cpu (macOS-14, fabric, 3.12, 2.5.1)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.12, 2.5.1)	success	✅
fabric-cpu (windows-2022, fabric, 3.12, 2.5.1)	success	✅

These checks are required after the changes to tests/tests_fabric/conftest.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/fp8_distributed_transformer/requirements.txt, examples/fabric/fp8_distributed_transformer/train.py, tests/tests_fabric/conftest.py, tests/tests_fabric/strategies/test_model_parallel_integration.py, tests/run_standalone_tests.sh.

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-11-25T23:47:48Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 88%. Comparing base (8ce5287) to head (71a4b97).
Report is 8 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master   #20440   +/-   ##
=======================================
  Coverage      88%      88%           
=======================================
  Files         267      267           
  Lines       23203    23274   +71     
=======================================
+ Hits        20313    20381   +68     
- Misses       2890     2893    +3

lantiga · 2024-11-26T10:26:47Z

Investigating CI failures in isolation, the same code succeeds when running standalone but something's up.

lantiga · 2024-11-26T10:28:59Z

Also, updates to docs at https://lightning.ai/docs/pytorch/stable/advanced/model_parallel/ will be coming as a follow up

Borda

I think would be easier and more robust to refer the examples directly in docs and not copy the code which maybe yield in accidently updating only one of them...

lantiga · 2024-11-26T11:42:13Z

I think would be easier and more robust to refer the examples directly in docs and not copy the code which maybe yield in accidently updating only one of them...

That's fair, will code examples expand on the page?

Also, docs skip the torch ao part, since it's an optional / cutting-edge dependency (it depends on triton nightly) so it only appears in the examples

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

lantiga · 2024-11-26T14:29:28Z

Ok CI has been fixed. Merging, the only failing tests are doc links to master (which don't exist yet until this is merged).

Minimal transformer examples

c6695f8

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

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

fcf115a

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

lantiga self-assigned this Nov 22, 2024

lantiga added the distributed Generic distributed-related topic label Nov 22, 2024

This was referenced Nov 24, 2024

Add fp8 (torchao)/fsdp2/torch_compile handlers and tests #20443

Closed

Add fp8 (torchao)/fsdp2/torch_compile handlers and tests #20445

Draft

lantiga added 3 commits November 25, 2024 22:33

Add tests for compile after fsdp2/tp

8be123a

Add README's

27c36d9

Add docs

ca3c4cf

lantiga marked this pull request as ready for review November 25, 2024 23:29

lantiga requested review from Borda, tchaton, justusschock and ethanwharris as code owners November 25, 2024 23:29

Rename folder, add cross-reference

55a6fde

lantiga force-pushed the luca/fp8-fsdp2-compile-examples branch from 1cf53e6 to 55a6fde Compare November 25, 2024 23:36

github-actions bot added the docs Documentation related label Nov 25, 2024

Fix link

06c6af6

lantiga changed the title ~~Minimal FP8 + FSDP2 + compile examples for PyTorch Lightning and Fabric~~ FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric Nov 25, 2024

lantiga and others added 4 commits November 26, 2024 01:13

Newline after code-block directive

acb3125

Update section name

1646eca

Fix reference

3d6cea6

Half standalone tests batch size

4eaac17

tchaton approved these changes Nov 26, 2024

View reviewed changes

Borda approved these changes Nov 26, 2024

View reviewed changes

mergify bot added the ready PRs ready to be merged label Nov 26, 2024

lantiga and others added 2 commits November 26, 2024 12:37

Fix integration tests

a8310cf

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

71a4b97

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

lantiga merged commit 87565cb into master Nov 26, 2024
70 of 72 checks passed

lantiga deleted the luca/fp8-fsdp2-compile-examples branch November 26, 2024 14:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric #20440

FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric #20440

lantiga commented Nov 22, 2024 •

edited

Loading

github-actions bot commented Nov 25, 2024 •

edited

Loading

codecov bot commented Nov 25, 2024 •

edited

Loading

lantiga commented Nov 26, 2024

lantiga commented Nov 26, 2024

Borda left a comment

lantiga commented Nov 26, 2024 •

edited

Loading

lantiga commented Nov 26, 2024

FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric #20440

FP8 + FSDP2 + torch.compile examples for PyTorch Lightning and Fabric #20440

Conversation

lantiga commented Nov 22, 2024 • edited Loading

What does this PR do?

github-actions bot commented Nov 25, 2024 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Nov 25, 2024 • edited Loading

Codecov Report

lantiga commented Nov 26, 2024

lantiga commented Nov 26, 2024

Borda left a comment

Choose a reason for hiding this comment

lantiga commented Nov 26, 2024 • edited Loading

lantiga commented Nov 26, 2024

lantiga commented Nov 22, 2024 •

edited

Loading

github-actions bot commented Nov 25, 2024 •

edited

Loading

codecov bot commented Nov 25, 2024 •

edited

Loading

lantiga commented Nov 26, 2024 •

edited

Loading