Utility to consolidate sharded checkpoints #19213

awaelchli · 2023-12-26T22:40:23Z

What does this PR do?

Adds a utility function and CLI to consolidate sharded checkpoints saved with FSDP. See the docs page added here for an explanation of how to use it.

The name for the CLI option needs to be decided.

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

cc @Borda @awaelchli @carmocca @justusschock

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

…ightning-AI/lightning into feature/consolidate-sharded-checkpoint

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

github-actions · 2024-01-21T23:12:19Z

⚡ Required checks status: All passing 🟢

Groups summary

🟢 pytorch_lightning: Tests workflow

Check ID	Status
pl-cpu (macOS-11, lightning, 3.8, 1.12, oldest)	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.10, 2.1)	success	✅
pl-cpu (ubuntu-20.04, lightning, 3.8, 1.12, oldest)	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.10, 2.1)	success	✅
pl-cpu (windows-2022, lightning, 3.8, 1.12, oldest)	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.10, 2.1)	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 (macOS-12, pytorch, 3.11, 2.1)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.0)	success	✅
pl-cpu (ubuntu-22.04, pytorch, 3.11, 2.1)	success	✅
pl-cpu (windows-2022, pytorch, 3.11, 2.0)	success	✅
pl-cpu (windows-2022, pytorch, 3.11, 2.1)	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, src/lightning/pytorch/utilities/consolidate_checkpoint.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/utilities/test_consolidate_checkpoint.py.

🟢 pytorch_lightning: Azure GPU

Check ID	Status
pytorch-lightning (GPUs) (testing Lightning \| latest)	success	✅
pytorch-lightning (GPUs) (testing PyTorch \| latest)	success	✅

These checks are required after the changes to src/lightning/pytorch/utilities/consolidate_checkpoint.py, tests/tests_pytorch/strategies/test_fsdp.py, tests/tests_pytorch/utilities/test_consolidate_checkpoint.py, src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.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/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, src/lightning/pytorch/utilities/consolidate_checkpoint.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/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, docs/source-fabric/guide/checkpoint/distributed_checkpoint.rst.

🟢 pytorch_lightning: Docs

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

These checks are required after the changes to src/lightning/pytorch/utilities/consolidate_checkpoint.py, docs/source-pytorch/common/checkpointing_expert.rst.

🟢 lightning_fabric: CPU workflow

Check ID	Status
fabric-cpu (macOS-11, lightning, 3.8, 1.12, oldest)	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.11, 2.1)	success	✅
fabric-cpu (ubuntu-20.04, lightning, 3.8, 1.12, oldest)	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.11, 2.1)	success	✅
fabric-cpu (windows-2022, lightning, 3.8, 1.12, oldest)	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.11, 2.1)	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 (macOS-12, fabric, 3.11, 2.1)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.0)	success	✅
fabric-cpu (ubuntu-22.04, fabric, 3.11, 2.1)	success	✅
fabric-cpu (windows-2022, fabric, 3.11, 2.0)	success	✅
fabric-cpu (windows-2022, fabric, 3.11, 2.1)	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/utilities/test_consolidate_checkpoint.py, tests/tests_fabric/utilities/test_load.py.

🟢 lightning_fabric: Azure GPU

Check ID	Status
lightning-fabric (GPUs) (testing Fabric \| latest)	success	✅
lightning-fabric (GPUs) (testing Lightning \| latest)	success	✅

These checks are required after the changes to src/lightning/fabric/strategies/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, tests/tests_fabric/strategies/test_fsdp_integration.py, tests/tests_fabric/utilities/test_consolidate_checkpoint.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/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, src/lightning/pytorch/utilities/consolidate_checkpoint.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/fsdp.py, src/lightning/fabric/utilities/consolidate_checkpoint.py, src/lightning/fabric/utilities/load.py, src/lightning/pytorch/utilities/consolidate_checkpoint.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.

codecov · 2024-01-21T23:32:47Z

Codecov Report

Merging #19213 (7b97057) into master (ed367ca) will decrease coverage by 29%.
The diff coverage is 89%.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #19213      +/-   ##
==========================================
- Coverage      83%      54%     -29%     
==========================================
  Files         446      443       -3     
  Lines       37702    37697       -5     
==========================================
- Hits        31242    20282   -10960     
- Misses       6460    17415   +10955

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

