Relax the assumption that the root module is FSDP wrapped #18054

carmocca · 2023-07-11T00:54:11Z

What does this PR do?

cc @Borda @carmocca @justusschock @awaelchli

github-actions · 2023-07-15T00:29:39Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

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

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/pytorch/strategies/fsdp.py, tests/tests_pytorch/strategies/test_fsdp.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

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

🟢 fabric: Docs

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

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

🟢 pytorch_lightning: Docs

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

These checks are required after the changes to src/lightning/pytorch/strategies/fsdp.py, docs/source-pytorch/advanced/model_parallel.rst.

🟢 lightning_fabric: CPU workflow

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

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, tests/tests_fabric/strategies/test_fsdp.py, tests/tests_fabric/strategies/test_fsdp_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/fsdp.py, tests/tests_fabric/strategies/test_fsdp.py, tests/tests_fabric/strategies/test_fsdp_integration.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.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.10)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.8)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.10)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.8)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.10)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.8)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.10)	success	✅
install-pkg (ubuntu-22.04, notset, 3.8)	success	✅
install-pkg (ubuntu-22.04, notset, 3.10)	success	✅
install-pkg (macOS-12, app, 3.8)	success	✅
install-pkg (macOS-12, app, 3.10)	success	✅
install-pkg (macOS-12, fabric, 3.8)	success	✅
install-pkg (macOS-12, fabric, 3.10)	success	✅
install-pkg (macOS-12, pytorch, 3.8)	success	✅
install-pkg (macOS-12, pytorch, 3.10)	success	✅
install-pkg (macOS-12, lightning, 3.8)	success	✅
install-pkg (macOS-12, lightning, 3.10)	success	✅
install-pkg (macOS-12, notset, 3.8)	success	✅
install-pkg (macOS-12, notset, 3.10)	success	✅
install-pkg (windows-2022, app, 3.8)	success	✅
install-pkg (windows-2022, app, 3.10)	success	✅
install-pkg (windows-2022, fabric, 3.8)	success	✅
install-pkg (windows-2022, fabric, 3.10)	success	✅
install-pkg (windows-2022, pytorch, 3.8)	success	✅
install-pkg (windows-2022, pytorch, 3.10)	success	✅
install-pkg (windows-2022, lightning, 3.8)	success	✅
install-pkg (windows-2022, lightning, 3.10)	success	✅
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.10)	success	✅

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

🟢 link-check

Check ID	Status
check-md-links / 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/fsdp.py

codecov · 2023-07-17T18:24:58Z

Codecov Report

Merging #18054 (d97c083) into master (7116a9f) will decrease coverage by 23%.
The diff coverage is 68%.

❗ Current head d97c083 differs from pull request most recent head ca55a97. Consider uploading reports for the commit ca55a97 to get more accurate results

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #18054     +/-   ##
=========================================
- Coverage      84%      61%    -23%     
=========================================
  Files         430      425      -5     
  Lines       32687    32585    -102     
=========================================
- Hits        27344    19811   -7533     
- Misses       5343    12774   +7431

tests/tests_fabric/strategies/test_fsdp.py

awaelchli · 2023-07-18T00:36:31Z

Glad it was relatively uncomplicated to do this! Great work

justusschock

Looks good to me, only one question :)

src/lightning/fabric/strategies/fsdp.py

carmocca added feature Is an improvement or enhancement fabric lightning.fabric.Fabric strategy: fsdp Fully Sharded Data Parallel pl Generic label for PyTorch Lightning package labels Jul 11, 2023

carmocca added this to the 2.1 milestone Jul 11, 2023

carmocca self-assigned this Jul 11, 2023

WIP

48153d4

carmocca force-pushed the carmocca/fsdp-non-root branch from 5806465 to 48153d4 Compare July 11, 2023 00:59

carmocca added 6 commits July 11, 2023 03:57

Fabric WIP

1857d83

Merge branch 'master' into carmocca/fsdp-non-root

62b7ad3

Update

fb6df52

CHANGELOG

1994d12

Turns out filter with sharded checkpoint doesnt work

ae67328

Fix

9b190e6

carmocca marked this pull request as ready for review July 15, 2023 00:29

carmocca requested review from awaelchli, justusschock, williamFalcon, edenlightning, lantiga and Borda as code owners July 15, 2023 00:29

carmocca added the bug Something isn't working label Jul 15, 2023

carmocca added 3 commits July 15, 2023 03:03

mypy

0f28b88

Fix

0c641c1

Merge branch 'master' into carmocca/fsdp-non-root

cd89211

mergify bot added the has conflicts label Jul 15, 2023

awaelchli reviewed Jul 17, 2023

View reviewed changes

src/lightning/fabric/strategies/fsdp.py Show resolved Hide resolved

src/lightning/fabric/strategies/fsdp.py Show resolved Hide resolved

Merge branch 'master' into carmocca/fsdp-non-root

d97c083

mergify bot removed the has conflicts label Jul 17, 2023

mergify bot added the has conflicts label Jul 17, 2023

carmocca added 4 commits July 18, 2023 01:53

Warning test

e152a99

Review comments

8a5204e

Mock

55711bb

Merge branch 'master' into carmocca/fsdp-non-root

ca55a97

carmocca requested a review from awaelchli July 18, 2023 00:08

mergify bot removed the has conflicts label Jul 18, 2023

mypy

1980f8a

awaelchli reviewed Jul 18, 2023

View reviewed changes

tests/tests_fabric/strategies/test_fsdp.py Outdated Show resolved Hide resolved

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

Revert

c230eca

carmocca requested a review from awaelchli July 18, 2023 00:34

awaelchli approved these changes Jul 18, 2023

View reviewed changes

Link to torch issue

03a99e1

mergify bot added the has conflicts label Jul 18, 2023

justusschock approved these changes Jul 19, 2023

View reviewed changes

src/lightning/fabric/strategies/fsdp.py Show resolved Hide resolved

Merge branch 'master' into carmocca/fsdp-non-root

53e7269

mergify bot added ready PRs ready to be merged and removed has conflicts ready PRs ready to be merged labels Jul 19, 2023

carmocca merged commit d653e4e into master Jul 19, 2023

carmocca deleted the carmocca/fsdp-non-root branch July 19, 2023 13:34

carmocca mentioned this pull request Jul 25, 2023

[TPU] Add Fabric support for PyTorch XLA FSDP #18126

Merged

asuglia-alana mentioned this pull request Nov 4, 2023

Root module wrapped when using custom FSDP strategy #18944

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Relax the assumption that the root module is FSDP wrapped #18054

Relax the assumption that the root module is FSDP wrapped #18054

carmocca commented Jul 11, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jul 15, 2023 •

edited

Loading

codecov bot commented Jul 17, 2023 •

edited

Loading

awaelchli commented Jul 18, 2023

justusschock left a comment

Relax the assumption that the root module is FSDP wrapped #18054

Relax the assumption that the root module is FSDP wrapped #18054

Conversation

carmocca commented Jul 11, 2023 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented Jul 15, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Jul 17, 2023 • edited Loading

Codecov Report

awaelchli commented Jul 18, 2023

justusschock left a comment

Choose a reason for hiding this comment

carmocca commented Jul 11, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jul 15, 2023 •

edited

Loading

codecov bot commented Jul 17, 2023 •

edited

Loading