Repository: pytorch
Issue: BAR1 memory of GPU is not released when main process is killed.
Labels: module: dataloader,module: cuda,module: memory usage,triaged

## 🐛 Bug

When the main process receives a `SIGABRT`, `SIGTERM` or `SIGKILL` signal while
iterating through a `DataLoader` the `BAR1 memory` of the GPU is not released.
This happens only with `num_workers>0`.

The [doc-string](https://github.com/pytorch/pytorch/blob/213c3f45da4bf9c862095cf05a55df0ae30214d3/torch/utils/data/dataloader.py#L570-L875) of `_MultiProcessingDataLoaderIter` mentions that `Terminating multiprocessing logic requires very careful design.` However, the `FB memory` of the GPU is released as expected and all processes are terminated.

Is there a way to make sure everything is cleaned up properly or is this the expected behavior?


## To Reproduce

```python
import torch
import os
from torch.utils.data import DataLoader, Dataset


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


if __name__ == '__main__':
    pid = os.getpid()
    data = DataLoader(RandomDataset(32, 64), batch_size=2, num_workers=1)

    for i in range(100):
        print(i)

        for batch in data:
            batch = batch.cuda(0)

            if i == 40:
                # kill the process after an arbitrary number of epochs
                # SIGABRT, SIGTERM and SIGKILL lead to same behaviour.
                # SIGINT (CTRL-C) and sys.exit() works fine
                os.system(f"kill -s 9 {pid}")
```
Memory before run (output from nvidia-smi -q -d MEMORY)
```
Driver Version                            : 460.73.01
CUDA Version                              : 11.2

Attached GPUs                             : 14
GPU 00000000:1B:00.0
    FB Memory Usage
        Total                             : 15109 MiB
        Used                              : 0 MiB
        Free                              : 15109 MiB
    BAR1 Memory Usage
        Total                             : 256 MiB
        Used                              : 230 MiB
        Free                              : 26 MiB
```
Memory after the main process is killed
```
Driver Version                            : 460.73.01
CUDA Version                              : 11.2

Attached GPUs                             : 14
GPU 00000000:1B:00.0
    FB Memory Usage
        Total                             : 15109 MiB
        Used                              : 0 MiB
        Free                              : 15109 MiB
    BAR1 Memory Usage
        Total                             : 256 MiB
        Used                              : 234 MiB
        Free                              : 22 MiB
```
## Expected behavior

`BAR1 memory` is released when main process receives `SIGABRT`, `SIGTERM` or `SIGKILL` signal.

## Environment
```
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 10 (buster) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: version 3.13.4
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.19.0-17-amd64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla T4
GPU 1: Tesla T4
GPU 2: Tesla T4
GPU 3: Tesla T4
GPU 4: Tesla T4
GPU 5: Tesla T4
GPU 6: Tesla T4
GPU 7: Tesla T4
GPU 8: Tesla T4
GPU 9: Tesla T4
GPU 10: Tesla T4
GPU 11: Tesla T4
GPU 12: Tesla T4
GPU 13: Tesla T4

Nvidia driver version: 460.73.01
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.910
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.2
[pip3] pytest-mypy==0.8.1
[pip3] pytorch-lightning==1.4.2
[pip3] pytorch-ranger==0.1.1
[pip3] torch==1.9.0
[pip3] torch-fidelity==0.3.0
[pip3] torch-optimizer==0.1.0
[pip3] torchaudio==0.9.0a0+33b2469
[pip3] torchmetrics==0.5.0
[pip3] torchtext==0.10.0
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.1.74              h6bb024c_0    nvidia
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py38h27cfd23_1  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] mypy                      0.910                    pypi_0    pypi
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.20.2           py38h2d18471_0  
[conda] numpy-base                1.20.2           py38hfae3a4d_0  
[conda] pytest-mypy               0.8.1                    pypi_0    pypi
[conda] pytorch                   1.9.0           py3.8_cuda11.1_cudnn8.0.5_0    pytorch
[conda] pytorch-lightning         1.4.2                    pypi_0    pypi
[conda] pytorch-ranger            0.1.1                    pypi_0    pypi
[conda] torch-fidelity            0.3.0                    pypi_0    pypi
[conda] torch-optimizer           0.1.0                    pypi_0    pypi
[conda] torchaudio                0.9.0                      py38    pytorch
[conda] torchmetrics              0.5.0                    pypi_0    pypi
[conda] torchtext                 0.10.0                     py38    pytorch
[conda] torchvision               0.10.0               py38_cu111    pytorch
```

## Additional context

This error may occur when using parallelized hyperparameter search (e.g. [Ray Tune](https://docs.ray.io/en/latest/tune/index.html)) with aggressive early stopping criteria. 


cc @SsnL @VitalyFedyunin @ejguan @NivekT @ngimel