Combined setup of model and optimizer with FSDP #17305

awaelchli · 2023-04-08T15:20:51Z

What does this PR do?

User does no longer have to change their code when switching to FSDP.
Before:

model = fabric.setup_module(model)
# you had to reference parameters from wrapped model
optimizer = Adam(model.parameters())
optimizer = fabric.setup_optimizers(optimizer)

Now:

model = ...
optimizer = ...

# same for all strategies now
model, optimizer = fabric.setup(model, optimizer)

User can specify multiple param groups in the optimizer. This was not possible before.
User does no longer see FlatParameter when inspecting the parameters of optimizer or model, just normal parameters.
This is only possible through the FSDP(use_orig_params=True) which was introduced in PyTorch 2.0. The old code stays until PyTorch 2.0 becomes our minimum supported version.

The corresponding update for PL is here: #17309

cc @Borda @carmocca @justusschock @awaelchli

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

…nto use-orig-params

github-actions · 2023-04-11T04:28:13Z

⚡ 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/fabric.py, src/lightning/fabric/strategies/fsdp.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

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

🟢 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/fabric.py, src/lightning/fabric/strategies/fsdp.py, tests/tests_fabric/helpers/models.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/fabric.py, src/lightning/fabric/strategies/fsdp.py, tests/tests_fabric/helpers/models.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/fabric.py, src/lightning/fabric/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/fabric.py, src/lightning/fabric/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.

codecov · 2023-04-11T04:53:16Z

Codecov Report

Merging #17305 (349dcab) into master (fd4697c) will decrease coverage by 9%.
The diff coverage is 88%.

Additional details and impacted files

@@            Coverage Diff            @@
##           master   #17305     +/-   ##
=========================================
- Coverage      68%      59%     -9%     
=========================================
  Files         435      410     -25     
  Lines       31658    31376    -282     
=========================================
- Hits        21650    18569   -3081     
- Misses      10008    12807   +2799

…nto use-orig-params

awaelchli added 3 commits April 7, 2023 21:09

use-orig-params

d0e3bb5

use-orig-params

1f087a9

testing

a05a3cd

awaelchli added feature Is an improvement or enhancement fabric lightning.fabric.Fabric labels Apr 8, 2023

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

69a2e77

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

awaelchli added this to the 2.1 milestone Apr 8, 2023

awaelchli added 3 commits April 8, 2023 11:36

testing

e662c71

changelog

86a43de

Merge branch 'use-orig-params' of github.com:Lightning-AI/lightning i…

7b11a12

…nto use-orig-params

awaelchli mentioned this pull request Apr 9, 2023

Simplified setup of optimizers in FSDP #17309

Merged

awaelchli added the fun Staff contributions outside working hours - to differentiate from the "community" label label Apr 9, 2023

update

1301158

awaelchli marked this pull request as ready for review April 11, 2023 04:27

awaelchli requested review from carmocca and justusschock as code owners April 11, 2023 04:27

Borda approved these changes Apr 11, 2023

View reviewed changes

awaelchli added 3 commits April 11, 2023 08:22

update tests

9ce51aa

Merge branch 'use-orig-params' of github.com:Lightning-AI/lightning i…

17330c1

…nto use-orig-params

update tests

7edae3e

justusschock approved these changes Apr 11, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Apr 11, 2023

udpate import

e226b6d

carmocca approved these changes Apr 11, 2023

View reviewed changes

awaelchli enabled auto-merge (squash) April 11, 2023 13:56

awaelchli added 2 commits April 11, 2023 18:52

skip test that requires rework

b456c8d

fix test

349dcab

awaelchli merged commit 51697a8 into master Apr 11, 2023

awaelchli deleted the use-orig-params branch April 11, 2023 19:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Combined setup of model and optimizer with FSDP #17305

Combined setup of model and optimizer with FSDP #17305

awaelchli commented Apr 8, 2023 •

edited

Loading

github-actions bot commented Apr 11, 2023 •

edited

Loading

codecov bot commented Apr 11, 2023 •

edited

Loading

Combined setup of model and optimizer with FSDP #17305

Combined setup of model and optimizer with FSDP #17305

Conversation

awaelchli commented Apr 8, 2023 • edited Loading

What does this PR do?

github-actions bot commented Apr 11, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Apr 11, 2023 • edited Loading

Codecov Report

awaelchli commented Apr 8, 2023 •

edited

Loading

github-actions bot commented Apr 11, 2023 •

edited

Loading

codecov bot commented Apr 11, 2023 •

edited

Loading