Repository: pytorch
Issue: CUDACachingAllocator is not GC-aware
Labels: module: cuda,module: memory usage,triaged

## 🐛 Bug

Dead objects holding `Tensor`s are considered active memory by `CUDACachingAllocator`. Under GPU memory pressure, the allocator makes no attempt to run a garbage collection cycle on dead but as yet uncollected `Tensor`s. As a result, the allocator has no recourse but to raise an OOM exception despite having potentially significant amounts of memory available to service the allocation request.

## To Reproduce

Run the following code on an 11GB GPU and observe an OOM exception:

```python
import gc
import torch


class Garbage:
  def __init__(self, obj):
    self.obj = obj
    self.cycle = self


gc.collect()
for i in range(100):
  Garbage(torch.empty(1024 * 100, 1024).cuda())
  print(i)
```

Sample output:
```
0
[...snip...]
25
Traceback (most recent call last):
  File "oom.py", line 13, in <module>
    Garbage(torch.empty(1024 * 100, 1024).cuda())
RuntimeError: CUDA out of memory. Tried to allocate 400.00 MiB (GPU 0; 10.92 GiB total capacity; 10.16 GiB already allocated; 287.06 MiB free; 10.16 GiB reserved in total by PyTorch)
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

This code snippet should execute successfully as there are no live objects holding GPU memory at the beginning of each loop iteration.

## Environment

```
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 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

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti
GPU 2: GeForce GTX 1080 Ti
GPU 3: GeForce GTX 1080 Ti

Nvidia driver version: 440.33.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.5.1
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] haste-pytorch==0.5.0rc0
[pip3] numpy==1.18.5
[pip3] torch==1.6.0
[pip3] torchaudio==0.6.0
[pip3] torchvision==0.7.0
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] haste-pytorch             0.5.0rc0                 pypi_0    pypi
[conda] mkl                       2020.1                      217
[conda] mkl-service               2.3.0            py38he904b0f_0
[conda] mkl_fft                   1.1.0            py38h23d657b_0
[conda] mkl_random                1.1.1            py38h0573a6f_0
[conda] numpy                     1.18.5                   pypi_0    pypi
[conda] pytorch                   1.6.0           py3.8_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] torchaudio                0.6.0                    pypi_0    pypi
[conda] torchvision               0.7.0                    pypi_0    pypi
```

In addition to PyTorch 1.6, I have reproduced this failure mode with the same machine configuration as above with PyTorch 1.7.1 and commit `ee80b45` on `master` (corresponding to version `1.7.0a0+57bffc3`).

## Additional context

Under normal circumstances, `Tensor`s are freed immediately when their refcount goes to 0. There are some situations in which reference cycles exist (e.g. self-reference) where one or more of the objects participating in the cycle are also holding a reference to a `Tensor`. The memory held by that `Tensor` will be considered allocated until the garbage collector runs.

Unfortunately, I ran into exactly this problem when passing around `Tensor` objects inside a structure provided by a 3rd party library. My training loop would slowly accumulate garbage until OOMing. Further analysis showed that I had approximately 3.5GiB of dead (but collectible) GPU memory at the beginning of the training step that failed.

I've implemented a fix which calls the garbage collector when `CUDACachingAllocator` is under memory pressure and has exhausted all other means of servicing the allocation request. This solution imposes no additional performance overhead since it invokes a collection only if the allocator's next fallback option was to raise an OOM exception. Running the above code snippet with the patch succeeds, as expected. PR coming shortly.

Adding some folks who might have an opinion on this: @colesbury, @dojoteef, @ezyang, @ngimel, @vadimkantorov

cc @ngimel