Repository: pytorch
Issue: non-rentrant checkpointing uses same memory as non-checkpointed code
Labels: high priority,oncall: distributed,module: checkpoint,triaged

### 🐛 Describe the bug

The newly addressed [checkpoint_without_reentrant](https://github.com/pytorch/pytorch/blob/master/torch/utils/checkpoint.py#L309) never clears the `storage` List. So it uses the same memory as non-checkpointed code.

Changing the `unpack` to should fix it:

```
    def unpack(x):
        nonlocal storage
        unpack_counter = 0
        if len(storage) == 0:
            def inner_pack(inner):
                nonlocal unpack_counter
                storage[unpack_counter] = inner
                unpack_counter += 1
                return None

            def inner_unpack(packed):
                raise RuntimeError("You are calling backwards on a tensor that is never exposed. Please open an issue.")

            # Stash the surrounding rng state, and mimic the state that was
            # present at this time during forward.  Restore the surrounding state
            # when we're done.
            rng_devices = []
            if preserve_rng_state and had_cuda_in_fwd:
                rng_devices = fwd_gpu_devices
            with torch.random.fork_rng(devices=rng_devices, enabled=preserve_rng_state):
                if preserve_rng_state:
                    torch.set_rng_state(fwd_cpu_state)
                    if had_cuda_in_fwd:
                        set_device_states(fwd_gpu_devices, fwd_gpu_states)
                with torch.enable_grad(), torch.cuda.amp.autocast(had_autocast_in_fwd):
                    with torch.autograd.graph.saved_tensors_hooks(inner_pack, inner_unpack):
                        _unused = function(*args, **kwargs)

        return storage.pop(x)
```


### Versions


Versions
```
Collecting environment information...
PyTorch version: 1.11.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
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: Could not collect
CMake version: version 3.10.2
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-1063-azure-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.4.152
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 470.82.01
cuDNN version: Could not collect
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.11.0+cu113
[pip3] torchviz==0.0.2
[conda] numpy                     1.22.3                   pypi_0    pypi
[conda] torch                     1.11.0+cu113             pypi_0    pypi
[conda] torchviz                  0.0.2                    pypi_0    pypi
```



 

cc @ezyang @gchanan @zou3519 @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang @kwen2501