Repository: pytorch
Issue: Enable jit error when using FSDP
Labels: oncall: jit

### 🐛 Describe the bug

Execuse me! Recently I trained GPT-2 model with enabling `torch.jit,` if I using DDP, it works normally. but when I using `FSDP`, it cause error. 

relevant codes as follows:

parameters config:
`config.use_torch_jit = True`
`config.use_cuda = True`
`config.dp_mode = "DDP" or "FSDP"`
`auto_wrap_policy = None`
`sharding_strategy = None`
`backward_prefetch = None`
`mixed_precision = None`
`cpu_offload = None`

```python
if config.use_torch_jit:
    logging.info("JIT optimization in progress...")
    fB = config.model.batch_size // config.dp
    fS = config.data.max_seq_length
    fV = config.model.vocab_size
    fake_input = torch.ones([fB, fS], dtype=torch.long)
    fake_words = torch.ones([fB, fS], dtype=torch.long)

    if config.use_cuda:
        fake_input = fake_input.cuda()
        fake_words = fake_words.cuda()

    gpt2_example_inputs = (
        fake_input,
        fake_words
    )
    gpt2 = torch.jit.trace(gpt2, gpt2_example_inputs, check_trace=False)
    fake_input = None
    fake_words = None
    logging.info("JIT optimization completed \033[0;32m[OK]")

dp_process_group = layout.group(0).process_group
if config.dp_mode == "DDP":
    model = DDP(gpt2, process_group=dp_process_group)
else:
    auto_wrap_policy = WARP_POLICES[config.fsdp.wrap_type]
    sharding_strategy = SHARDING_STRATEGYS[config.fsdp.sharding_strategy]
    backward_prefetch = BW_PREFETCH_TYPES[config.fsdp.bw_prefetch_type]
    if config.use_fp16_compress_dp:
        reduce_dtype = torch.float16
    all_dtypes = (param_dtype, buffer_dtype, reduce_dtype)
    if all(dtype is None for dtype in all_dtypes):
        mixed_precision = None
    else:
        mixed_precision = MixedPrecision(
            param_dtype=param_dtype,
            reduce_dtype=reduce_dtype,
            buffer_dtype=buffer_dtype
        )
    cpu_offload = None
    if config.fsdp.use_cpu_offload:
        cpu_offload = CPUOffload(offload_params=True)
    model = FSDP(gpt2, process_group=dp_process_group,
                 auto_wrap_policy=auto_wrap_policy,
                 sharding_strategy=sharding_strategy,
                 backward_prefetch=backward_prefetch,
                 mixed_precision=mixed_precision,
                 cpu_offload=cpu_offload)
```

Errors:

```python
    model = FSDP(gpt2, process_group=dp_process_group,
  File "/usr/local/conda/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 814, in __init__
    self._fsdp_wrapped_module: FlattenParamsWrapper = FlattenParamsWrapper(
  File "/usr/local/conda/lib/python3.9/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 329, in __init__
    self._flatten_params()
  File "/usr/local/conda/lib/python3.9/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 393, in _flatten_params
    delattr(m, n)
  File "/usr/local/conda/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1259, in __delattr__
    del self._parameters[name]
  File "/usr/local/conda/lib/python3.9/site-packages/torch/jit/_script.py", line 194, in __delitem__
    raise RuntimeError("cannot delete methods or parameters of a script module")
RuntimeError: cannot delete methods or parameters of a script module
```

### 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: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 7.3.1 20180303 (Red Hat 7.3.1-5)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.9.12 (main, Apr  5 2022, 06:56:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.18.0-147.mt20200626.413.el8_1.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.0.221
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 470.82.01
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.3
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.3
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.3
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.3
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.3
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.3
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.3
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.3
[pip3] torch==1.12.0
[pip3] torchaudio==0.12.0
[pip3] torchvision==0.13.0
[conda] blas                      1.0                         mkl    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] cudatoolkit               11.3.1               h2bc3f7f_2    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] ffmpeg                    4.3                  hf484d3e_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
[conda] mkl                       2021.4.0           h06a4308_640    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] mkl-service               2.4.0            py39h7f8727e_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] mkl_fft                   1.3.1            py39hd3c417c_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] mkl_random                1.2.2            py39h51133e4_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] numpy                     1.22.3           py39he7a7128_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] numpy-base                1.22.3           py39hf524024_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
[conda] pytorch                   1.12.0          py3.9_cuda11.3_cudnn8.3.2_0    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
[conda] pytorch-mutex             1.0                        cuda    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
[conda] torchaudio                0.12.0               py39_cu113    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
[conda] torchvision               0.13.0               py39_cu113    https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch