Repository: pytorch
Issue: Abnormal GPU memory usage when using CUDA tensors with multiprocessing
Labels: module: multiprocessing,module: cuda,module: memory usage,triaged,module: jiterator

### 🐛 Describe the bug

Hi! I wanted to do some inference with a network in multiple processes, so I followed this and used `torch.multiprocessing` and `share_memory`, hoping all of the processes can share the same network to avoid copying the network for each process which may cause OOM on GPU.

However, I found that the `share_memory` didn't work. I ran the script below and monitored the GPU memory usage, finding that the GPU memory usage can be up to ~ 4GB with number of processes equaled to 4, which was abnormal for a quite simple linear layer with shape `(3, 5)`. If I increase the number of processes, the GPU memory usage will grow linearly.

This issue happens for both of the `spawn` and `forkserver` multiprocessing starting method which are needed for passing CUDA tensors as arguments to multi-processed functions.

Code to reproduce the issue:

```python
import torch
import torch.nn as nn
import torch.multiprocessing as mp
from functools import partial

def train(task, net):
    print(f'task {task} : {list(net.parameters())[0].device}')
    # x = torch.rand(1, 3).cuda(3)
    # y = net(x)
    # print(y.shape)


def main():
    mp.set_start_method('spawn')
    net = nn.Linear(3, 5).cuda(3)
    net = net.share_memory()
    p_train = partial(train, net=net)
    
    with mp.Pool(
        processes=8,
        maxtasksperchild=1,
    ) as pool:
        pool.map(p_train, list(range(16)), chunksize=1)
    

if __name__ == '__main__':
    main()
```

I tried to make the network global and use `fork` starting method, then the GPU memory usage is normal (takes about ~1GB, which is equal to the program with the same network but without any multiprocessing). However, I couldn't create CUDA tensors inside sub-processes with `fork`, which makes this method usefulness for most of the cases (after all we need to compute something rather than just print it out).

```python
import torch
import torch.nn as nn
import torch.multiprocessing as mp
from functools import partial

net: nn.Module

def train(task):
    global net
    print(f'task {task} : {net} {list(net.parameters())[0].device}')
    x = torch.rand(1, 3).cuda(3) # not allowed
    # RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
    y = net(x)
    print(y.shape)


def main():
    mp.set_start_method('fork')
    global net
    net = nn.Linear(3, 5).cuda(3)
    # net = net.share_memory()
    p_train = partial(train)#, net=net)
    
    with mp.Pool(
        processes=8,
        maxtasksperchild=1,
    ) as pool:
        pool.map(p_train, list(range(16)), chunksize=1)
    

if __name__ == '__main__':
    main()
```


### Versions

```
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 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.10.4 (main, Mar 31 2022, 08:41:55) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1045-aws-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: Tesla V100-SXM2-16GB
GPU 1: Tesla V100-SXM2-16GB
GPU 2: Tesla V100-SXM2-16GB
GPU 3: Tesla V100-SXM2-16GB

Nvidia driver version: 460.91.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] torch==1.11.0+cu113
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] torch                     1.11.0+cu113             pypi_0    pypi
```

cc @VitalyFedyunin @ngimel @mruberry