Adopt state_dict_pre_hook in FSDP #90436
Conversation
[ghstack-poisoned]
rohan-varma
requested review from
H-Huang,
awgu,
kwen2501,
mrshenli,
pritamdamania87,
wanchaol and
zhaojuanmao
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/90436
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 108c6bd: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
rohan-varma
changed the title
[ghstack-poisoned]
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. [ghstack-poisoned]
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. [ghstack-poisoned]
rohan-varma
changed the title
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. [ghstack-poisoned]
| StateDictType.FULL_STATE_DICT: functools.partial( | ||
| _full_pre_state_dict_hook, fsdp_state | ||
| ), | ||
| StateDictType.LOCAL_STATE_DICT: functools.partial( | ||
| _local_pre_state_dict_hook, fsdp_state | ||
| ), | ||
| StateDictType.SHARDED_STATE_DICT: functools.partial( | ||
| _sharded_pre_state_dict_hook, fsdp_state | ||
| ), |
There was a problem hiding this comment.
Prefer not to use partial here but just pass the fsdp_state to _pre_state_dict_hook_fn[fsdp_state._state_dict_type]().
| try: | ||
| state_dict.pop(f"{prefix}{FLAT_PARAM}") | ||
| except: | ||
| pass |
There was a problem hiding this comment.
This is pretty strange. We don't have this try-except before, right? I'm not sure if there will be a hidden bug if we use try-except to avoid the issue.
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. [ghstack-poisoned]
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. [ghstack-poisoned]
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. [ghstack-poisoned]
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. This removes `def state_dict` entirely and paves the path for composable API as well. [ghstack-poisoned]
|
@pytorchbot merge -f "CI passed" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
| ) | ||
| from ._state_dict_utils import ( | ||
| _post_load_state_dict_hook, | ||
| _pre_state_dict_hook, |
There was a problem hiding this comment.
It looks like this is going to get sorted in the next ufmt.
4 participants
Stack from ghstack (oldest at bottom):
Use register_state_dict_pre_hook in FSDP to simplify state_dict implementations & remove hacks. This removes
def state_dictentirely and paves the path for composable API as well.