Repository: pytorch
Issue: [FSDP] Revisit meta device initialization
Labels: oncall: distributed,triaged,module: fsdp

### Context
Today, `FullyShardedDataParallel` (FSDP) supports meta device initialization via two paths, where the precondition is that the `module` passed to FSDP has some parameter on meta device:
1. The user passes `param_int_fn=None` (e.g. by default).
2. The user passes an explicit `param_init_fn`.

https://github.com/pytorch/pytorch/blob/dc40b6d04320baf26b940225a65f40466ebf3664/torch/distributed/fsdp/fully_sharded_data_parallel.py#L249

In case 1, FSDP has its own code to provide a default parameter initialization function. In case 2, FSDP calls the passed-in function.

### Issue
Both the default parameter initialization function and the example explicit `param_init_fn` in the unit tests (`test_fsdp_meta.py`) have a silent correctness issue.

The (undocumented) invariant required of the `param_init_fn` is as follows:
> A `param_init_fn` should only modify the parameters managed by the FSDP instance to which the function is passed.

The contrapositive is that a `param_init_fn` **should not** modify any parameters not managed by the FSDP instance, which we currently violate.

#### 1: Default Meta Device Initialization

https://github.com/pytorch/pytorch/blob/dc40b6d04320baf26b940225a65f40466ebf3664/torch/distributed/fsdp/_init_utils.py#L529-L537

#### 2: `test_fsdp_meta.py` Example `param_init_fn`

https://github.com/pytorch/pytorch/blob/dc40b6d04320baf26b940225a65f40466ebf3664/test/distributed/fsdp/test_fsdp_meta.py#L96-L105

**The problematic line is `module.to_empty()`.**

FSDP supports nested wrapping, in which case initialization runs bottom-up. This means that for a nested case, the child materializes its meta-device parameters, and then the parent materializes its meta-device parameters. When the parent does so, it calls `parent_module.to_empty()`, which affects **all** parameters in `parent_module`, including those in `child_module`.

At this point in the execution, there are two choices, both of which are wrong. Either the specific `nn.Module.reset_parameters()` implementation affects all parameters in its tree (i.e. `.parameters(recurse=True)`), or it affects only immediately-owned parameters (i.e. `.parameters(recurse=False)`).
- [`recurse=False`] Any child FSDP instance's parameters are now left as uninitialized memory (like `torch.empty()`). Moreover, if an FSDP instance is applied to a module with submodules with parameters, then `recurse=False` will not include those submodules' parameters.
- [`recurse=True`] Any child FSDP instance's parameters are re-initialized. This is arguably "less wrong" than the `recurse=False` case, but it is still wrong. Re-initialization can silently advance the random seed, and random in-place initialization depends on parameter shape, which means that re-initializing parameters after flattening/sharding produces different results (thanks @min-xu-ai for flagging this to Fairscale users!).

A correct `param_init_fn` should skip parameters already processed by FSDP. However, we do not have any public API for detecting this, which means that writing a correct `param_init_fn` is impractical today.

Here is an example `param_init_fn` I wrote for `test_fully_shard.py`:
```
def param_init_fn(module: nn.Module):
    for submodule in module.modules():
        for param_name, param in submodule.named_parameters(recurse=False):
            if not _is_fsdp_flattened(param) and param.is_meta:
                materialized_param = nn.Parameter(
                    torch.empty_like(param, device=torch.device("cuda"))
                )
                nn.init.uniform_(materialized_param)
                setattr(submodule, param_name, materialized_param)
```
This is manually implementing the parameter swapping logic from the `module.to()` variants and requires the private `_is_fsdp_flattened()` function, making this unsatisfactory to me.

This issue is meant to track the progress on revisiting this meta device initialization. Ultimately, our goal is just to support a robust and usable deferred initialization API, which does not strictly require us to stay with this meta device initialization (e.g. we could migrate to PyTorch-native `FakeTensor` if that works).


cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501