-
Notifications
You must be signed in to change notification settings - Fork 423
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Fix fsdp checkpoint strategy #1734
Merged
Merged
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
abhi-mosaic
approved these changes
Nov 18, 2022
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you for the quick fix! I think there is some duplicate code to cleanup, but approving now to unblock.
bmosaicml
pushed a commit
to bmosaicml/composer
that referenced
this pull request
Dec 13, 2022
* Adding in performance registration * clearing optimizer param groups before FSDP wrapping * fixing comments * Adding a conditional to ignore checkpointing FSDP and FlattenedWrappers * removing print comments * Adding in print for my own sanity * removing extraneous print statement * removing duplicate code due to a merge Co-authored-by: Brandon Cui <bcui@Brandons-MBP.hsd1.ca.comcast.net>
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?
This PR allows us to now use method 2 for activation checkpointing. Basically FSDP recursively dives down to see if there are modules that should have activation checkpointing. However, we shouldn't be trying to wrap 'FullyShardedDataParallel' or 'FlattenedParamsWrappers' modules with the 'CheckpointWrapper'. By trying to wrap these two methods, it leads to a strange assert
Ran with 7b param gpt-2 style models:
run with method 1: https://wandb.ai/mosaic-ml/fsdp-activation-checkpointing/runs/17oacy01
run with method 2: https://wandb.ai/mosaic-ml/fsdp-activation-checkpointing/runs/1f1ij64d
What issue(s) does this change relate to?
Fixing CO-1377
Before submitting
pre-commit
on your change? (see thepre-commit
section of prerequisites)