Repository: pytorch
Issue: Is memory-efficient FSDP initialization intended to be possible with torch.device('meta')?
Labels: triaged,module: fsdp

### 🐛 Describe the bug

Looking at the FSDP API, I was expecting that we could use the torch meta device in conjunction with the `param_init_fn` and `sync_module_states` arguments to `FSDP` to instantiate a large model in a memory-efficient manner.

Specifically, I was imagining a workflow like:

- Only loading model parameters on the rank 0 device (load on the meta device on the others)
- Use a param_init_fn like module.to_empty(device=f'cuda:{rank}') when creating the FSDP module on non-rank 0 devices
- Use sync_module_states=True when creating the FSDP instance to sync parameter values with rank 0

For example:


    import contextlib
    import transformers
    import functools
    import torch
    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

    init_context = contextlib.nullcontext() if local_rank == 0 else torch.device('meta')
    with init_context:
        model = transformers.AutoModelForCausalLM.from_pretrained(
            'EleutherAI/pythia-160m', low_cpu_mem_usage=True)
    if local_rank == 0:
        print(model)

    model_auto_wrap_policy = functools.partial(transformer_auto_wrap_policy, transformer_layer_cls={model.gpt_neox.layers[0].__class__},)
    model = FSDP(model,
                 auto_wrap_policy=model_auto_wrap_policy,
                 sharding_strategy=ShardingStrategy.FULL_SHARD,
                 cpu_offload=CPUOffload(offload_params=False),
                 backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
                 device_id=local_rank,
                 sync_module_states=True,
                 param_init_fn=lambda mod: mod.to_empty(device=f'cuda:{local_rank}'))

But this code fails on the `FSDP()` line. It seems the [model sharding](https://github.com/pytorch/pytorch/blob/b689128db3564ca46c2e6f2f272653f3801a9b37/torch/distributed/fsdp/fully_sharded_data_parallel.py#L436) happens before modules on the meta device [are re-instantiated](https://github.com/pytorch/pytorch/blob/b689128db3564ca46c2e6f2f272653f3801a9b37/torch/distributed/fsdp/fully_sharded_data_parallel.py#L453). Is this behavior intentional?

### Versions

Collecting environment information...                              
PyTorch version: 2.0.1+cu117                                       
Is debug build: False                                    
CUDA used to build PyTorch: 11.7                                   
ROCM used to build PyTorch: N/A                                    
                                                                       
OS: Ubuntu 20.04.3 LTS (x86_64)                                    
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0                  
Clang version: Could not collect                                   
CMake version: version 3.26.3                                  
Libc version: glibc-2.31                                                                                                                       
                                                                                                                                               
Python version: 3.8.10 (default, Jun  2 2021, 10:49:15)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.29                                                                                 
Is CUDA available: True                                                                                                                        
CUDA runtime version: Could not collect                                                                                                        
CUDA_MODULE_LOADING set to: LAZY                               
GPU models and configuration:                                                                                                                  
GPU 0: NVIDIA A100-SXM4-40GB                                                                                                                   
GPU 1: NVIDIA A100-SXM4-40GB                                                                                                                   
                                                                                                                                               
Nvidia driver version: 515.43.04
cuDNN version: Probably one of the following:                                                                                                  
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/scr/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/scr/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/scr/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn.so.8.4.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.4.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.4.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.4.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.4.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.4.1
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.4.1
/usr/local/cuda-11.7/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.4.1
HIP runtime version: N/A                     
MIOpen runtime version: N/A                  
Is XNNPACK available: True 

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   43 bits physical, 48 bits virtual
CPU(s):                          256
On-line CPU(s) list:             0-255
Thread(s) per core:              2
Core(s) per socket:              64
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD EPYC 7662 64-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         3296.299
CPU max MHz:                     2154.2959
CPU min MHz:                     1500.0000
BogoMIPS:                        3999.91
Virtualization:                  AMD-V
L1d cache:                       4 MiB
L1i cache:                       4 MiB
L2 cache:                        64 MiB
L3 cache:                        512 MiB
NUMA node0 CPU(s):               0-63,128-191
NUMA node1 CPU(s):               64-127,192-255
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Mitigation; untrained return thunk; SMT enabled with STIBP protection
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP always-on, RSB filling
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr rdpru wbnoinvd amd_ppin arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif v_spec_ctrl umip rdpid overflow_recov succor smca sme sev sev_es

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.0.1
[pip3] triton==2.0.0
[conda] Could not collect

cc @zhaojuanmao @mrshenli @rohan-varma @awgu