Support saving and loading stateful objects in Fabric #18513

awaelchli · 2023-09-10T00:14:04Z

What does this PR do?

Sofar we've handled state_dict() and load_state_dict() automatically for the user on the module and optimizer(s) when saving and loading a state with Fabric. This PR extends this to any object in the state that is "stateful", for example learning rate schedulers:

state = {"model": ..., "lr_scheduler": my_torch_scheduler}

# auomatically saves my_torch_scheduler.state_dict() now
fabric.save(path, state)

# auomatically calls my_torch_scheduler.load_state_dict() now
fabric.load(path, state)

cc @Borda @justusschock @awaelchli @carmocca

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

src/lightning/fabric/strategies/strategy.py

github-actions · 2023-09-10T12:45:14Z

⚡ 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	✅
pl-cpu (macOS-12, pytorch, 3.11, 2.0)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.0)	success	✅
pl-cpu (windows-2022, pytorch, 3.11, 2.0)	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.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=174138&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=174138&view=logs&jobId=3f274fac-2e11-54ca-487e-194c91f3ae9f)	success

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.py.

🟢 fabric: Docs

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

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.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	✅
fabric-cpu (macOS-12, fabric, 3.11, 2.0)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.0)	success	✅
fabric-cpu (windows-2022, fabric, 3.11, 2.0)	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.py, tests/tests_fabric/strategies/test_deepspeed.py, tests/tests_fabric/strategies/test_strategy.py, tests/tests_fabric/utilities/test_load.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=174140&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=174140&view=logs&jobId=47e66f3c-897a-5428-da11-bf5c7745762e)	success

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.py, tests/tests_fabric/strategies/test_deepspeed.py, tests/tests_fabric/strategies/test_strategy.py, tests/tests_fabric/utilities/test_load.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.py.

🟢 install

Check ID	Status
install-pkg (ubuntu-22.04, app, 3.8)	success	✅
install-pkg (ubuntu-22.04, app, 3.11)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.8)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.11)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.8)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.11)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.8)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.11)	success	✅
install-pkg (ubuntu-22.04, notset, 3.8)	success	✅
install-pkg (ubuntu-22.04, notset, 3.11)	success	✅
install-pkg (macOS-12, app, 3.8)	success	✅
install-pkg (macOS-12, app, 3.11)	success	✅
install-pkg (macOS-12, fabric, 3.8)	success	✅
install-pkg (macOS-12, fabric, 3.11)	success	✅
install-pkg (macOS-12, pytorch, 3.8)	success	✅
install-pkg (macOS-12, pytorch, 3.11)	success	✅
install-pkg (macOS-12, lightning, 3.8)	success	✅
install-pkg (macOS-12, lightning, 3.11)	success	✅
install-pkg (macOS-12, notset, 3.8)	success	✅
install-pkg (macOS-12, notset, 3.11)	success	✅
install-pkg (windows-2022, app, 3.8)	success	✅
install-pkg (windows-2022, app, 3.11)	success	✅
install-pkg (windows-2022, fabric, 3.8)	success	✅
install-pkg (windows-2022, fabric, 3.11)	success	✅
install-pkg (windows-2022, pytorch, 3.8)	success	✅
install-pkg (windows-2022, pytorch, 3.11)	success	✅
install-pkg (windows-2022, lightning, 3.8)	success	✅
install-pkg (windows-2022, lightning, 3.11)	success	✅
install-pkg (windows-2022, notset, 3.8)	success	✅
install-pkg (windows-2022, notset, 3.11)	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/deepspeed.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/strategies/strategy.py, src/lightning/fabric/utilities/load.py.

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/utilities/load.py

tests/tests_fabric/strategies/test_deepspeed_integration.py

src/lightning/fabric/utilities/load.py

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

tests/tests_fabric/strategies/test_deepspeed_integration.py

…fabric/stateful-save-load

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

src/lightning/fabric/strategies/deepspeed.py

Co-authored-by: Jirka Borovec <[email protected]>

save and load stateful objects

a181bde

github-actions bot added the fabric lightning.fabric.Fabric label Sep 10, 2023

awaelchli added fun Staff contributions outside working hours - to differentiate from the "community" label feature Is an improvement or enhancement checkpointing Related to checkpointing refactor labels Sep 10, 2023

awaelchli added this to the 2.1 milestone Sep 10, 2023

pre-commit-ci bot and others added 9 commits September 10, 2023 00:15

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

fceff40

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

fix text

63e98e3

mini refactor

8348b3d

chlog

62660eb

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

252cde3

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

fix test

8b8c9d9

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

43885bb

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

add comment

4ea75da

deepspeed flaky test

f83b293

awaelchli commented Sep 10, 2023

View reviewed changes

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

rename

f28348b

awaelchli marked this pull request as ready for review September 10, 2023 12:44

awaelchli requested review from carmocca and justusschock as code owners September 10, 2023 12:44

carmocca approved these changes Sep 10, 2023

View reviewed changes

awaelchli and others added 4 commits September 10, 2023 12:56

Update src/lightning/fabric/utilities/load.py

82bba02

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

avoid .keys() in sets

c86bcf1

review

40f9b4b

review

e4b580f

awaelchli commented Sep 10, 2023

View reviewed changes

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

awaelchli added 2 commits September 10, 2023 16:47

make test flaky

62bc0d6

rename

3434c99

awaelchli and others added 2 commits September 10, 2023 22:54

Merge remote-tracking branch 'origin/fabric/stateful-save-load' into …

73acb70

…fabric/stateful-save-load

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

77cd46d

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

Borda approved these changes Sep 12, 2023

View reviewed changes

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

mergify bot added the ready PRs ready to be merged label Sep 12, 2023

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

5949f65

Co-authored-by: Jirka Borovec <[email protected]>

awaelchli merged commit c959df7 into master Sep 12, 2023

awaelchli deleted the fabric/stateful-save-load branch September 12, 2023 11:58

awaelchli mentioned this pull request Oct 3, 2023

Pickle error when saving the lr scheduler defined after fabric.setup or by LambdaLR in Fabric #18688

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support saving and loading stateful objects in Fabric #18513

Support saving and loading stateful objects in Fabric #18513

awaelchli commented Sep 10, 2023 •

edited

Loading

github-actions bot commented Sep 10, 2023 •

edited

Loading

Support saving and loading stateful objects in Fabric #18513

Support saving and loading stateful objects in Fabric #18513

Conversation

awaelchli commented Sep 10, 2023 • edited Loading

What does this PR do?

github-actions bot commented Sep 10, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

awaelchli commented Sep 10, 2023 •

edited

Loading

github-actions bot commented Sep 10, 2023 •

edited

Loading