Repository: pytorch
Issue: CPU Memory Leak for JIT ScriptModule on DataParallel
Labels: oncall: jit,triaged

## 🐛 Bug

I am encountering CPU memory leakage issues when training my custom torchscript RNN models in a DataParallel environment. I am implementing custom RNN cells, using torchscript for speedups. These models perform great in a single-GPU environment, and torchscript gives me an ~2.0x speedup. However, when I deploy my models on multiple GPUs using DataParallel, I encounter a CPU memory leak. This leak eventually leads to an overflow that kills the training.

I believe that I have determined that the issue stems from torch internal rather than from my own code. A similar issue page https://github.com/pytorch/pytorch/issues/4865, provides the advice that DataParallel leads to CPU leakage when there are cycles in data structures. I wanted to distill whether the issue comes from my own custom RNN layers (which are fairly complex) or if it originates in the interaction between torchscript and DataParallel.

To this end, I implemented a simple Linear layer as both a `jit.ScriptModule` and a `nn.Module`. This serves as a minimal stand-in for my more complicated RNN layers to replicate the issue that I am encountering. I ran these Linear layers in inference mode (just forward pass, no training) as DataParallel on a g3.8xlarge AWS instance (2 NVIDIA Tesla M60s). 

As you will see below, I have validated that the interaction between torchscript and DataParallel leads to CPU memory leakage at a constant rate, independent of model/minibatch size. Although the models explored in this study are toy models, I believe that this gets to the core of the issue that I am experiencing.


## To Reproduce

Steps to reproduce the behavior:

1. Implement Linear layer as `nn.Module` and `jit.ScriptModule`:
```python
import torch
from torch import nn, jit
import math
import psutil # for tracking CPU memory usage

# Regular nn.Module implementation of a Linear layer
class Linear(nn.Module):
    def __init__(self, Din, Dout):
        super(Linear, self).__init__()
        
        self.weight = nn.Parameter(torch.Tensor(Din, Dout))
        self.bias = nn.Parameter(torch.Tensor(Dout))
        
        stdv = 1.0/math.sqrt(0.5*(Dout+Din))
        
    def forward(self, x):
        return torch.mm(x, self.weight) + self.bias

# Torchscript implementation of a Linear Layer
class jitLinear(jit.ScriptModule):
    def __init__(self, Din, Dout):
        super(jitLinear, self).__init__()
        
        self.weight = nn.Parameter(torch.Tensor(Din, Dout))
        self.bias = nn.Parameter(torch.Tensor(Dout))
        
        stdv = 1.0/math.sqrt(0.5*(Dout+Din))
        
    @jit.script_method
    def forward(self, x):
        # type: (Tensor) -> (Tensor)
        return torch.mm(x, self.weight) + self.bias
```
2. Define a simple inference loop that measures CPU memory
```python
def inference(model, batch_size, Din, device='cuda', N=10000, Nprint=500):
    model.train(False)
    cpu_memory = []
    for i in range(N):
        x = torch.randn(batch_size, Din, device=device)
        y = model(x)
        # measure CPU memory
        memory = psutil.virtual_memory().used*1e-9
        cpu_memory.append(memory)
        if i % Nprint == 0:
            print('Memory: {:.3f} Gb'.format(memory))
    return cpu_memory
```
3. Run inference loop to measure CPU memory for different combinations of layer implementation ('nn.Module' or 'jit.ScriptModule') and multi-GPU (DataParallel or not).
```python
# Batch and model size
B = 1024
Din = 1024
Dout = 1024

#-----nn.Module------
model = Linear(Din, Dout).to('cuda')
nn_memory = inference(model, B, Din)

#-----nn.Module + DataParallel-----
model = Linear(Din, Dout).to('cuda')
model = nn.DataParallel(model)
nn_DP_memory = inference(model, B, Din)

#-----jit.ScriptModule-----
model = jitLinear(Din, Dout).to('cuda')
jit_memory = inference(model, B, Din)

#-----jit.ScriptModule + DataParallel-----
model = jitLinear(Din, Dout).to('cuda')
model = nn.DataParallel(model)
jit_DP_memory = inference(model, B, Din)

#-----jit.ScriptModule + DataParallel, larger minibatch-----
# batch size of 2048 instead of 1024
model = jitLinear(Din, Dout).to('cuda')
model = nn.DataParallel(model)
large_minibatch_jit_DP_memory = inference(model, 2048, Din)

#-----jit.ScriptModule + DataParallel, larger model-----
# Din=2048 instead of 1024
# Dout = 2048 instead of 1024
model = jitLinear(2048, 2048).to('cuda')
model = nn.DataParallel(model)
large_model_jit_DP_memory = inference(model, 1024, 2048)
```

![image](https://user-images.githubusercontent.com/11778463/69099045-b0305980-0a0e-11ea-8807-8a26eb9a0cd5.png)


## Expected behavior

I expect the CPU memory to be constant throughout the inference loop. I observe the expected behavior for all combinations except for those that combine torchscript and DataParallel. An interesting observation is that the rate of CPU memory leakage does not seem to be correlated with model size or minibatch size. For all combinations of torchscript and DataParallel, memory leaked at a rate of ~148kb/minibatch.

## Environment

PyTorch version: 1.3.0
Is debug build: No
CUDA used to build PyTorch: 10.1.243

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
CMake version: version 3.13.3

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.0.130
GPU models and configuration: 
GPU 0: Tesla M60
GPU 1: Tesla M60

Nvidia driver version: 418.87.00
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.15.4
[conda] _pytorch_select           0.2                       gpu_0  
[conda] blas                      1.0                         mkl  
[conda] mkl                       2019.4                      243  
[conda] mkl-service               2.3.0            py36he904b0f_0  
[conda] mkl_fft                   1.0.14           py36ha843d7b_0  
[conda] mkl_random                1.0.2            py36hd81dba3_0  
[conda] pytorch                   1.3.1           py3.6_cuda10.1.243_cudnn7.6.3_0    pytorch
[conda] torch                     1.3.0                     <pip>
[conda] torchaudio                0.3.1                     <pip>
[conda] torchvision               0.4.2                py36_cu101    pytorch


## Additional context

<!-- Add any other context about the problem here. -->


cc @suo