Repository: pytorch
Issue: Memory leak in distributions.multivariate_normal.MultivariateNormal
Labels: high priority,needs reproduction,module: distributions,module: memory usage,triaged

### Description:

CPU memory is not released when sequentially training independent neural networks that output a distributions.multivariate_normal.MultivariateNormal object. This is in contrast to an equivalent procedure utilizing a distributions.normal.Normal object, where memory is released as expected after each training loop.

### Minimal example:
```
import os
import numpy as np
import torch
import torch.optim
import torch.nn as nn

class Model_N(nn.Module):
    # output of neural net is a normal.Normal distribution
    # otherwise equivalent to Model_MN
    def __init__(self):
        super().__init__()
        self.layers1 = nn.Sequential(nn.Linear(2,100), nn.ReLU())
        self.layers2 = nn.Sequential(nn.Linear(100,100), nn.ReLU(), nn.Linear(100,2))
        self.layers3 = nn.Sequential(nn.Linear(100,100), nn.ReLU(), nn.Linear(100,2))

    def forward(self, x):
        x = self.layers1(x)
        mu = self.layers2(x)
        sigma = self.layers3(x).exp().sqrt()
        dist = torch.distributions.normal.Normal(mu, sigma)
        return dist

class Model_MN(nn.Module):
    # output of neural net is a multivariate_normal.MultivariateNormal distribution
    # otherwise equivalent to Model_N
    def __init__(self):
        super().__init__()
        self.layers1 = nn.Sequential(nn.Linear(2,100), nn.ReLU())
        self.layers2 = nn.Sequential(nn.Linear(100,100), nn.ReLU(), nn.Linear(100,2))
        self.layers3 = nn.Sequential(nn.Linear(100,100), nn.ReLU(), nn.Linear(100,2))

    def forward(self, x):
        x = self.layers1(x)
        mu = self.layers2(x)
        cov = self.layers3(x).exp().diag_embed()
        dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, cov)
        return dist

def calc_loss_N(p):
    # log_prob result is equivalent but requires slightly different treatment between cases
    out = p.sample()
    loss = p.log_prob(out).sum(-1) * (out**2).sum(-1)
    return loss.mean()

def calc_loss_MN(p):
    # log_prob result is equivalent but requires slightly different treatment between cases
    out = p.sample()
    loss = p.log_prob(out) * (out**2).sum(-1)
    return loss.mean()

def train_loop(model, opt, loss_fn):
    model.train()
    for _ in range(1000):
        opt.zero_grad()
        inp = torch.ones((100,2))
        out = model(inp)
        loss = loss_fn(out)
        loss.backward()
        opt.step()
    return

seeds = np.arange(20)

print('Launching 20 sequential training loops with normal.Normal distribution as output of neural net. Results show expected behavior:')
_, used_m_init, _ = map(int, os.popen('free -t -m').readlines()[-1].split()[1:])
for seed in seeds:
    torch.manual_seed(seed)
    np.random.seed(seed)
    model = Model_N()
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    train_loop(model, opt, calc_loss_N)
    _, used_m, _ = map(int, os.popen('free -t -m').readlines()[-1].split()[1:])
    print('Excess memory at iter %d (MB): %d' % (seed+1, used_m-used_m_init))

print('Launching 20 sequential training loops with multivariate_normal.MultivariateNormal distribution as output of neural net. Results show unexpected behavior:')
_, used_m_init, _ = map(int, os.popen('free -t -m').readlines()[-1].split()[1:])
for seed in seeds:
    torch.manual_seed(seed)
    np.random.seed(seed)
    model = Model_MN()
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    train_loop(model, opt, calc_loss_MN)
    _, used_m, _ = map(int, os.popen('free -t -m').readlines()[-1].split()[1:])
    print('Excess memory at iter %d (MB): %d' % (seed+1, used_m-used_m_init))
```

### Output showing expected and unexpected behavior
```
Launching 20 sequential training loops with normal.Normal distribution as output of neural net. Results show expected behavior:
Excess memory at iter 1 (MB): 52
Excess memory at iter 2 (MB): 51
Excess memory at iter 3 (MB): 51
Excess memory at iter 4 (MB): 51
Excess memory at iter 5 (MB): 51
Excess memory at iter 6 (MB): 51
Excess memory at iter 7 (MB): 52
Excess memory at iter 8 (MB): 51
Excess memory at iter 9 (MB): 51
Excess memory at iter 10 (MB): 51
Excess memory at iter 11 (MB): 50
Excess memory at iter 12 (MB): 72
Excess memory at iter 13 (MB): 51
Excess memory at iter 14 (MB): 51
Excess memory at iter 15 (MB): 51
Excess memory at iter 16 (MB): 51
Excess memory at iter 17 (MB): 51
Excess memory at iter 18 (MB): 51
Excess memory at iter 19 (MB): 51
Excess memory at iter 20 (MB): 51
Launching 20 sequential training loops with multivariate_normal.MultivariateNormal distribution as output of neural net. Results show unexpected behavior:
Excess memory at iter 1 (MB): 52
Excess memory at iter 2 (MB): 102
Excess memory at iter 3 (MB): 152
Excess memory at iter 4 (MB): 202
Excess memory at iter 5 (MB): 252
Excess memory at iter 6 (MB): 300
Excess memory at iter 7 (MB): 350
Excess memory at iter 8 (MB): 400
Excess memory at iter 9 (MB): 450
Excess memory at iter 10 (MB): 501
Excess memory at iter 11 (MB): 548
Excess memory at iter 12 (MB): 599
Excess memory at iter 13 (MB): 648
Excess memory at iter 14 (MB): 698
Excess memory at iter 15 (MB): 748
Excess memory at iter 16 (MB): 798
Excess memory at iter 17 (MB): 848
Excess memory at iter 18 (MB): 898
Excess memory at iter 19 (MB): 947
Excess memory at iter 20 (MB): 1000
```

### Versions

Collecting environment information...
PyTorch version: 1.9.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: (GCC) 7.3.1 20180303 (Red Hat 7.3.1-5)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.8.11 (default, Aug  3 2021, 15:09:35)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.49.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.4.152
GPU models and configuration: 
GPU 0: Tesla V100-PCIE-32GB
GPU 1: Tesla V100-PCIE-32GB

Nvidia driver version: 495.29.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.9.1
[pip3] torchaudio==0.9.0a0+a85b239
[pip3] torchvision==0.10.1
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] numpy                     1.20.3           py38hf144106_0  
[conda] numpy-base                1.20.3           py38h74d4b33_0  
[conda] pytorch                   1.9.1           py3.8_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchaudio                0.9.1                      py38    pytorch
[conda] torchvision               0.10.1               py38_cu102    pytorch

cc @ezyang @gchanan @zou3519 @bdhirsh @fritzo @neerajprad @alicanb @nikitaved