…ightning-AI/lightning into feature/consolidate-sharded-checkpoint

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

docs/source-pytorch/common/checkpointing_expert.rst

src/lightning/fabric/utilities/load.py

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

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

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

awaelchli added 2 commits December 26, 2023 17:22

add utility

4005b46

update

1145247

github-actions bot added the fabric lightning.fabric.Fabric label Dec 26, 2023

todos

b5bbd20

awaelchli force-pushed the feature/consolidate-sharded-checkpoint branch from 5cb990a to b5bbd20 Compare December 26, 2023 22:52

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

5506767

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

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

awaelchli added this to the 2.2 milestone Dec 26, 2023

awaelchli added 4 commits December 26, 2023 18:03

todo

8d770ef

refactor

023bd89

conversion

c5376f2

handle extra

670039f

github-actions bot added the pl Generic label for PyTorch Lightning package label Dec 28, 2023

pre-commit-ci bot and others added 9 commits December 28, 2023 23:18

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

76dab90

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

fix

b08608e

Merge branch 'feature/consolidate-sharded-checkpoint' of github.com:L…

660be54

…ightning-AI/lightning into feature/consolidate-sharded-checkpoint

reset

73d8a31

handle errors

8fbcd7d

handle flattened keys

61ed213

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

0c8ae6a

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

move

7a7a5d8

tests

9ed053a

awaelchli force-pushed the feature/consolidate-sharded-checkpoint branch from d62cf6a to 9ed053a Compare January 2, 2024 09:57

pre-commit-ci bot and others added 5 commits January 2, 2024 09:58

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

b463021

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

update

a3b0eec

merge script

a167727

cli

d425c6f

types

17580a8

awaelchli added 3 commits January 21, 2024 23:52

update docs

3d9e8c7

update test

7ad3705

Merge branch 'master' into feature/consolidate-sharded-checkpoint

85d0fe0

awaelchli marked this pull request as ready for review January 21, 2024 23:11

awaelchli requested review from carmocca, justusschock, Borda, williamFalcon and lantiga as code owners January 21, 2024 23:11

update monkeypatch

7af24fe

awaelchli and others added 5 commits January 22, 2024 05:19

windows

b8d9952

rename

64e983b

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

e92252b

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

rename

3342c0b

Merge branch 'feature/consolidate-sharded-checkpoint' of github.com:L…

4649b8e

…ightning-AI/lightning into feature/consolidate-sharded-checkpoint

carmocca approved these changes Jan 22, 2024

View reviewed changes

awaelchli and others added 4 commits January 22, 2024 15:54

Update docs/source-fabric/guide/checkpoint/distributed_checkpoint.rst

f5496aa

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

postfix

c2ac631

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

e2f5ec1

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

add the pytorch guide

331f222

Borda approved these changes Jan 23, 2024

View reviewed changes

mergify bot added the ready PRs ready to be merged label Jan 23, 2024

awaelchli and others added 3 commits January 23, 2024 22:03

validate that the checkpoint folder is an fsdp checkpoint

db5ecf7

Merge branch 'master' into feature/consolidate-sharded-checkpoint

a191835

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

7b97057

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

awaelchli merged commit b1127e3 into master Jan 23, 2024
118 checks passed

awaelchli deleted the feature/consolidate-sharded-checkpoint branch January 23, 2024 22:15

awaelchli mentioned this pull request Jan 26, 2024

Support consolidating sharded checkpoints in PyTorch >= 2.2 #19351

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Utility to consolidate sharded checkpoints #19213

Utility to consolidate sharded checkpoints #19213

awaelchli commented Dec 26, 2023 •

edited

Loading

github-actions bot commented Jan 21, 2024 •

edited

Loading

codecov bot commented Jan 21, 2024 •

edited

Loading

Utility to consolidate sharded checkpoints #19213

Utility to consolidate sharded checkpoints #19213

Conversation

awaelchli commented Dec 26, 2023 • edited Loading

What does this PR do?

github-actions bot commented Jan 21, 2024 • edited Loading

⚡ Required checks status: All passing 🟢

Groups summary

codecov bot commented Jan 21, 2024 • edited Loading

Codecov Report

awaelchli commented Dec 26, 2023 •

edited

Loading

github-actions bot commented Jan 21, 2024 •

edited

Loading

codecov bot commented Jan 21, 2024 •

edited

Loading