Repository: pytorch
Issue: `inference_mode` before training results in FSDP AssertionError
Labels: triaged,module: fsdp

### 🐛 Describe the bug

```python
import os
import torch
import torch.cuda
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

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)

    with device:
        model = torch.nn.Linear(10, 10)
        input_ids = torch.randn(2, 10)
    model = FSDP(model)

    with torch.inference_mode():
        print("Validating ...")
        model(input_ids)
    print("Training ...")
    model(input_ids)

if __name__ == "__main__":
    mp.spawn(work, nprocs=2)
```

The resulting error is:

```python
Validating ...
Validating ...
Training ...
The `grad_fn` is needed to access the `AccumulateGrad` and register the post-backward hook
Training ...
The `grad_fn` is needed to access the `AccumulateGrad` and register the post-backward hook
  File "<string>", line 1, in <module>
  File "/usr/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "/usr/lib/python3.10/multiprocessing/spawn.py", line 129, in _main
    return self._bootstrap(parent_sentinel)
  File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 74, in _wrap
    fn(i, *args)
  File "/home/carlos/lit-parrot/repro.py", line 25, in work
    model(input_ids)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 849, in forward
    args, kwargs = _pre_forward(
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 416, in _pre_forward
    _register_post_backward_hook(state, handle)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 1419, in _register_post_backward_hook
    _p_assert(
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/utils.py", line 145, in _p_assert
    traceback.print_stack()
  File "<string>", line 1, in <module>
  File "/usr/lib/python3.10/multiprocessing/spawn.py", line 116, in spawn_main
    exitcode = _main(fd, parent_sentinel)
  File "/usr/lib/python3.10/multiprocessing/spawn.py", line 129, in _main
    return self._bootstrap(parent_sentinel)
  File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 74, in _wrap
    fn(i, *args)
  File "/home/carlos/lit-parrot/repro.py", line 25, in work
    model(input_ids)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 849, in forward
    args, kwargs = _pre_forward(
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 416, in _pre_forward
    _register_post_backward_hook(state, handle)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 1419, in _register_post_backward_hook
    _p_assert(
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/utils.py", line 145, in _p_assert
    traceback.print_stack()
Traceback (most recent call last):
  File "/home/carlos/lit-parrot/repro.py", line 28, in <module>
    mp.spawn(work, nprocs=2)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 246, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 202, in start_processes
    while not context.join():
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 163, in join
    raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException: 

-- Process 0 terminated with the following error:
Traceback (most recent call last):
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 74, in _wrap
    fn(i, *args)
  File "/home/carlos/lit-parrot/repro.py", line 25, in work
    model(input_ids)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 849, in forward
    args, kwargs = _pre_forward(
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 416, in _pre_forward
    _register_post_backward_hook(state, handle)
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/fsdp/_runtime_utils.py", line 1419, in _register_post_backward_hook
    _p_assert(
  File "/home/carlos/nightly-env/lib/python3.10/site-packages/torch/distributed/utils.py", line 147, in _p_assert
    raise AssertionError(s)
AssertionError: The `grad_fn` is needed to access the `AccumulateGrad` and register the post-backward hook
```

### Versions

`torch==2.2.0.dev20231102+cu121`

Initially reported in https://github.com/Lightning-AI/lit-gpt/issues/699

cc @zhaojuanmao @mrshenli @rohan-varma @awgu @fegin @penguinwu