Repository: pytorch
Issue: FSDP fails to load state dict under inference_mode
Labels: triaged,enhancement,inference mode,module: fsdp

### 🐛 Describe the bug

A runtime error occurs when attempting to load the state dict of an fsdp model under `torch.inference_mode()`:

```py
import os

import torch.cuda
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed.fsdp import FullyShardedDataParallel
from torch.nn.parallel import DistributedDataParallel


def work(rank):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "1234"
    dist.init_process_group("nccl", world_size=2, rank=rank)

    torch.cuda.set_device(rank)
    device = torch.device("cuda", rank)

    model = nn.Linear(100, 50).to(device)
    model = FullyShardedDataParallel(model)

    # no error with DDP
    # model = DistributedDataParallel(model, device_ids=[rank])

    with torch.inference_mode():
        model(torch.rand(2, 100))

    torch.save(model.state_dict(), "fsdp_model.pt")

    with torch.inference_mode():
        model.load_state_dict(torch.load("fsdp_model.pt"))
        model(torch.rand(2, 100))


def run():
    mp.spawn(work, nprocs=2)


if __name__ == "__main__":
    run()

```

Error:

```
Traceback (most recent call last):
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 69, in _wrap
    fn(i, *args)
  File "/home/adrian/repositories/lightning/repro.py", line 25, in work
    torch.save(model.state_dict(), "fsdp_model.pt")
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 2402, in state_dict
    with summon_ctx:
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 3002, in _summon_full_params
    stack.enter_context(self._fsdp_wrapped_module.unflatten_as_params())
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/contextlib.py", line 448, in enter_context
    result = _cm_type.__enter__(cm)
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/distributed/fsdp/flatten_params_wrapper.py", line 144, in unflatten_as_params
    with self._flat_param_handle.unflatten_as_params():
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 1019, in unflatten_as_params
    self._unflatten(as_params=True)
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/distributed/fsdp/flat_param.py", line 990, in _unflatten
    module.register_parameter(param_name, nn.Parameter(view))
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.9/site-packages/torch/nn/parameter.py", line 36, in __new__
    return torch.Tensor._make_subclass(cls, data, requires_grad)
RuntimeError: Setting requires_grad=True on inference tensor outside InferenceMode is not allowed.
```

This only happens with FSDP, and other wrappers like DDP don't experience this issue. 
PyTorch 1.13 and 2.0.


First encountered in PyTorch Lightning https://github.com/Lightning-AI/lightning/issues/16908

### Versions

PyTorch version: 1.13.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.9.12 (main, Jun  1 2022, 11:38:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-135-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.1.105
CUDA_MODULE_LOADING set to: LAZY

Nvidia driver version: 520.61.05
cuDNN version: Could not collect
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):                          64
On-line CPU(s) list:             0-63
Thread(s) per core:              2
Core(s) per socket:              16
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD EPYC 7302 16-Core Processor
Stepping:                        0
Frequency boost:                 enabled
CPU MHz:                         1499.605
CPU max MHz:                     3000.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        5999.81
Virtualization:                  AMD-V
L1d cache:                       1 MiB
L1i cache:                       1 MiB
L2 cache:                        16 MiB
L3 cache:                        256 MiB
NUMA node0 CPU(s):               0-15,32-47
NUMA node1 CPU(s):               16-31,48-63
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:          Vulnerable
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 conditional, RSB filling, PBRSB-eIBRS Not affected
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 pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 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 sme ssbd mba sev 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 wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca

Versions of relevant libraries:
[pip3] mypy==0.971
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.5
[pip3] torch==1.13.1
[pip3] torchmetrics==0.11.1
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] pytorch-cuda              11.8                 h7e8668a_3    pytorch-nightly
[conda] pytorch-triton            2.0.0+b8b470bc59          pypi_0    pypi
[conda] torch                     1.13.1                   pypi_0    pypi
[conda] torchmetrics              0.11.1                   pypi_0    pypi

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