[WIP] Composable FSDP Follow-Ups
Labels: oncall: distributed, triaged, module: fsdp

This tracks the follow-ups from my PR stack(s) for implementing composable FSDP.

**High Priority**
- Investigate behavior around `_apply()` and whether pre- and post-`_apply()` hooks are needed/reasonable
    - Since composable FSDP requires `use_orig_params=True` and `use_orig_params=True` usually overrides `nn.Module._apply()` to expose the `FlatParameter`s before `_apply()` and restore the original parameters after:
https://github.com/pytorch/pytorch/blob/3c6bddc3f6347ce7d1ed33aee94cdaa953cbc387/torch/distributed/fsdp/fully_sharded_data_parallel.py#L1605-L1610
    - The reason for this is to allocate contiguous memory for each `FlatParameter` when performing `.cuda()` (or similar) instead of allocating memory per original parameter and requiring a separate allocation and write back.
    - One possibility is to enforce the `device_id` argument for composable FSDP so that FSDP initialization can ensure the sharded `FlatParameter` memory is allocated on GPU directly before setting the original parameters to be views into that memory.

**Medium Priority**
- Define FSDP's contract for buffers (e.g. for ignored modules)
https://github.com/pytorch/pytorch/pull/87923#discussion_r1008537088
- Add sibling shared parameter support
https://github.com/pytorch/pytorch/pull/87923#discussion_r1008550533
- Investigate enabling side streams in `_to_kwargs()` in the root pre-forward
https://github.com/pytorch/pytorch/pull/87915#discussion_r1008406766
- Migrate to module pre-backward hook to avoid registering the pre-backward hook every iteration (on the forward pass output tensors that require gradient)
https://github.com/pytorch/pytorch/pull/87927#discussion_r1009505679
- Investigate the padded unsharded gradient allocation in the post-backward (i.e. can we avoid making a copy by calling `F.pad()`)
https://github.com/pytorch/pytorch/pull/87927#discussion_r1009516921
- Investigate if reference cycle involving `_post_backward_hook_state()` is a concern
https://github.com/pytorch/pytorch/pull/87927#discussion_r1009512801

**Low Priority**
- [BE] Investigate `tree_map()` instead of `_apply_to_tensors()` for mapping the low precision cast for forward pass inputs
https://github.com/pytorch/pytorch/pull/87915#discussion_r1008408406
    - `tree_map()` does not seem to currently support `dataclass` and `PackedSequence`, which `_apply_to_tensor()` does. This is not high priority and can be revisited in the future.
- [Done] ~~Add note to `BackwardPrefetch` docs pointing to using the rate limiter if encountering issues from increased memory when enabling the prefetching
https://github.com/pytorch/pytorch/pull/87917#discussion_r1008455071~~
- Verify why `CPUOffload` is a `dataclass` and if there are any plans to add additional attributes
https://github.com/pytorch/pytorch/pull/87920#discussion_r1008470789
- [Done] ~~[BE] Investigate why offloading to CPU in the FSDP constructor happens in `torch.no_grad()` context but not the preceding `module.to()`
https://github.com/pytorch/pytorch/pull/87921#discussion_r1008508533~~
- Decide if the `device_id` constructor argument should support `str` or not (e.g. `device_id="cuda"`)
https://github.com/pytorch/pytorch/pull/87921#discussion_r1008505562
- Investigate if FSDP should permit `FullyShardedDataParallel` instances in `ignored_modules` but simply ignore them and recurse into their child modules to be consistent with the behavior for nested FSDP instances not directly passed to `ignored_modules`
https://github.com/pytorch/pytorch/pull/87921#discussion_r1008502560
- [BE] Investigate if the manual BFS in `_apply_to_modules` is necessary instead of simply looping over `named_modules()` (see `_get_param_to_unflat_param_names()` and `_get_buffer_names()`)
https://github.com/pytorch/pytorch/pull/87921#discussion_r1008494191
- [Done] ~~[BE] Refactor `_get_param_to_unflat_param_names()` to unify implementation naming conventions
https://github.com/pytorch/pytorch/pull/87921#discussion_r1008483348~~
- Make the usage of preceding underscore (`_`) for class, method, and function names consistent (i.e. unify what it means to be private or public)

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501