[FSDP] Error when wrapping FSDP inside `checkpoint_wrapper`
Labels: oncall: distributed, triaged, module: fsdp

### 🐛 Describe the bug

This is the issue when running FSDP along with activation checkpointing. When FSDP units are wrapped inside `checkpoint_wrapper`, running checkpointing with both `NO_REENTRANT` and `REENTRANT` will fail. In particular, checkpointing with `NO_REENTRANT` since `_pre_forward_hook` is triggered ahead of `forward`, thus parameters are not correctly gather after `reshard`.

Below is an example for reproducing the error, this can be run via `gpurun torchrun --nnodes 1 --nproc_per_node 2 ./test_checkpoint.py`.
```python
import functools, os, torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import always_wrap_policy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
    apply_activation_checkpointing_wrapper,
)

class Model(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.layer0 = torch.nn.Linear(6, 6)
        self.layer1 = torch.nn.Linear(6, 6, bias=False)

    def forward(self, x):
        z = self.layer0(x)
        z = self.layer1(z)
        return z

    def get_input(self, device: torch.device):
        return (torch.randn((8, 6)).to(device),)

    def get_loss(self, input, output):
        return (output - input[0]).sum()


def fsdp_main():
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.distributed.init_process_group("nccl")
    torch.cuda.empty_cache()
    if torch.distributed.is_initialized():
        torch.cuda.set_device(local_rank)
    model = FSDP(
        Model(),
        auto_wrap_policy=always_wrap_policy,
        backward_prefetch=None,
        forward_prefetch=False,
        device_id=torch.cuda.current_device(),
    )
    non_reentrant_wrapper = functools.partial(
        checkpoint_wrapper,
        offload_to_cpu=False,
        checkpoint_impl=CheckpointImpl.NO_REENTRANT,
    )
    check_fn = lambda submodule: isinstance(submodule, FSDP)
    apply_activation_checkpointing_wrapper(
        model, checkpoint_wrapper_fn=non_reentrant_wrapper, check_fn=check_fn
    )
    if local_rank == 0:
        print(model)
    input = model.module.get_input(local_rank)
    output = model(*input)
    loss = model.module.get_loss(input, output).to(local_rank)
    loss.backward()


if __name__ == "__main__":
    fsdp_main()

```

### Versions

```
Collecting environment information...
PyTorch version: 1.13.0a0+git0fcdf93
Is debug build: False
CUDA used to build PyTorch: 11.1
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: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.22.1
Libc version: glibc-2.27

Python version: 3.9.12 (main, Apr  5 2022, 06:56:58)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1069-aws-x86_64-with-glibc2.27
Is CUDA available: False
CUDA runtime version: 11.1.105
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.3
[pip3] torch==1.10.1+cu111
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.10.1+cu111
[pip3] torchvision==0.11.2+cu111
[pip3] vit-pytorch==0.35.5
[conda] blas                      1.0                         mkl  
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] magma-cuda111             2.5.2                         1    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.22.3           py39he7a7128_0  
[conda] numpy-base                1.22.3           py39hf524024_0  
[conda] torch                     1.13.0a0+git0fcdf93           dev_0    <develop>
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.10.1+cu111             pypi_0    pypi
[conda] torchvision               0.11.2+cu111             pypi_0    pypi
[conda] vit-pytorch               0.35.5                   pypi_0    pypi
```


cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @pietern @SciPioneer @ezyang