[FSDP][5/N] Add manual "wrapping" support for fully_shard#90874
Closed
awgu wants to merge 13 commits intogh/awgu/279/basefrom
Closed
[FSDP][5/N] Add manual "wrapping" support for fully_shard#90874awgu wants to merge 13 commits intogh/awgu/279/basefrom
fully_shard#90874awgu wants to merge 13 commits intogh/awgu/279/basefrom
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/90874
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1924517: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Dec 14, 2022
This is not ready for review. [ghstack-poisoned]
This is not ready for review. [ghstack-poisoned]
This is not ready for review. [ghstack-poisoned]
This is not ready for review. [ghstack-poisoned]
This is not ready for review. [ghstack-poisoned]
mrshenli
reviewed
Dec 16, 2022
| flat_param = handle.flat_param | ||
| if flat_param not in self.flat_param_to_prefixed_param_names: | ||
| if flat_param not in self.param_to_fqn: | ||
| continue |
Contributor
There was a problem hiding this comment.
when is a handle invalid?
Collaborator
Author
There was a problem hiding this comment.
I traced through the history of my PRs, and it looks like this check was added arbitrarily. However, reasoning about it now, I feel like this check is important for execution to not crash for use_orig_params=True since self.param_to_fqn is constructed via _get_param_to_fqns(root_module) and for use_orig_params=True, _get_param_to_fqns does not include any FlatParameters since they are not registered.
I think _exec_order_utils.py needs to be revisited for use_orig_params=True.
|
|
||
| states = [state] if _is_composable(state) else _get_fsdp_states(state) | ||
| for state in states: | ||
| for state in _get_fsdp_states(module): |
Contributor
There was a problem hiding this comment.
hmm, is this intentional to abuse the var name state?
Contributor
There was a problem hiding this comment.
and on line 887, is it a bug? I assume your intention is to set it on the original state var instead of the one used in the loop?
Collaborator
Author
There was a problem hiding this comment.
Oh man. Thanks for the great catch. Let me fix this.
This PR adds manual "wrapping" support for `fully_shard`. For example, for ``` fully_shard(mod.sub) fully_shard(mod) ``` `mod.sub` and `mod` will share the same FSDP data structures. To have parity with wrapper FSDP, this PR only checks support for when each manual application of `fully_shard` passes `policy=None`. Hybrid auto / manual wrapping is not in scope for this PR since it is not supported for wrapper FSDP either. I can follow up to either add support properly or raise and error early. [ghstack-poisoned]
awgu
pushed a commit
to awgu/pytorch
that referenced
this pull request
Dec 17, 2022
ghstack-source-id: 4e1e0b4 Pull Request resolved: pytorch#90874
This PR adds manual "wrapping" support for `fully_shard`. For example, for ``` fully_shard(mod.sub) fully_shard(mod) ``` `mod.sub` and `mod` will share the same FSDP data structures. To have parity with wrapper FSDP, this PR only checks support for when each manual application of `fully_shard` passes `policy=None`. Hybrid auto / manual wrapping is not in scope for this PR since it is not supported for wrapper FSDP either. I can follow up to either add support properly or raise and error early. [ghstack-poisoned]
This PR adds manual "wrapping" support for `fully_shard`. For example, for ``` fully_shard(mod.sub) fully_shard(mod) ``` `mod.sub` and `mod` will share the same FSDP data structures. To have parity with wrapper FSDP, this PR only checks support for when each manual application of `fully_shard` passes `policy=None`. Hybrid auto / manual wrapping is not in scope for this PR since it is not supported for wrapper FSDP either. I can follow up to either add support properly or raise and error early. [ghstack-poisoned]
This PR adds manual "wrapping" support for `fully_shard`. For example, for ``` fully_shard(mod.sub) fully_shard(mod) ``` `mod.sub` and `mod` will share the same FSDP data structures. To have parity with wrapper FSDP, this PR only checks support for when each manual application of `fully_shard` passes `policy=None`. Hybrid auto / manual wrapping is not in scope for this PR since it is not supported for wrapper FSDP either. I can follow up to either add support properly or raise and error early. [ghstack-poisoned]
awgu
pushed a commit
to awgu/pytorch
that referenced
this pull request
Dec 19, 2022
ghstack-source-id: 03633b9 Pull Request resolved: pytorch#90874
This PR adds manual "wrapping" support for `fully_shard`. For example, for ``` fully_shard(mod.sub) fully_shard(mod) ``` `mod.sub` and `mod` will share the same FSDP data structures. To have parity with wrapper FSDP, this PR only checks support for when each manual application of `fully_shard` passes `policy=None`. Hybrid auto / manual wrapping is not in scope for this PR since it is not supported for wrapper FSDP either. I can follow up to either add support properly or raise and error early. [ghstack-poisoned]
Contributor
|
This pull request has been merged in 32fde53. |
awgu
pushed a commit
to awgu/pytorch
that referenced
this pull request
Dec 20, 2022
ghstack-source-id: 97a14bc Pull Request resolved: pytorch#90874
This was referenced Dec 20, 2022
This was referenced Jan 5, 2023
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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:
replicateinfully_shard#91044 [FSDP][7/N] Supportreplicateinfully_shard_FSDPStatetraversal #90959 [FSDP][6/N] Add note explaining idioms for_FSDPStatetraversalfully_shard#90874 [FSDP][5/N] Add manual "wrapping" support forfully_shardThis PR adds manual "wrapping" support for
fully_shard. For example, formod.subandmodwill share the same FSDP data structures.To have parity with wrapper FSDP, this PR only checks support for when each manual application of
fully_shardpassespolicy=None. Hybrid auto / manual wrapping is not in scope for this PR since it is not supported for wrapper FSDP either. I can follow up to either add support properly or raise and error early.