[FSDP] `fully_shard` Follow-Ups & Known Issues
Labels: oncall: distributed, triaged, module: fsdp

This issue is tracking work items for `fully_shard`. The list will be dynamically evolving.
I am assigning this to myself for organizational purposes, but others may feel free to also assign themselves.

**Core Features**
- [x] Support manual "wrapping": For `fully_shard`, this means having multiple applications of `fully_shard` share data structures and only have a single local FSDP root module instead of having each application be its own distinct FSDP.
- [x] Support model checkpointing (i.e. `module.state_dict` / `module.load_state_dict`).
- [x] Support optimizer state checkpointing (i.e. `optim.state_dict` / `optim.load_state_dict`).

**Composability**
- [x] Ignore modules annotated `replicate` and test `fully_shard` above `replicate`.
    - [x] Change `_get_fsdp_states()` to not traverse down through any module annotated with a different composable API. This requires the global registry to check a module's annotation.
- [ ] Investigate `fully_shard` and `checkpoint(reentrant=True)` compatibility.
- [ ] Investigate `fully_shard` and `checkpoint(reentrant=False)` compatibility.


**Better Engineering**
- [ ] Add comprehensive docstring to `fully_shard` explaining the usage and composability.
- [x] Refactor `_lazy_init` to take a sequence of constructor arguments that need to be uniform across all FSDP instances / `fully_shard` applications and verify that. Currently, we have the check only for `limit_all_gathers`. However, `use_orig_params` needs to be uniform. In this way, we will know explicitly which arguments we expect to be uniform and which we do not.
- [x] Remove `HandleConfig` and store its attributes directly on the `FlatParamHandle`.
- [x] https://github.com/pytorch/pytorch/issues/90788
- [ ] Move all FSDP attributes to the `_FSDPState` class to enable direct type checking.
- [x] Split `test_fully_shard.py` by sub-functionality. Eventually, `test_fully_shard.py` will become too long, which will make it hard to read through. I have been keeping tests grouped by sub-functionality (e.g. initialization, runtime, checkpointing). We can separate these into their own files.
- [x] After introducing "fully sharded module", I missed updating the variable names in `torch/distributed/fsdp/_wrap_utils.py` and `test/distributed/fsdp/test_utils.py` -- fix these variable names to use "fully sharded module" instead of "submodule".
https://github.com/pytorch/pytorch/blob/212873c615dd3455a24d390605335aeeebd76236/torch/distributed/fsdp/_wrap_utils.py#L119
https://github.com/pytorch/pytorch/blob/212873c615dd3455a24d390605335aeeebd76236/test/distributed/fsdp/test_utils.py#L163

**For Discussion**
- FSDP currently raises a `ValueError` if a user passes in a `FullyShardedDataParallel` instance into the `ignored_modules` argument. We can relax this by skipping the `FullyShardedDataParallel` instance instead taking its wrapped module.
https://github.com/pytorch/pytorch/blob/212873c615dd3455a24d390605335aeeebd76236/torch/distributed/fsdp/_init_utils.py#L541-L542
- `fully_shard` handles shared parameters slightly differently than `FullyShardedDataParallel` in its pseudo-auto-wrapping. We should decide on how we want to proceed with this in the ideal state and add unit tests to capture the differences.



cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501