Fabric checkpointing 1/n: base implementation #16434

awaelchli · 2023-01-19T12:12:17Z

What does this PR do?

Adds the base implementation for checkpoint saving and loading, with the new API. Details of how it is used are in the docs pages of this PR.
Follow up PRs will add the implementation in the DeepSpeed and FSDP strategies.

Part of #14816

# Saving
state = {
    "model1": model1,
    "model2": model2,
    "optimizer": optimizer,
    "iteration": iteration,
}
fabric.save("./checkpoints/version_2/checkpoint.ckpt", state)

# Loading (selective, in-place)
state = {"model1": model1}
fabric.load("./checkpoints/version_2/checkpoint.ckpt", state)

# Loading (all)
all_states = fabric.load("./checkpoints/version_2/checkpoint.ckpt")

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you make sure to update the documentation with your changes? (if necessary)
Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

I made sure I had fun coding 🙃

cc @Borda @awaelchli @carmocca @justusschock

fix types api docs

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

…ite/dist-checkpoint-api

src/lightning_fabric/strategies/strategy.py

docs/source-pytorch/fabric/guide/checkpoint.rst

src/lightning_fabric/strategies/strategy.py

github-actions · 2023-01-19T13:40:29Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, pytorch, 3.8, 1.11)	success	✅
pl-cpu (macOS-11, pytorch, 3.9, 1.12)	success	✅
pl-cpu (macOS-11, pytorch, 3.10, 1.13)	success	✅
pl-cpu (macOS-11, pytorch, 3.8, 1.10, oldest)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.8, 1.10)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.9, 1.11)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.10, 1.12)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.10, 1.13)	success	✅
pl-cpu (ubuntu-20.04, pytorch, 3.7, 1.10, oldest)	success	✅
pl-cpu (windows-2022, pytorch, 3.9, 1.11)	success	✅
pl-cpu (windows-2022, pytorch, 3.10, 1.12)	success	✅
pl-cpu (windows-2022, pytorch, 3.10, 1.13)	success	✅
pl-cpu (windows-2022, pytorch, 3.7, 1.10, oldest)	success	✅
pl-cpu (slow, macOS-11, pytorch, 3.7, 1.11)	success	✅
pl-cpu (slow, ubuntu-20.04, pytorch, 3.7, 1.11)	success	✅
pl-cpu (slow, windows-2022, pytorch, 3.7, 1.11)	success	✅
pl-cpu (macOS-11, lightning, 3.8, 1.13)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.13)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.13)	success	✅

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

🟢 pytorch_lightning: Azure HPU

Check ID	Status
pytorch-lightning (HPUs)	success	✅

These checks are required after the changes to src/lightning_fabric/fabric.py, src/lightning_fabric/strategies/deepspeed.py, src/lightning_fabric/strategies/strategy.py, src/lightning_fabric/strategies/xla.py, src/lightning_fabric/wrappers.py.

🟢 pytorch_lightning: Azure IPU

Check ID	Status
pytorch-lightning (IPUs)	success	✅

These checks are required after the changes to src/lightning_fabric/fabric.py, src/lightning_fabric/strategies/deepspeed.py, src/lightning_fabric/strategies/strategy.py, src/lightning_fabric/strategies/xla.py, src/lightning_fabric/wrappers.py.

🟢 pytorch_lightning: Docs

Check ID	Status
make-doctest (pytorch)	success	✅
make-html (pytorch)	success	✅

These checks are required after the changes to docs/source-pytorch/fabric/api/fabric_methods.rst, docs/source-pytorch/fabric/fabric.rst, docs/source-pytorch/fabric/guide/checkpoint.rst.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, fabric, 3.8, 1.11)	success	✅
fabric-cpu (macOS-11, fabric, 3.9, 1.12)	success	✅
fabric-cpu (macOS-11, fabric, 3.10, 1.13)	success	✅
fabric-cpu (macOS-11, fabric, 3.7, 1.10, oldest)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.8, 1.10)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.9, 1.11)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.10, 1.12)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.10, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, fabric, 3.7, 1.10, oldest)	success	✅
fabric-cpu (windows-2022, fabric, 3.9, 1.11)	success	✅
fabric-cpu (windows-2022, fabric, 3.10, 1.12)	success	✅
fabric-cpu (windows-2022, fabric, 3.10, 1.13)	success	✅
fabric-cpu (windows-2022, fabric, 3.7, 1.10, oldest)	success	✅
fabric-cpu (macOS-11, lightning, 3.8, 1.13)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.13)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.13)	success	✅

These checks are required after the changes to src/lightning_fabric/fabric.py, src/lightning_fabric/strategies/deepspeed.py, src/lightning_fabric/strategies/strategy.py, src/lightning_fabric/strategies/xla.py, src/lightning_fabric/wrappers.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/strategies/test_strategy.py, tests/tests_fabric/test_fabric.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/deepspeed.py, src/lightning_fabric/strategies/strategy.py, src/lightning_fabric/strategies/xla.py, src/lightning_fabric/wrappers.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/strategies/test_strategy.py, tests/tests_fabric/test_fabric.py.

🟢 mypy

Check ID	Status
mypy	success	✅

These checks are required after the changes to src/lightning_fabric/fabric.py, src/lightning_fabric/strategies/deepspeed.py, src/lightning_fabric/strategies/strategy.py, src/lightning_fabric/strategies/xla.py, src/lightning_fabric/wrappers.py.

🟢 install

