Support true 16-bit precision with FSDP in Trainer #18219

awaelchli · 2023-08-02T23:57:37Z

What does this PR do?

Follow up to #17609

Removes inheritance of MixedPrecisionPlugin from the FSDP precision plugin
The plugin now handles both the full-precision, mixed-precision, and true half precision.

Changes are similar to #17807

cc @Borda @justusschock @awaelchli @carmocca

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

… into feature/fsdp-half

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

github-actions · 2023-08-09T15:17:09Z

⚡ 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/connector.py, src/lightning/fabric/plugins/precision/fsdp.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/amp.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/deprecated_api/test_no_removal_version.py, tests/tests_pytorch/plugins/precision/test_fsdp.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
[pytorch-lightning (GPUs) (testing Lightning	latest)](https://dev.azure.com/Lightning-AI/72ab7ed8-b00f-4b6e-b131-3388f7ffafa7/_build/results?buildId=169246&view=logs&jobId=47e66f3c-897a-5428-da11-bf5c7745762e)	success
[pytorch-lightning (GPUs) (testing PyTorch	latest)](https://dev.azure.com/Lightning-AI/72ab7ed8-b00f-4b6e-b131-3388f7ffafa7/_build/results?buildId=169246&view=logs&jobId=3f274fac-2e11-54ca-487e-194c91f3ae9f)	success

These checks are required after the changes to src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/amp.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, tests/tests_pytorch/deprecated_api/test_no_removal_version.py, tests/tests_pytorch/plugins/precision/test_fsdp.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/trainer/connectors/test_accelerator_connector.py, src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/fsdp.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/fsdp.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/amp.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.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/connector.py, src/lightning/fabric/plugins/precision/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/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/amp.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py, docs/source-pytorch/api_references.rst, docs/source-pytorch/extensions/plugins.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/connector.py, src/lightning/fabric/plugins/precision/fsdp.py, tests/tests_fabric/plugins/precision/test_fsdp.py, tests/tests_fabric/test_connector.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
[lightning-fabric (GPUs) (testing Fabric	latest)](https://dev.azure.com/Lightning-AI/72ab7ed8-b00f-4b6e-b131-3388f7ffafa7/_build/results?buildId=169248&view=logs&jobId=3f274fac-2e11-54ca-487e-194c91f3ae9f)	success
[lightning-fabric (GPUs) (testing Lightning	latest)](https://dev.azure.com/Lightning-AI/72ab7ed8-b00f-4b6e-b131-3388f7ffafa7/_build/results?buildId=169248&view=logs&jobId=47e66f3c-897a-5428-da11-bf5c7745762e)	success

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/fsdp.py, tests/tests_fabric/plugins/precision/test_fsdp.py, tests/tests_fabric/test_connector.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/connector.py, src/lightning/fabric/plugins/precision/fsdp.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/amp.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/fsdp.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.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/connector.py, src/lightning/fabric/plugins/precision/fsdp.py, src/lightning/pytorch/plugins/__init__.py, src/lightning/pytorch/plugins/precision/__init__.py, src/lightning/pytorch/plugins/precision/amp.py, src/lightning/pytorch/plugins/precision/fsdp.py, src/lightning/pytorch/strategies/fsdp.py, src/lightning/pytorch/trainer/connectors/accelerator_connector.py.

🟢 link-check

Check ID	Status
check-md-links / markdown-link-check	success	✅

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

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

tests/tests_pytorch/strategies/test_fsdp.py

src/lightning/pytorch/plugins/precision/fsdp.py

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

src/lightning/pytorch/plugins/precision/fsdp.py

wip

7653c61

awaelchli added feature Is an improvement or enhancement precision: amp Automatic Mixed Precision strategy: fsdp Fully Sharded Data Parallel labels Aug 2, 2023

github-actions bot added fabric lightning.fabric.Fabric pl Generic label for PyTorch Lightning package labels Aug 2, 2023

awaelchli added this to the 2.1 milestone Aug 2, 2023

awaelchli added the refactor label Aug 2, 2023

awaelchli changed the title ~~Support all half-precision modes in FSDP precision plugin for Trainer~~ Support true 16-bit precision with FSDP in Trainer #18217 Aug 2, 2023

awaelchli changed the title ~~Support true 16-bit precision with FSDP in Trainer #18217~~ Support true 16-bit precision with FSDP in Trainer Aug 2, 2023

pre-commit-ci bot and others added 6 commits August 2, 2023 23:59

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

cd83eff

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

changelog

655023c

update tests

bf192ac

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

48a3bd2

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

fix deprecation message

20f9d6b

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

abaa249

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

awaelchli mentioned this pull request Aug 3, 2023

True half-precision support in Trainer #18193

Merged

7 tasks

awaelchli and others added 8 commits August 3, 2023 15:55

fix imports

33c19ac

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

8ce3463

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

Merge branch 'master' into feature/fsdp-half

7f997f8

fix

ff22615

Merge branch 'feature/fsdp-half' of github.com:Lightning-AI/lightning…

128115c

… into feature/fsdp-half

Merge branch 'master' into feature/fsdp-half

5923ef0

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

a1a9826

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

fix tests

469d8b1

awaelchli marked this pull request as ready for review August 9, 2023 15:16

awaelchli requested review from carmocca, justusschock, Borda and williamFalcon as code owners August 9, 2023 15:16

awaelchli requested review from edenlightning and lantiga as code owners August 9, 2023 15:16

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

14b9a78

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

awaelchli commented Aug 9, 2023

View reviewed changes

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

awaelchli commented Aug 9, 2023

View reviewed changes

src/lightning/pytorch/plugins/precision/fsdp.py Outdated Show resolved Hide resolved

carmocca approved these changes Aug 9, 2023

View reviewed changes

awaelchli and others added 5 commits August 9, 2023 17:39

mypy fixes

fe9d91e

deprecation

e07d3dd

mypy ignore, torch 2.0.1 has wrong type

9c6f332

mypy

0984798

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

3653b94

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

justusschock approved these changes Aug 10, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Aug 10, 2023

Borda approved these changes Aug 10, 2023

View reviewed changes

src/lightning/pytorch/plugins/precision/fsdp.py Show resolved Hide resolved

awaelchli merged commit 888466b into master Aug 10, 2023
103 checks passed

awaelchli deleted the feature/fsdp-half branch August 10, 2023 08:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support true 16-bit precision with FSDP in Trainer #18219

Support true 16-bit precision with FSDP in Trainer #18219

awaelchli commented Aug 2, 2023 •

edited by github-actions bot

github-actions bot commented Aug 9, 2023 •

edited

Support true 16-bit precision with FSDP in Trainer #18219

Support true 16-bit precision with FSDP in Trainer #18219

Conversation

awaelchli commented Aug 2, 2023 • edited by github-actions bot

What does this PR do?

github-actions bot commented Aug 9, 2023 • edited

⚡ Required checks status: All passing 🟢

Groups summary

awaelchli commented Aug 2, 2023 •

edited by github-actions bot

github-actions bot commented Aug 9, 2023 •

edited