Repository: pytorch
Issue: CheckpointError with checkpoint(..., use_reentrant=False) & autocast()
Labels: high priority,module: checkpoint,triaged,module: amp (automated mixed precision)

### 🐛 Describe the bug

When running the code below, I get the following error:

```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/workspaces/pytorch/grad/checkpoint.py", line 45, in main
    loss.backward()
  File "/opt/conda/lib/python3.11/site-packages/torch/_tensor.py", line 581, in backward
    torch.autograd.backward(
  File "/opt/conda/lib/python3.11/site-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/opt/conda/lib/python3.11/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/lib/python3.11/site-packages/torch/utils/checkpoint.py", line 1129, in unpack_hook
    frame.check_recomputed_tensors_match(gid)
  File "/opt/conda/lib/python3.11/site-packages/torch/utils/checkpoint.py", line 903, in check_recomputed_tensors_match
    raise CheckpointError(
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
tensor at position 5:
saved metadata: {'shape': torch.Size([64, 64]), 'dtype': torch.bfloat16, 'device': device(type='cpu')}
recomputed metadata: {'shape': torch.Size([6, 64]), 'dtype': torch.bfloat16, 'device': device(type='cpu')}
```

Note, this is a (very) minimized version of [this code](https://github.com/dptech-corp/Uni-Fold/blob/7b55789e8c27c83cbe3b2b71065299f9edc4c3d8/unifold/modules/template.py#L256-L340) from Uni-Fold, an open-source reimplementation of AlphaFold 2.

The essential parts seem to be
- Checkpointing with `checkpoint(..., use_reentrant=False)`
- Running a layer twice, once without gradients enabled (called "recycling" in the AF2 paper)
- Use of `torch.autocast`

```python3
import torch
from torch.utils.checkpoint import checkpoint


class PairStack(torch.nn.Module):
    def __init__(self, size: int):
        super().__init__()
        self.block = torch.nn.Sequential(
            torch.nn.LayerNorm(size),
            torch.nn.Linear(size, size)
        )

    def forward(self, x):
        # Issue also occurs if we unconditionally checkpoint()
        if torch.is_grad_enabled():
            return checkpoint(self.block, x, use_reentrant=False, debug=True)
        else:
            return self.block(x)


class Mod(torch.nn.Module):
    def __init__(self, size: int):
        super().__init__()
        self.stack = PairStack(size)
        self.linear = torch.nn.Linear(size, size)

    def forward(self, x):
        with torch.set_grad_enabled(False):
            x = self.stack(x)
        x = self.stack(x)
        return self.linear(x)


def main():
    device = torch.device("cpu")
    size = 64

    m = Mod(size=size).to(device)
    x = torch.linspace(0, 1, 2 * 3 * size).reshape(2, 3, size).to(device)

    with torch.autocast(device.type, dtype=torch.bfloat16):
        output = m(x)

    loss = output.sum()
    loss.backward()

if __name__ == "__main__":
    main()
```

adding `checkpoint(..., debug=True)` shows that part of the linear layer seems to be lost. IIUC, the linear layer consists of:

1. Reshape tensor into 2D
2. Matmul with weights + bias
3. Reshape tensor back into original shape

Strangely, (1) is still done during the recomputation, but (2) & (3) are not, leading to the shape mismatch.

```
Operations executed during the original forward:

('$3: bf16[2, 3, 64]', '$4: f32[2, 3, 1]', '$5: f32[2, 3, 1]') = torch._ops.aten.native_layer_norm.default($0, ['64'], $1, $2, 1e-05)
$6: bf16[6, 64] = torch._ops.aten.view.default($3, ['6', '64'])
$8: bf16[64, 64] = torch._ops.aten.t.default($7)
$10: bf16[6, 64] = torch._ops.aten.addmm.default($9, $6, $8)
$11: bf16[2, 3, 64] = torch._ops.aten.view.default($10, ['2', '3', '64'])

Operations executed during recomputation:

$1: f32[64] = torch._ops.aten.detach.default($0)
$2: f32[64] = torch._ops.aten.detach.default($1)
$4: f32[64] = torch._ops.aten.detach.default($3)
$5: f32[64] = torch._ops.aten.detach.default($4)
('$7: bf16[2, 3, 64]', '$8: f32[2, 3, 1]', '$9: f32[2, 3, 1]') = torch._ops.aten.native_layer_norm.default($6, ['64'], $3, $0, 1e-05)
$11: bf16[64] = torch._ops.aten._to_copy.default($10, dtype=torch.bfloat16)
$13: bf16[64, 64] = torch._ops.aten._to_copy.default($12, dtype=torch.bfloat16)
$14: bf16[6, 64] = torch._ops.aten.view.default($7, ['6', '64'])
$15: bf16[64, 64] = torch._ops.aten.t.default($13)
$16: bf16[6, 64] = torch._ops.aten.detach.default($14)
$17: bf16[6, 64] = torch._ops.aten.detach.default($16)
```

### Versions

This is running in the `pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel` Docker image. Currently testing on a M1 MacBook Pro, but I have reproduced the issue on an AMD CPU/Nvidia GPU server as well.

```
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.30.5
Libc version: glibc-2.35

Python version: 3.11.10 | packaged by conda-forge | (main, Oct 16 2024, 01:27:36) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-6.10.14-linuxkit-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: 12.4.131
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
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
Byte Order:                           Little Endian
CPU(s):                               8
On-line CPU(s) list:                  0-7
Vendor ID:                            Apple
Model:                                0
Thread(s) per core:                   1
Core(s) per socket:                   8
Socket(s):                            1
Stepping:                             0x0
BogoMIPS:                             48.00
Flags:                                fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 asimddp sha512 asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp flagm2 frint
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Vulnerable
Vulnerability Spectre v1:             Mitigation; __user pointer sanitization
Vulnerability Spectre v2:             Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==2.1.2
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] optree==0.13.0
[pip3] torch==2.5.1+cu124
[pip3] torchaudio==2.5.1+cu124
[pip3] torchelastic==0.2.2
[pip3] torchvision==0.20.1+cu124
[pip3] triton==3.1.0
[conda] numpy                     2.1.2           py311h71ddf71_0    conda-forge
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] optree                    0.13.0                   pypi_0    pypi
[conda] torch                     2.5.1+cu124              pypi_0    pypi
[conda] torchaudio                2.5.1+cu124              pypi_0    pypi
[conda] torchelastic              0.2.2                    pypi_0    pypi
[conda] torchvision               0.20.1+cu124             pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi
```

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @soulitzer @mcarilli @ptrblck @leslie-fang-intel @jgong5