[FSDP] Allow different optim_input orders across ranks
#78599
Closed
Conversation
This file contains hidden or 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
[ghstack-poisoned]
awgu
requested review from
mrshenli,
pritamdamania87,
zhaojuanmao,
rohan-varma,
H-Huang and
mingzhe09088
as code owners
🔗 Helpful links
❌ 1 New FailuresAs of commit edef523 (more details on the Dr. CI page): Expand to see more
🕵️ 1 new failure recognized by patternsThe following CI failures do not appear to be due to upstream breakages
|
facebook-github-bot
added
the
oncall: distributed
Differential Revision: [D36798482](https://our.internmc.facebook.com/intern/diff/D36798482) [ghstack-poisoned]
awgu
changed the title
optim_input across ranksoptim_input orders across ranks
rohan-varma
reviewed
Jun 1, 2022
awgu
commented
Jun 1, 2022
rohan-varma
reviewed
Jun 1, 2022
This enables the first argument to the optimizer constructor, `optim_input`, to have different orders across ranks, e.g. if the parameters in one parameter group are permuted. This requires modification to `full_optim_state_dict()`, `shard_full_optim_state_dict()`, and `scatter_full_optim_state_dict()`. The high-level algorithmic change is that the state dicts are kept as being keyed by unflattened parameter name until after sharding/unsharding and flattening/unflattening and are rekeyed to be by parameter ID according to each rank's own `optim_input` only at the end. Differential Revision: [D36798482](https://our.internmc.facebook.com/intern/diff/D36798482) [ghstack-poisoned]
rohan-varma
approved these changes
Jun 1, 2022
| raise RuntimeError( | ||
| "FSDP currently requires each rank to have at least the " | ||
| "optimizer states needed by rank 0's optimizer but some ranks " | ||
| "are missing some of those states" |
There was a problem hiding this comment.
Will it also be useful to log the missing keys?
There was a problem hiding this comment.
I can add that at the cost of an all_gather_object().
There was a problem hiding this comment.
Update: The error now looks like:
RuntimeError: FSDP currently requires each rank to have at least the optimizer states needed by rank 0's optimizer but some ranks are missing some of those states
Rank 1 is missing states for the parameters: [('block2.2.weight', 'block2.2.bias_module0.bias', 'block2.2.bias_module1.bias')]
Rank 2 is missing states for the parameters: [('block2.2.weight', 'block2.2.bias_module0.bias', 'block2.2.bias_module1.bias')]
Rank 3 is missing states for the parameters: [('block2.2.weight', 'block2.2.bias_module0.bias', 'block2.2.bias_module1.bias')]
This enables the first argument to the optimizer constructor, `optim_input`, to have different orders across ranks, e.g. if the parameters in one parameter group are permuted. This requires modification to `full_optim_state_dict()`, `shard_full_optim_state_dict()`, and `scatter_full_optim_state_dict()`. The high-level algorithmic change is that the state dicts are kept as being keyed by unflattened parameter name until after sharding/unsharding and flattening/unflattening and are rekeyed to be by parameter ID according to each rank's own `optim_input` only at the end. Differential Revision: [D36798482](https://our.internmc.facebook.com/intern/diff/D36798482) [ghstack-poisoned]
This enables the first argument to the optimizer constructor, `optim_input`, to have different orders across ranks, e.g. if the parameters in one parameter group are permuted. This requires modification to `full_optim_state_dict()`, `shard_full_optim_state_dict()`, and `scatter_full_optim_state_dict()`. The high-level algorithmic change is that the state dicts are kept as being keyed by unflattened parameter name until after sharding/unsharding and flattening/unflattening and are rekeyed to be by parameter ID according to each rank's own `optim_input` only at the end. Differential Revision: [D36798482](https://our.internmc.facebook.com/intern/diff/D36798482) [ghstack-poisoned]
|
Backward compatibility error due to broken ops seems unrelated. |
|
@pytorchbot merge |
|
Hey @awgu. |
awgu
added
release notes: distributed (fsdp)
facebook-github-bot
pushed a commit
that referenced
this pull request
Jun 3, 2022
Summary: Pull Request resolved: #78599 Approved by: https://github.com/rohan-varma Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/4615738a3d5aee256a9f1929c846dcae7d20a041 Reviewed By: rohan-varma Differential Revision: D36798482 fbshipit-source-id: de05b37db8ed41b6cf11a9fc526a0e03a98f570d
zhaojuanmao
reviewed
Jun 5, 2022
There was a problem hiding this comment.
cc @fegin as sharded optimizer states are built on top of it
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.
Stack from ghstack:
optim_inputorders across ranks #78599 [FSDP] Allow differentoptim_inputorders across ranksfull_optim_state_dict()#78784 [FSDP][Docs] Fix typo infull_optim_state_dict()This enables the first argument to the optimizer constructor,
optim_input, to have different orders across ranks, e.g. if the parameters in one parameter group are permuted. This requires modification tofull_optim_state_dict(),shard_full_optim_state_dict(), andscatter_full_optim_state_dict().The high-level algorithmic change is that the state dicts are kept as being keyed by unflattened parameter name until after sharding/unsharding and flattening/unflattening and are rekeyed to be by parameter ID according to each rank's own
optim_inputonly at the end.Because this PR adds non-parameter-specific collectives to
full_optim_state_dict(), it adds agroup=Noneargument to the method to have a process group to default to when running those common collectives.