Support loading distributed checkpoints for FSDP in Trainer #18358
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
awaelchli
changed the title
for more information, see https://pre-commit.ci
…to feature/dist-checkpoint-fsdp
awaelchli
force-pushed
the
feature/dist-checkpoint-fsdp
branch
from
98b2528 to
1cdf15b
Compare
awaelchli
added
feature
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
…I/lightning into feature/dist-checkpoint-fsdp
for more information, see https://pre-commit.ci
…to feature/dist-checkpoint-fsdp-debug
awaelchli
requested review from
carmocca,
justusschock,
Borda and
williamFalcon
as code owners
awaelchli
changed the title
carmocca
approved these changes
Aug 23, 2023
Co-authored-by: Carlos Mocholí <[email protected]>
Borda
approved these changes
Aug 23, 2023
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Part 2 (#18364) completes the support for sharded / distributed checkpointing with FSDP in Trainer. The FSDP user guide will be updated in a follow-up PR.
There are minor breaking changes (when using FSDP only):
configure_sharded_model(), that is, after the model is sharded and put on the GPU. This is to avoid expensively loading the checkpoint in CPU memory before sharding (the model may not fit!). The other strategy that also does this is DeepSpeed. This change is necessary now to support sharded checkpoints, but will also be needed for supporting meta device later (also for full-state checkpoints).FSDPStrategy.load_optimizer_state_dictmethod becomes a no-op.FSDPStrategy.load_model_state_dictmethod becomes a no-op.PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:
Reviewer checklist
cc @Borda @awaelchli @carmocca