Repository: pytorch
Issue: [FSDP] deepcopy FSDP model for EMA results in error
Labels: oncall: distributed,module: fsdp

### 🐛 Describe the bug

The model I want to train is more stably trained with EMA, I want to apply the FSDP to the model so as to train with much larger model sizes, with code like the below:

```python
import copy
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.fsdp.fully_sharded_data_parallel import (
    FullyShardedDataParallel as FSDP,
)
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

dist.init_process_group(backend="nccl", init_method="env://")
model = nn.Linear(3, 1024)
fsdp_model = FSDP(model, device_id=0, auto_wrap_policy=size_based_auto_wrap_policy)
fsdp_model_ema = copy.deepcopy(fsdp_model)
```

I run into the following problem:
```
(base) root@ccab4f7a6a8b4707903b4b491c065374000001:~/scripts# torchrun --nproc_per_node 1 bug_repro.py 
Traceback (most recent call last):
  File "bug_repro.py", line 12, in <module>
    fsdp_model_ema = copy.deepcopy(fsdp_model)
  File "/opt/conda/lib/python3.7/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/opt/conda/lib/python3.7/copy.py", line 281, in _reconstruct
    state = deepcopy(state, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 150, in deepcopy
    y = copier(x, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 241, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/opt/conda/lib/python3.7/copy.py", line 307, in _reconstruct
    value = deepcopy(value, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/opt/conda/lib/python3.7/copy.py", line 281, in _reconstruct
    state = deepcopy(state, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 150, in deepcopy
    y = copier(x, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 241, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 180, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/opt/conda/lib/python3.7/copy.py", line 307, in _reconstruct
    value = deepcopy(value, memo)
  File "/opt/conda/lib/python3.7/copy.py", line 161, in deepcopy
    y = copier(memo)
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/parameter.py", line 56, in __deepcopy__
    result = type(self)(self.data.clone(memory_format=torch.preserve_format), self.requires_grad)
  File "/opt/conda/lib/python3.7/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 106, in __new__
    raise ValueError("An non-empty list or tuple argument is needed")
ValueError: An non-empty list or tuple argument is needed
```

Is this expected? What would be the suggested way to implement EMA with FSDP model?

### Versions

PyTorch version: 1.12.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.7.13 (default, Mar 29 2022, 02:18:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1063-azure-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 470.82.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.6
[pip3] pytorch-lightning==1.6.4
[pip3] torch==1.12.0
[pip3] torchaudio==0.12.0
[pip3] torchmetrics==0.9.2
[pip3] torchvision==0.13.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py37h7f8727e_0  
[conda] mkl_fft                   1.3.1            py37hd3c417c_0  
[conda] mkl_random                1.2.2            py37h51133e4_0  
[conda] numpy                     1.21.6                   pypi_0    pypi
[conda] numpy-base                1.21.5           py37ha15fc14_3  
[conda] pytorch                   1.12.0          py3.7_cuda11.3_cudnn8.3.2_0    pytorch
[conda] pytorch-lightning         1.6.4                    pypi_0    pypi
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch                     1.12.0                   pypi_0    pypi
[conda] torchaudio                0.12.0               py37_cu113    pytorch
[conda] torchmetrics              0.9.2                    pypi_0    pypi
[conda] torchvision               0.13.0                   pypi_0    pypi

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501 @ezyang