Enable passing `load_state_dict(..., assign=True|False)` in FabricModule #18690

awaelchli · 2023-10-02T16:40:08Z

What does this PR do?

Allows you to pass the new arguments assign to the load_state_dict() of a Fabric-wrapped module. Previously, this would fail with this error

TypeError: _FabricModule.load_state_dict() got an unexpected keyword argument 'assign'

since our signature doesn't accept this argument.

📚 Documentation preview 📚: https://pytorch-lightning--18690.org.readthedocs.build/en/18690/

cc @Borda @carmocca @justusschock @awaelchli

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

…-dict' into bugfix/fabric-wrapper-load-state-dict

src/lightning/fabric/wrappers.py

github-actions · 2023-10-02T16:47:42Z

⚡ 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/wrappers.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=177305&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=177305&view=logs&jobId=3f274fac-2e11-54ca-487e-194c91f3ae9f)	success

These checks are required after the changes to src/lightning/fabric/wrappers.py.

🟢 pytorch_lightning: Benchmarks

Check ID	Status
lightning.Benchmarks	success	✅

These checks are required after the changes to src/lightning/fabric/wrappers.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/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	✅
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/wrappers.py, tests/tests_fabric/test_wrappers.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=177307&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=177307&view=logs&jobId=47e66f3c-897a-5428-da11-bf5c7745762e)	success

These checks are required after the changes to src/lightning/fabric/wrappers.py, tests/tests_fabric/test_wrappers.py.

🟢 mypy

Check ID	Status
mypy	success	✅

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

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

update

3764a97

github-actions bot added the fabric lightning.fabric.Fabric label Oct 2, 2023

awaelchli added bug Something isn't working feature Is an improvement or enhancement and removed bug Something isn't working labels Oct 2, 2023

awaelchli added this to the 2.1 milestone Oct 2, 2023

pre-commit-ci bot and others added 4 commits October 2, 2023 16:41

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

3437d5a

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

test

081cc86

Merge remote-tracking branch 'origin/bugfix/fabric-wrapper-load-state…

b8d7c68

…-dict' into bugfix/fabric-wrapper-load-state-dict

chlog

1310551

awaelchli commented Oct 2, 2023

View reviewed changes

src/lightning/fabric/wrappers.py Outdated Show resolved Hide resolved

awaelchli marked this pull request as ready for review October 2, 2023 16:47

awaelchli requested review from carmocca and justusschock as code owners October 2, 2023 16:47

carmocca approved these changes Oct 2, 2023

View reviewed changes

Borda approved these changes Oct 2, 2023

View reviewed changes

mergify bot added the ready PRs ready to be merged label Oct 2, 2023

awaelchli and others added 4 commits October 2, 2023 11:25

Merge branch 'master' into bugfix/fabric-wrapper-load-state-dict

f3fe145

Merge branch 'master' into bugfix/fabric-wrapper-load-state-dict

8dd8e84

ignore override

8bc7b2a

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

bbc19c5

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

awaelchli merged commit 256f16e into master Oct 3, 2023
110 of 113 checks passed

awaelchli deleted the bugfix/fabric-wrapper-load-state-dict branch October 3, 2023 17:49

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable passing `load_state_dict(..., assign=True|False)` in FabricModule #18690

Enable passing `load_state_dict(..., assign=True|False)` in FabricModule #18690

awaelchli commented Oct 2, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Oct 2, 2023 •

edited

Loading

Enable passing load_state_dict(..., assign=True|False) in FabricModule #18690

Enable passing load_state_dict(..., assign=True|False) in FabricModule #18690

Conversation

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

What does this PR do?

github-actions bot commented Oct 2, 2023 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

Enable passing `load_state_dict(..., assign=True|False)` in FabricModule #18690

Enable passing `load_state_dict(..., assign=True|False)` in FabricModule #18690

awaelchli commented Oct 2, 2023 •

edited by github-actions bot

Loading

github-actions bot commented Oct 2, 2023 •

edited

Loading