[FSDP][3/N] Unify fully_shard auto wrap
#104408
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/104408
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit c501068: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
awgu
added a commit
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: 9e0cf806b4bc63ef5bb5361c6a4c1cb33cd80c7c Pull Request resolved: #104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: dede963dea8977e420b32c5b10b0587688c9f693 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: 457c67dc90ed19b789b827fd99a06b4bcd5951b6 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 29, 2023
ghstack-source-id: 457c67dc90ed19b789b827fd99a06b4bcd5951b6 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 30, 2023
ghstack-source-id: f811cfa160eb549249eda792e376dc0277053373 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 30, 2023
ghstack-source-id: 141ba57ed15da65b7052e81f706b82925a376d33 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jun 30, 2023
ghstack-source-id: 5185a43b01c99015c96e04ea13a87f2449f39022 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. [ghstack-poisoned]
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jul 5, 2023
ghstack-source-id: 4ad898a9777d75b08b32a786fcaab38388900ae9 Pull Request resolved: pytorch#104408
awgu
added a commit
to awgu/pytorch
that referenced
this pull request
Jul 5, 2023
ghstack-source-id: 4ad898a9777d75b08b32a786fcaab38388900ae9 Pull Request resolved: pytorch#104408
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. This includes several important fixes: - We should register the pre/post-forward hooks on the module regardless of it has managed parameters. - We can permit `_module_handles` to return `[]` in the composable path (for when the module has no managed parameters). - We should unify the paths for `_get_buffers_and_dtypes_for_computation()` (previously, composable path was buggy in some cases). [ghstack-poisoned]
awgu
requested review from
zhaojuanmao,
rohan-varma,
H-Huang,
kwen2501,
wanchaol,
fegin,
fduwjj,
yhcharles,
kiukchung and
d4l3k
as code owners
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. This includes several important fixes: - We should register the pre/post-forward hooks on the module regardless of it has managed parameters. - We can permit `_module_handles` to return `[]` in the composable path (for when the module has no managed parameters). - We should unify the paths for `_get_buffers_and_dtypes_for_computation()` (previously, composable path was buggy in some cases). [ghstack-poisoned]
This was referenced Jul 6, 2023
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. This includes several important fixes: - We should register the pre/post-forward hooks on the module regardless of it has managed parameters. - We can permit `_module_handles` to return `[]` in the composable path (for when the module has no managed parameters). - We should unify the paths for `_get_buffers_and_dtypes_for_computation()` (previously, composable path was buggy in some cases). [ghstack-poisoned]
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. This includes several important fixes: - We should register the pre/post-forward hooks on the module regardless of it has managed parameters. - We can permit `_module_handles` to return `[]` in the composable path (for when the module has no managed parameters). - We should unify the paths for `_get_buffers_and_dtypes_for_computation()` (previously, composable path was buggy in some cases). [ghstack-poisoned]
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. This includes several important fixes: - We should register the pre/post-forward hooks on the module regardless of it has managed parameters. - We can permit `_module_handles` to return `[]` in the composable path (for when the module has no managed parameters). - We should unify the paths for `_get_buffers_and_dtypes_for_computation()` (previously, composable path was buggy in some cases). [ghstack-poisoned]
rohan-varma
approved these changes
Jul 7, 2023
| # A valid FSDP state may have no managed parameters and hence no | ||
| # handles, meaning no entry in `_fully_sharded_module_to_handles` | ||
| if len(state._handles) == 0: | ||
| return [] |
There was a problem hiding this comment.
did we add a test for this + test to ensure that if a composable FSDP module manages no params, it is still marked as FSDP managed?
There was a problem hiding this comment.
I think the _has_fsdp_params() check is still valid. Before, the composable path would raise an error when it did not need to, which is why I had to add this case.
In other words, this is covered by the existing tests. summon_full_params() on a module with fully_shard() applied but no managed parameters would error otherwise.
This moves `fully_shard` to use `_auto_wrap()` just like `FullyShardedDataParallel`. This means that `fully_shard` goes through the `_init_param_handle_from_module()` path (i.e. 1 `fully_shard` per "wrap"), removing the need for `_init_param_handles_from_module()` (which was 1 `fully_shard` for all "wraps" of a given policy). `_auto_wrap()` simply calls `fully_shard` on target submodules. This includes several important fixes: - We should register the pre/post-forward hooks on the module regardless of it has managed parameters. - We can permit `_module_handles` to return `[]` in the composable path (for when the module has no managed parameters). - We should unify the paths for `_get_buffers_and_dtypes_for_computation()` (previously, composable path was buggy in some cases). [ghstack-poisoned]
3 tasks
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 moves
fully_shardto use_auto_wrap()just likeFullyShardedDataParallel. This means thatfully_shardgoes through the_init_param_handle_from_module()path (i.e. 1fully_shardper "wrap"), removing the need for_init_param_handles_from_module()(which was 1fully_shardfor all "wraps" of a given policy)._auto_wrap()simply callsfully_shardon target submodules.This includes several important fixes:
_module_handlesto return[]in the composable path (for when the module has no managed parameters)._get_buffers_and_dtypes_for_computation()(previously, composable path was buggy in some cases).