Enable loading full state dict checkpoints with FSDP #17623

awaelchli · 2023-05-13T04:03:58Z

What does this PR do?

The last missing piece to support both full-state and sharded checkpoint saving/loading.

Saving sharded state: Save and load sharded checkpoints with FSDP in Fabric #17323
Loading sharded state: Save and load sharded checkpoints with FSDP in Fabric #17323
Saving full-state: Add support for saving with full state-dict in Fabric's FSDP #17526
Loading full-state: This PR 🎉

cc @Borda @awaelchli @carmocca @justusschock

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

src/lightning/fabric/strategies/fsdp.py

github-actions · 2023-05-30T03:34:37Z

⚡ 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.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs)	success	✅

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

🟢 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.

🟢 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.

🟢 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.

🟢 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.

carmocca

Nice job!

src/lightning/fabric/strategies/fsdp.py

Co-authored-by: Carlos Mocholí <[email protected]>

awaelchli added 2 commits May 11, 2023 19:40

load full state dict with fsdp

3a84dc5

update

5cf586a

github-actions bot added the fabric lightning.fabric.Fabric label May 13, 2023

awaelchli added feature Is an improvement or enhancement strategy: fsdp Fully Sharded Data Parallel checkpointing Related to checkpointing labels May 13, 2023

awaelchli added this to the 2.1 milestone May 13, 2023

pre-commit-ci bot and others added 4 commits May 13, 2023 04:05

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

f0c6e72

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

Merge branch 'master' into fabric/fsdp-load-full-state

308a513

wip test

9e7570d

Merge branch 'master' into fabric/fsdp-load-full-state

9c5fb8e

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

awaelchli added 2 commits May 29, 2023 23:04

update

a38c397

update

4d4aae0

awaelchli commented May 30, 2023

View reviewed changes

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

awaelchli commented May 30, 2023

View reviewed changes

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

awaelchli added 2 commits May 29, 2023 23:19

add test

b57fb3c

precommit

d1b69dd

awaelchli changed the title ~~WIP: Enable loading full state dict checkpoints with FSDP~~ Enable loading full state dict checkpoints with FSDP May 30, 2023

awaelchli added 5 commits May 29, 2023 23:25

extend test

a264242

update changelog

f07bae7

nit

40cf7b5

meta constant

e8c9c7a

docs

969048a

awaelchli marked this pull request as ready for review May 30, 2023 03:34

awaelchli requested review from carmocca and justusschock as code owners May 30, 2023 03:34

carmocca approved these changes May 30, 2023

View reviewed changes

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

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

run test with min torch

ef5c665

awaelchli and others added 2 commits May 30, 2023 06:13

move warning up

72ccf95

Update src/lightning/fabric/strategies/fsdp.py

e913e8e

Co-authored-by: Carlos Mocholí <[email protected]>

justusschock approved these changes May 31, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label May 31, 2023

awaelchli merged commit fd296e0 into master May 31, 2023

awaelchli deleted the fabric/fsdp-load-full-state branch May 31, 2023 15:30

leng-yue mentioned this pull request Jun 3, 2023

Enable loading full optimizer checkpoints with FSDP #17747

Merged

10 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable loading full state dict checkpoints with FSDP #17623

Enable loading full state dict checkpoints with FSDP #17623

awaelchli commented May 13, 2023 •

edited by github-actions bot

Loading

github-actions bot commented May 30, 2023 •

edited

Loading

carmocca left a comment

Enable loading full state dict checkpoints with FSDP #17623

Enable loading full state dict checkpoints with FSDP #17623

Conversation

awaelchli commented May 13, 2023 • edited by github-actions bot Loading

What does this PR do?

github-actions bot commented May 30, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

carmocca left a comment

Choose a reason for hiding this comment

awaelchli commented May 13, 2023 •

edited by github-actions bot

Loading

github-actions bot commented May 30, 2023 •

edited

Loading