[BE] Enabled mypy in common_fsdp.py
#118755
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/118755
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit fe502d8 with merge base 278a0e1 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
[ghstack-poisoned]
|
|
||
| def get_future(): | ||
| future = torch.futures.Future() | ||
| future: torch.futures.Future = torch.futures.Future() |
There was a problem hiding this comment.
For some reason mypy wanted this type annotation.
| ] | ||
| world_size = dist.get_world_size(process_group) | ||
| olist = [None for _ in range(world_size)] | ||
| dist.all_gather_object(olist, named_module_states, group=process_group) |
There was a problem hiding this comment.
all_gather_object fills in the olist destructively. Another approach could be to initialize olist with some dummy object of the expected type.
There was a problem hiding this comment.
would it be a good BE item (maybe for myself), to allow [] ?
olist = []
dist.all_gather_object(olist
There was a problem hiding this comment.
This seems reasonable to me. all_gather_object can destructively modify the olist as it already does today and append world_size many elements.
This sounds like a good BE task!
| dist.reduce_scatter_tensor = orig_reduce_scatter | ||
|
|
||
|
|
||
| @no_type_check |
There was a problem hiding this comment.
These patching methods complain about assigning to a method like FSDPParamGroup.unshard = new_unshard. Since we have two of these assignments (to set to new and to restore to old) per patch context, I preferred to just ignore type checking, as it is not too valuable here.
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
ghstack-source-id: 5430730 Pull Request resolved: pytorch#118755
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
|
@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 |
This PR adds a way to do gradient accumulation without collectives (i.e. reduce-scatter for FSDP and reduce-scatter/all-reduce for HSDP, though HSDP is not yet implemented). Since the `no_sync()` context manager has received some feedback, we simply define a method on the module to set whether the module requires gradient synchronization or not, where this method can recurse or not.
```
# Before with `no_sync()`:
with fsdp_model.no_sync() if not is_last_microbatch else contextlib.nullcontext():
# Forward/backward
# After with a setter:
fsdp_model.set_requires_gradient_sync(not is_last_microbatch)
# Forward/backward
```
Having the method be able to recurse or not also gives some flexibility. For example, some large modules can still reduce-scatter, while some smaller modules can avoid it to save communication bandwidth:
```
fsdp_modules_to_reduce_scatter: Set[nn.Module] = ...
for module in fsdp_model.modules():
if isinstance(module, FSDP) and module not in fsdp_modules_to_reduce_scatter:
module.set_requires_gradient_sync(not is_last_microbatch)
# Forward/backward
```
(Separately, we may expose a helper for `return [module for model.modules() if isinstance(module, FSDP)]`.)
---
To show the spirit of this API choice, I also included `set_requires_all_reduce` that would give us the ability to only reduce-scatter but not all-reduce for HSDP (originally from the MiCS paper). If we want to flexibly support heterogeneous sharding where FSDP is applied to some modules and HSDP to others in the same model, then having a module-level method that has the option to not recurse makes sense to me.
Pull Request resolved: #118298
Approved by: https://github.com/wconstab, https://github.com/wanchaol
ghstack dependencies: #119550, #118136, #118223, #118755, #119825
5 participants
Stack from ghstack (oldest at bottom):
common_fsdp.py#118755no_grad; removedno_grad#119550