Check ID	Status
install-pkg (ubuntu-22.04, app, 3.7)	success	✅
install-pkg (ubuntu-22.04, app, 3.10)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.7)	success	✅
install-pkg (ubuntu-22.04, fabric, 3.10)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.7)	success	✅
install-pkg (ubuntu-22.04, pytorch, 3.10)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.7)	success	✅
install-pkg (ubuntu-22.04, lightning, 3.10)	success	✅
install-pkg (ubuntu-22.04, notset, 3.7)	success	✅
install-pkg (ubuntu-22.04, notset, 3.10)	success	✅
install-pkg (macOS-12, app, 3.7)	success	✅
install-pkg (macOS-12, app, 3.10)	success	✅
install-pkg (macOS-12, fabric, 3.7)	success	✅
install-pkg (macOS-12, fabric, 3.10)	success	✅
install-pkg (macOS-12, pytorch, 3.7)	success	✅
install-pkg (macOS-12, pytorch, 3.10)	success	✅
install-pkg (macOS-12, lightning, 3.7)	success	✅
install-pkg (macOS-12, lightning, 3.10)	success	✅
install-pkg (macOS-12, notset, 3.7)	success	✅
install-pkg (macOS-12, notset, 3.10)	success	✅
install-pkg (windows-2022, app, 3.7)	success	✅
install-pkg (windows-2022, app, 3.10)	success	✅
install-pkg (windows-2022, fabric, 3.7)	success	✅
install-pkg (windows-2022, fabric, 3.10)	success	✅
install-pkg (windows-2022, pytorch, 3.7)	success	✅
install-pkg (windows-2022, pytorch, 3.10)	success	✅
install-pkg (windows-2022, lightning, 3.7)	success	✅
install-pkg (windows-2022, lightning, 3.10)	success	✅
install-pkg (windows-2022, notset, 3.7)	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/deepspeed.py, src/lightning_fabric/strategies/strategy.py, src/lightning_fabric/strategies/xla.py, src/lightning_fabric/wrappers.py.

🟢 link-check

Check ID	Status
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/strategy.py

docs/source-pytorch/fabric/api/fabric_methods.rst

src/lightning_fabric/wrappers.py

src/lightning_fabric/strategies/strategy.py

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

src/lightning_fabric/strategies/strategy.py

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

src/lightning_fabric/fabric.py

awaelchli added this to the 2.0 milestone Jan 19, 2023

awaelchli self-assigned this Jan 19, 2023

awaelchli added checkpointing Related to checkpointing fabric lightning.fabric.Fabric labels Jan 19, 2023

github-actions bot added the pl Generic label for PyTorch Lightning package label Jan 19, 2023

awaelchli added 3 commits January 19, 2023 14:18

squash all commits

8a524a2

fix types api docs

checkpoint docs

e215033

changelog

46109fb

awaelchli force-pushed the lite/dist-checkpoint-api branch from aa41249 to 46109fb Compare January 19, 2023 13:19

pre-commit-ci bot and others added 4 commits January 19, 2023 13:23

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

60766e4

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

add backlinks

77bceed

Merge remote-tracking branch 'origin/lite/dist-checkpoint-api' into l…

1262595

…ite/dist-checkpoint-api

fix

056a73f

awaelchli commented Jan 19, 2023

View reviewed changes

src/lightning_fabric/strategies/strategy.py Show resolved Hide resolved

docs/source-pytorch/fabric/guide/checkpoint.rst Show resolved Hide resolved

src/lightning_fabric/strategies/strategy.py Show resolved Hide resolved

src/lightning_fabric/strategies/strategy.py Show resolved Hide resolved

awaelchli marked this pull request as ready for review January 19, 2023 13:39

awaelchli requested review from carmocca, justusschock, edenlightning, lantiga and Borda as code owners January 19, 2023 13:39

fix tests

b94f817

justusschock approved these changes Jan 19, 2023

View reviewed changes

src/lightning_fabric/strategies/strategy.py Show resolved Hide resolved

carmocca reviewed Jan 19, 2023

View reviewed changes

awaelchli and others added 2 commits January 19, 2023 12:02

Update src/lightning_fabric/wrappers.py

f0d975f

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

typing unwrap_collection

fe94b9e

carmocca approved these changes Jan 19, 2023

View reviewed changes

src/lightning_fabric/strategies/strategy.py Outdated Show resolved Hide resolved

mergify bot added the ready PRs ready to be merged label Jan 19, 2023

awaelchli added 2 commits January 19, 2023 18:32

clarify

1f3133e

add stateful type

a664bca

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

f55f259

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

awaelchli enabled auto-merge (squash) January 19, 2023 17:45

awaelchli added 2 commits January 19, 2023 18:54

make mock args check python 3.7 compatible

4273264

Merge branch 'master' into lite/dist-checkpoint-api

68016ff

Borda approved these changes Jan 19, 2023

View reviewed changes

src/lightning_fabric/fabric.py Show resolved Hide resolved

test for strict loading and error

cb23365

awaelchli merged commit 39acb81 into master Jan 19, 2023

awaelchli deleted the lite/dist-checkpoint-api branch January 19, 2023 20:40

This was referenced Jan 19, 2023

Make strict loading configurable in Fabric.load() #16447

Closed

Fabric checkpointing 2/n: DeepSpeed implementation #16452

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fabric checkpointing 1/n: base implementation #16434

Fabric checkpointing 1/n: base implementation #16434

awaelchli commented Jan 19, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jan 19, 2023 •

edited

Loading

Fabric checkpointing 1/n: base implementation #16434

Fabric checkpointing 1/n: base implementation #16434

Conversation

awaelchli commented Jan 19, 2023 • edited by github-actions bot Loading

What does this PR do?

Before submitting

PR review

Did you have fun?

github-actions bot commented Jan 19, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

awaelchli commented Jan 19, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Jan 19, 2023 •

edited

Loading