Repository: pytorch
Issue: PyTorch gradient checkpointing consumes a lot more memory than JAX
Labels: module: checkpoint,module: memory usage,triaged

### 🐛 Describe the bug

Here is the sample code for Pytorch and JAX: 
```python
import torch
from torch.utils.checkpoint import checkpoint
import memory_profiler

torch.manual_seed(0)

def mul(A, B, chkpt_size):
    for i in range(chkpt_size):
        B = A @ B
        B = B / torch.norm(B)
    return B

@profile
def func():
    size_ = 200
    num_steps = 10000
    A = torch.randn(size_, size_, dtype=torch.float64)
    A = A / torch.norm(A)
    B = torch.randn(size_, size_, dtype=torch.float64)
    B = B / torch.norm(B)

    lambda_ = torch.tensor([0.], dtype=torch.float64, requires_grad=True)

    C = A + lambda_ * B

    D = torch.randn(size_, size_, dtype=torch.float64)

    # for i in range(10000):
    #     D = C @ D
    #     D = D / torch.norm(D)
    chkpt_size = 20
    num_chkpt = num_steps // chkpt_size
    for i in range(num_chkpt):
        D = checkpoint(mul, C, D, chkpt_size, use_reentrant=False)
    E = torch.norm(B.t() @ D)
    print(E)

    E.backward()
    print(lambda_.grad)
    return

func()
```
This code is essentially taking the derivative of $(A+\lambda B)^n D$ with respect to $\lambda$ (but normalize every step). I also wrote a code that does the same thing in JAX:
```python
import jax
import jax.random as random
import jax.numpy as jnp
import jax.lax as lax
import memory_profiler

jax.config.update("jax_enable_x64", True)


def mul_jax(A, B, chkpt_size):
    def matmul_update(B, _):
        new_B = A @ B  
        new_B = new_B / jnp.linalg.norm(new_B)
        return new_B, new_B 
    final_B, B_list = lax.scan(matmul_update, B, None, length=chkpt_size)
    return final_B


def func(lambda_):
    seed = 0
    num_steps = 10000
    key = random.PRNGKey(seed)
    key, subkey = random.split(key)
    A = random.normal(subkey, (200, 200))
    A = A / jnp.linalg.norm(A)

    key, subkey = random.split(key)
    B = random.normal(subkey, (200, 200))
    B = B / jnp.linalg.norm(B)

    C = A + lambda_ * B

    key, subkey = random.split(key)
    D = random.normal(subkey, (200, 200))
    checkpoint = True

    if checkpoint:
        chkpt_size = 20
        num_chkpt = num_steps// chkpt_size
        def wrap_mul_jax(D, _):
            return mul_jax(C, D, chkpt_size), None
        chkpted_wrap_mul_jax = jax.checkpoint(wrap_mul_jax)
        D, _ = lax.scan(chkpted_wrap_mul_jax, D, None, length=num_chkpt)

    else:
        D = mul_jax(C, D, num_steps)
    
    E = jnp.linalg.norm(B.T @ D)
    return E

@profile
def main():
    lambda_ = 0.0
    primals, vjp_fun = jax.vjp(func, lambda_)
    grad = vjp_fun(1.0)[0]
    print(grad)

main()
```
They consume similar amount of memory when gradient checkpointing is turned off, but it seems that Pytorch consumes a lot more memory when the gradient checkpointing is turned on. There is a plot that shows the difference:
![memory](https://github.com/pytorch/pytorch/assets/66626987/8b794f2f-fbbc-4fb8-ae99-40d9896354ec)

I wonder why this happens and how it can be improved. Thank you!

### Versions

```
Collecting environment information...
PyTorch version: 2.2.0.post100
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.4.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.27.5
Libc version: N/A

Python version: 3.11.4 (main, Jul  5 2023, 08:54:11) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.4.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] flake8==7.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] numpydoc==1.5.0
[pip3] torch==2.2.0.post100
[pip3] torch-tb-profiler==0.4.3
[pip3] torchviz==0.0.2
[conda] numpy                     1.24.3          py311hb57d4eb_0  
[conda] numpy-base                1.24.3          py311h1d85a46_0  
[conda] numpydoc                  1.5.0           py311hca03da5_0  
[conda] pytorch                   2.2.0           gpu_mps_py311hf322ab5_100  
[conda] torch-tb-profiler         0.4.3                    pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi
```

cc @soulitzer