[FSDP]Skip unshard call during checkpointing for NO_SHARD sharding strategy #101095
Conversation
…rategy [ghstack-poisoned]
zhaojuanmao
requested review from
mrshenli,
rohan-varma,
H-Huang,
awgu,
kwen2501,
wanchaol,
fegin,
kiukchung and
d4l3k
as code owners
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/101095
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 5 PendingAs of commit 463b32d: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
…rategy ghstack-source-id: c2e57d97d5b5c9c48674983f2b6812418c231fca Pull Request resolved: #101095
|
Probably need to resubmit to trigger the CI, there was some infra issue. |
…sharding strategy" [ghstack-poisoned]
…rategy ghstack-source-id: 417efbb706677f99cf29ebb41cc394184229ffe1 Pull Request resolved: #101095
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
| @@ -146,6 +146,12 @@ def _common_unshard_pre_state_dict_hook( | |||
| Performs the pre-state_dict tasks shared by all state_dict types that require | |||
| ``_unshard_fsdp_state_params()``. FULL_STATE_DICT and SHARDED_STATE_DICT use this hook. | |||
| """ | |||
| # For composable `fully_shard`, it does not need to unshard parameters for `NO_SHARD` cases. | |||
| if ( | |||
| _is_composable(fsdp_state) | |||
There was a problem hiding this comment.
why does it not need to unshard for no_shard usecases only for composable, but this is fine for class based? @awgu would you happen to know?
| _is_composable(fsdp_state) | ||
| and fsdp_state.sharding_strategy == ShardingStrategy.NO_SHARD | ||
| ): | ||
| _exit_unshard_params_ctx(module, fsdp_state) |
There was a problem hiding this comment.
wonder if duplication of these chekcs can be reduced by putting it inside exit / enter unshard params ctx?
Stack from ghstack (oldest at bottom):