Repository: pytorch
Issue: CUDA memory leak in multi-processing
Labels: module: performance,module: cuda,module: memory usage,triaged

## 🐛 CUDA memory leak in multi-processing

I'm working on parallelizing my training tasks (that are independent) in a multi-GPU setting. For this, I have created several workers (processes) that take a list of tasks and work parallelly to finish them. Each task is assigned with a device and any process that executes it will use the assigned device.

**The problem I observe is**: once a process finishes a task, the memory that was allocated in a GPU is not released. The same worker might work on another task that uses a different GPU and once it is done, it holds that memory too. In this way, a process occupies some parts of memory in all GPUs and never releases them until the process is terminated. This leads to memory leak and eventually results in CUDA out of memory error if several worker processes are used.

**Note:** I have tried using `torch.cuda.empty_cache()` to release the memory, only unsuccessful.

**However**, if each worker process is assigned a device to work with and any task that is executed by that process would always use its designated device, then this behavior is not observed. Since a worker process always uses the same GPU for its tasks, it seems like it releases/reuses the same allocated memory.

## To Reproduce

To reproduce this behavior, let's consider two settings :

- **Device on task**: Each task is assigned a device and any tensor operations are performed on the specified device irrespective of which worker executes it **(this results in memory leak)**

```python
from torch import multiprocessing as mp
import torch
import random


class TaskWithDevice:
    def __init__(self, device: str):
        self.device = device

    def run(self):
       rand_tensor = torch.rand(1000, 1000).to(self.device)


class Worker(mp.Process):
    def __init__(self, queue: mp.Queue):
        super().__init__(target=self.work, args=(queue,))

    def work(self, queue):
        while queue.qsize() > 0:
            task = queue.get()
            task.run()
            del task


if __name__ == '__main__':
    mp.set_start_method('spawn')
    n_tasks = 10000

    # available devices
    available_devices = [f"cuda:{i}" for i in range(torch.cuda.device_count())]
    devices = random.choices(available_devices, k=n_tasks)

    # task queue
    task_queue = mp.Queue()
    for i in range(n_tasks):
        task = TaskWithDevice(devices[i])
        task_queue.put(task)

    # workers
    workers = [Worker(task_queue) for device in available_devices]
    for worker in workers:
        worker.start()
    for worker in workers:
        worker.join()

```

- **Device on worker**: Each worker is assigned a device and any tasks that are executed by it would always use the same device **(here there is no memory leak and it seems like the allocated memory is reused for tasks by that worker)**

```python
from torch import multiprocessing as mp
import torch


class Task:
    def run(self, device: str):
        rand_tensor = torch.rand(1000, 1000).to(device)


class WorkerWithDevice(mp.Process):
    def __init__(self, queue: mp.Queue, device: str):
        self.device = device
        super().__init__(target=self.work, args=(queue,))

    def work(self, queue):
        while queue.qsize() > 0:
            task = queue.get()
            task.run(self.device)
            del task


if __name__ == '__main__':
    mp.set_start_method('spawn')
    n_tasks = 10000

    # available devices
    available_devices = [f"cuda:{i}" for i in range(torch.cuda.device_count())]

    # task queue
    task_queue = mp.Queue()
    for i in range(10000):
        task = Task()
        task_queue.put(task)

    # workers attached to a particular device
    workers = [WorkerWithDevice(task_queue, device) for device in available_devices]
    for worker in workers:
        worker.start()
    for worker in workers:
        worker.join()

```

## Expected behavior

In the first setting, it is expected that once the task is finished, the GPU memory should be released as soon as the program finishes `run()`

## Environment

PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.2 LTS (x86_64)
GCC version: (Ubuntu 7.4.0-1ubuntu1~18.04.1) 7.4.0
Clang version: Could not collect
CMake version: version 3.10.2

Python version: 3.7 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB

Nvidia driver version: 450.51.06
cuDNN version: Could not collect



cc @ngimel @VitalyFedyunin