[FSDP][1/N] Move wrapper ModuleWrapPolicy to new path
#104346
Closed
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
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/104346
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit ba95ffc: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
awgu
added a commit
that referenced
this pull request
Jun 28, 2023
ghstack-source-id: eafbb91bf2339aeea230b92a0356e52003c5ce7a Pull Request resolved: #104346
[ghstack-poisoned]
awgu
added a commit
that referenced
this pull request
Jun 28, 2023
ghstack-source-id: 3c7ff44f030362717f1ae2fbc2b4cd9014ebedc6 Pull Request resolved: #104346
The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook. We can directly store the global state as a mapping.
To-do in follow-ups (not in order):
- Add frozen parameter check before `_post_order_apply()`
- Add shared parameter check before `_post_order_apply()`
- Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg)
- Refactor `fully_shard()` auto wrap unify with `FullyShardedDataParallel` auto wrap, where the only difference should be `fn_to_apply` in `_post_order_apply()`
- This means that for `fully_shard()`'s auto wrap, it will call `fully_shard()` on the target submodules, constructing a new `_FSDPState` object for each just like for the wrapper path.
- This change prohibits extensions like non-module-aligned wrapping, but it allows for unifying the code paths to decrease the likelihood for bugs. I do not foresee us pursuing non-module-aligned wrapping in the near term.
- After this change, we can then revisit the `ignored_states` with auto wrapping fix and land that without changing `_unshard_params()`.
[ghstack-poisoned]
awgu
added a commit
that referenced
this pull request
Jun 28, 2023
ghstack-source-id: 892168e233a9e74b7750227c086cb4689ab23572 Pull Request resolved: #104346
This was referenced Jun 28, 2023
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: 892168e233a9e74b7750227c086cb4689ab23572 Pull Request resolved: pytorch#104346
The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.)
To-do in follow-ups (not in order):
- Add frozen parameter check before `_post_order_apply()`
- Add shared parameter check before `_post_order_apply()`
- Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg)
- Refactor `fully_shard()` auto wrap unify with `FullyShardedDataParallel` auto wrap, where the only difference should be `fn_to_apply` in `_post_order_apply()`
- This means that for `fully_shard()`'s auto wrap, it will call `fully_shard()` on the target submodules, constructing a new `_FSDPState` object for each just like for the wrapper path.
- This change prohibits extensions like non-module-aligned wrapping, but it allows for unifying the code paths to decrease the likelihood for bugs. I do not foresee us pursuing non-module-aligned wrapping in the near term.
- After this change, we can then revisit the `ignored_states` with auto wrapping fix and land that without changing `_unshard_params()`.
[ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: 304383832c2b6bc8eafbfc7730d8198365049780 Pull Request resolved: pytorch#104346
This PR is the first in refactoring the auto wrapping. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
This was referenced Jun 29, 2023
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: b62d622a3c0745bc3149202a81ac730a7aefe995 Pull Request resolved: pytorch#104346
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 30, 2023
ghstack-source-id: b62d622a3c0745bc3149202a81ac730a7aefe995 Pull Request resolved: pytorch#104346
This PR is the first in refactoring the auto wrapping, only affecting `ModuleWrapPolicy` for wrapper `FullyShardedDataParallel`. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 30, 2023
ghstack-source-id: 6600d7ad0d44834537abefee871178fd3cdd6ff7 Pull Request resolved: pytorch#104346
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jul 5, 2023
ghstack-source-id: 6600d7ad0d44834537abefee871178fd3cdd6ff7 Pull Request resolved: pytorch#104346
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jul 5, 2023
ghstack-source-id: 6600d7ad0d44834537abefee871178fd3cdd6ff7 Pull Request resolved: pytorch#104346
This PR is the first in refactoring the auto wrapping, only affecting `ModuleWrapPolicy` for wrapper `FullyShardedDataParallel`. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
awgu
requested review from
rohan-varma,
H-Huang,
kwen2501,
wanchaol,
fegin,
fduwjj,
kiukchung and
d4l3k
as code owners
This was referenced Jul 6, 2023
This PR is the first in refactoring the auto wrapping, only affecting `ModuleWrapPolicy` for wrapper `FullyShardedDataParallel`. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
rohan-varma
approved these changes
Jul 6, 2023
| # NOTE: If the forward did not have any floating-point tensors, | ||
| # then the dtype will not be set for this module, and we do not | ||
| # upcast the dtype. | ||
| if module in _MODULE_TO_INP_DTYPE: |
There was a problem hiding this comment.
so _MODULE_TO_INP_DTYPE generalizes old_dtype to be on a per-module basis?
wanchaol
reviewed
Jul 6, 2023
Good point! I wonder how we can fit more info given the title character limit though :/ |
fegin
reviewed
Jul 7, 2023
fegin
reviewed
Jul 7, 2023
fegin
approved these changes
Jul 7, 2023
This PR is the first in refactoring the auto wrapping, only affecting `ModuleWrapPolicy` for wrapper `FullyShardedDataParallel`. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
This PR is the first in refactoring the auto wrapping, only affecting `ModuleWrapPolicy` for wrapper `FullyShardedDataParallel`. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
This PR is the first in refactoring the auto wrapping, only affecting `ModuleWrapPolicy` for wrapper `FullyShardedDataParallel`. The end goal is to improve the auto wrapping infra to support: - Checking valid frozen parameters (uniform frozenness per FSDP) - Checking valid shared parameters (shared parameters assigned to their lowest-common-ancestor module or higher) - Writing auto wrapping policies that may take multiple passes over the module tree - Specifying different FSDP kwargs per FSDP instance (instead of enforcing the same for all FSDP instances constructed via an auto wrap policy) The way I envision achieving this is that, we decouple the actual "wrapping" (which is `_post_order_apply()` in this PR) from constructing the wrapping targets and kwargs (which is `target_module_to_kwargs` in this PR). In that way, a policy reduces to just constructing that latter `target_module_to_kwargs` mapping. I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple. The change to how `old_dtype` is handled is mainly to avoid keeping a reference to `_override_module_mixed_precision()` function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.) To-do in follow-ups (not in order): - Add frozen parameter check before `_post_order_apply()` - Add shared parameter check before `_post_order_apply()` - Expose wrapping policy that allows per module / per module class kwarg customization (where any unspecified kwarg adopts the root's kwarg) [ghstack-poisoned]
voznesenskym
pushed a commit
that referenced
this pull request
Jul 19, 2023
ghstack-source-id: 158dc95b1093450c647296ee05ef13192ed67fc9 Pull Request resolved: #104346
voznesenskym
pushed a commit
that referenced
this pull request
Jul 21, 2023
ghstack-source-id: 158dc95b1093450c647296ee05ef13192ed67fc9 Pull Request resolved: #104346
voznesenskym
pushed a commit
that referenced
this pull request
Aug 7, 2023
ghstack-source-id: 158dc95b1093450c647296ee05ef13192ed67fc9 Pull Request resolved: #104346
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 (oldest at bottom):
ModuleWrapPolicy#104427ignored_states+ auto wrap (for now) #104418_get_fully_sharded_module_to_states#104409fully_shardauto wrap #104408_auto_wrapforfully_shard#104407ModuleWrapPolicyto new path #104346This PR is the first in refactoring the auto wrapping, only affecting
ModuleWrapPolicyfor wrapperFullyShardedDataParallel. The end goal is to improve the auto wrapping infra to support:The way I envision achieving this is that, we decouple the actual "wrapping" (which is
_post_order_apply()in this PR) from constructing the wrapping targets and kwargs (which istarget_module_to_kwargsin this PR). In that way, a policy reduces to just constructing that lattertarget_module_to_kwargsmapping.I do not personally recommend the size-based policy, but if we wanted to implement that under this new organization, the tracking of wrapped/nonwrapped numel should be done in the pass over the module tree prior to the actual "wrapping". This modularization keeps the actual "wrapping" part simple.
The change to how
old_dtypeis handled is mainly to avoid keeping a reference to_override_module_mixed_precision()function closure in each hook and to allow the function to take in all module clases at once to return which ones actually got overridden for the downstream error message. (We can directly store the global state as a mapping.)To-do in follow-ups (not in order):
_post_order_apply()_post_order_apply()