Add support for saving with full state-dict in Fabric's FSDP #17526

awaelchli · 2023-04-29T03:13:16Z

What does this PR do?

Adds support for selecting the mode in which a checkpoint gets saved with the FSDPStrategy.

# saves a single file
fabric = Fabric(strategy=FSDPStrategy(state_dict_type="full"), ...)
fabric.save("path/to/checkpoint.pt", ...)

# saves a directory with one file per process
fabric = Fabric(strategy=FSDPStrategy(state_dict_type="sharded"), ...)
fabric.save("path/to/checkpoint.pt", ...)

The full state dict is easier to work with since it can immediately be loaded into other scripts or models without needing to initialize FSDP. However, sharded state dict is more efficient (and the default) when dealing with very large models that may not fit into CPU memory.

A follow-up PR will add some logic to FSDPStrategy.load_checkpoint() to detect whether a sharded or full checkpoint should be loaded based on the file saved at the given file path.

cc @Borda @awaelchli @carmocca @justusschock

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

…ric/full-state-dict

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

…ric/full-state-dict

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

…ric/full-state-dict

github-actions · 2023-05-11T09:28:00Z

⚡ 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/fabric/wrappers.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, src/lightning/fabric/wrappers.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, src/lightning/fabric/wrappers.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, src/lightning/fabric/wrappers.py, tests/tests_fabric/strategies/test_fsdp.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/test_wrappers.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, src/lightning/fabric/wrappers.py, tests/tests_fabric/strategies/test_fsdp.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/test_wrappers.py.

🟢 mypy

Check ID	Status
mypy	success	✅

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

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

carmocca

Very nice

wip

e59ba1a

github-actions bot added the fabric lightning.fabric.Fabric label Apr 29, 2023

pre-commit-ci bot and others added 19 commits April 29, 2023 03:14

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

10406b8

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

Merge branch 'master' into fabric/full-state-dict

9a002f4

create dir as late as possible

6a45155

Merge branch 'master' into fabric/full-state-dict

a0ae4c3

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

6d80b9e

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

update

927952a

Merge remote-tracking branch 'origin/fabric/full-state-dict' into fab…

aa06465

…ric/full-state-dict

dummy test

4d92165

simple test

f19d419

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

4752df2

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

update test

10bf8ad

Merge remote-tracking branch 'origin/fabric/full-state-dict' into fab…

dfd08b6

…ric/full-state-dict

path

60dbfde

to cpu

018753d

single file saving

3fd94d5

fix

42eb09d

update docs and typing

95059c4

more assertions

3187c77

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

7376bdc

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

awaelchli changed the title ~~WIP: Add support for saving with full state-dict in Fabric's FSDP~~ Add support for saving with full state-dict in Fabric's FSDP May 11, 2023

awaelchli added feature Is an improvement or enhancement checkpointing Related to checkpointing labels May 11, 2023

awaelchli added this to the 2.1 milestone May 11, 2023

awaelchli added the strategy: fsdp Fully Sharded Data Parallel label May 11, 2023

awaelchli added 3 commits May 11, 2023 11:16

add changelog

b598889

Merge remote-tracking branch 'origin/fabric/full-state-dict' into fab…

4631c43

…ric/full-state-dict

update test

b64f4ce

awaelchli marked this pull request as ready for review May 11, 2023 09:27

awaelchli requested review from carmocca and justusschock as code owners May 11, 2023 09:27

justusschock approved these changes May 11, 2023

View reviewed changes

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

awaelchli and others added 2 commits May 11, 2023 14:05

update

d1f95ae

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

5aad769

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

Borda approved these changes May 11, 2023

View reviewed changes

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

carmocca approved these changes May 11, 2023

View reviewed changes

awaelchli merged commit c712ec1 into master May 11, 2023

awaelchli deleted the fabric/full-state-dict branch May 11, 2023 17:02

awaelchli mentioned this pull request May 13, 2023

Enable loading full state dict checkpoints with FSDP #17623

Merged

leng-yue mentioned this pull request May 23, 2023

optimizer tensor size mismatch when resuming from ckpt with different number of workers using fsdp #16815

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add support for saving with full state-dict in Fabric's FSDP #17526

Add support for saving with full state-dict in Fabric's FSDP #17526

awaelchli commented Apr 29, 2023 •

edited by github-actions bot

Loading

github-actions bot commented May 11, 2023 •

edited

Loading

carmocca left a comment

Add support for saving with full state-dict in Fabric's FSDP #17526

Add support for saving with full state-dict in Fabric's FSDP #17526

Conversation

awaelchli commented Apr 29, 2023 • edited by github-actions bot Loading

What does this PR do?

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

⚡ Required checks status: All passing 🟢

Groups summary

carmocca left a comment

Choose a reason for hiding this comment

awaelchli commented Apr 29, 2023 •

edited by github-actions bot

Loading

github-actions bot commented May 11, 2023 •

edited

Loading