[FSDP][1/N] Refactor module materialization #94196
Conversation
[ghstack-poisoned]
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/94196
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 221ef73: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
ghstack-source-id: d064a75 Pull Request resolved: pytorch#94196
awgu
requested review from
H-Huang,
kwen2501,
mrshenli,
rohan-varma,
wanchaol and
zhaojuanmao
as code owners
| if is_meta_module or is_torchdistX_deferred_init: | ||
| materialized_module = True | ||
| # Save the parameter and buffer names to reacquire references after | ||
| # after materialization since their variables may change |
There was a problem hiding this comment.
Even after reading the PR description, I'm not 100% sure why the variables may change after materialization? I thought materialization is all about filling in meta parameters with their actual values?
There was a problem hiding this comment.
Yes, but at the implementation level, module.to(device) will replace all meta-device parameters with new Python parameter variables.
This function returns False for meta-device tensors:
pytorch/torch/nn/modules/module.py
Lines 799 to 800 in a064ce1
This leads to the
else branch:pytorch/torch/nn/modules/module.py
Lines 821 to 829 in a064ce1
| elif is_meta_module: | ||
| _materialize_meta_module(fully_sharded_module, device_id) | ||
| elif is_torchdistX_deferred_init: | ||
| deferred_init.materialize_module( |
There was a problem hiding this comment.
do we have unittests covering deferred init for the composable path?
There was a problem hiding this comment.
We do not because torchdistX does not support the latest PyTorch version, so I could not run torchdistX locally.
**Overview** This refactors module materialization (i.e. meta device or `torchdistX` deferred initialization) to compute the parameter and buffer names as needed instead of pre-computing them. These are needed to reacquire references to the states (e.g. `module.get_parameter(param_name)`) after materialization since the materialization may create new variables. This refactor simplifies `_get_fully_sharded_module_to_states()` (the core function for "pseudo auto wrapping") to better enable lowest common ancestor (LCA) module computation for shared parameters, for which tracking parameter and buffer names may complicate the already non-obvious implementation. **Discussion** The tradeoff is a worst case quadratic traversal over modules if materializing all of them. However, since (1) the number of modules is relatively small, (2) the computation per module in the quadratic traversal is negligible, (3) this runs only once per training session, and (4) module materialization targets truly large models, I think this tradeoff is tolerable. **For Reviewers** - `_init_param_handle_from_module()` initializes _one_ `FlatParamHandle` from a fully sharded module and represents the module wrapper code path. For this code path, there is no need to reacquire references to the parameters/buffers for now since the managed parameters are only computed after materialization. This works because the managed parameters have a simple definition: any parameter in the local root module's tree excluding those already marked as flattened by FSDP. Similarly, FSDP marks buffers to indicate that they have already been processed (synced if `sync_module_states`). - `_init_param_handles_from_module()` initializes _all_ `FlatParamHandle`s from a fully sharded module and represents the composable code path. For this code path, we must reacquire references to parameters/buffers because each logical wrapping is specified as a list of parameters/buffers to group together by those variables and because materialization may create new variables. [ghstack-poisoned]
ghstack-source-id: 3563d70 Pull Request resolved: pytorch#94196
|
@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 |
Stack from ghstack:
fully_shard#94198 [FSDP][3/N] Add LCA logic tofully_shardOverview
This refactors module materialization (i.e. meta device or
torchdistXdeferred initialization) to compute the parameter and buffer names as needed instead of pre-computing them. These are needed to reacquire references to the states (e.g.module.get_parameter(param_name)) after materialization since the materialization may create new variables.This refactor simplifies
_get_fully_sharded_module_to_states()(the core function for "pseudo auto wrapping") to better enable lowest common ancestor (LCA) module computation for shared parameters, for which tracking parameter and buffer names may complicate the already non-obvious implementation.Discussion
The tradeoff is a worst case quadratic traversal over modules if materializing all of them. However, since (1) the number of modules is relatively small, (2) the computation per module in the quadratic traversal is negligible, (3) this runs only once per training session, and (4) module materialization targets truly large models, I think this tradeoff is tolerable.
For Reviewers
_init_param_handle_from_module()initializes oneFlatParamHandlefrom a fully sharded module and represents the module wrapper code path. For this code path, there is no need to reacquire references to the parameters/buffers for now since the managed parameters are only computed after materialization. This works because the managed parameters have a simple definition: any parameter in the local root module's tree excluding those already marked as flattened by FSDP. Similarly, FSDP marks buffers to indicate that they have already been processed (synced ifsync_module_states)._init_param_handles_from_module()initializes allFlatParamHandles from a fully sharded module and represents the composable code path. For this code path, we must reacquire references to parameters/buffers because each logical wrapping is specified as a list of parameters/buffers to group together by those variables and because materialization may create new variables.