Repository: pyro-ppl/pyro

Issue:
Memory leak using TraceEnum_ELBO
I noticed a major memory leak when training SVI using `TraceEnum_ELBO`.
I initially noticed this in a custom model we are developing but then I found it seems a more general bug.

For example, it affects even the Pyro tutorials GMM example [here](https://pyro.ai/examples/gmm.html). Where memory usage rapidly goes from a couple of hundred MBs to a many **GB**s very quickly!

I have run this Macbook Pro 2019 running MacOS 10.15. To replicate the issue is enough running the notebook linked.

I have tried to comment out the following lines and add a garbage collector call, that reduces the entity of the memory accumulation of one order of magnitude but does not solve the problem completely, which becomes particularly severe for large datasets.

```python
# Register hooks to monitor gradient norms.
# gradient_norms = defaultdict(list)
# for name, value in pyro.get_param_store().named_parameters():
#     value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))

import gc
losses = []
for i in range(200000):
    loss = svi.step(data)
    #losses.append(loss)
    gc.collect()
``` 

(from [this forum post](https://forum.pyro.ai/t/memory-usage-grows-with-iterations-in-svi-step/4103))
[bug] Memory leak on GPU
### Issue Description
After successfully following the SS-VAE example in pyro I have moved over to training a variant of it on a GPU. A memory leak has arisen that has been boiled down to a minimum working example that replicates the error. I am at a loss as to how to proceed in rectifying this issue, any help or insight would be much appreciated!

### Environment
 - OS and python version. 
   - OS: Red Hat Enterprise Linux Workstation release 7.9 (Maipo) 
   - Python: 3.9.7
   - CUDA Version: 11.2
 - PyTorch version: 1.10.0
 - Pyro version: 1.7.0

### Code Snippet
``` python
from datetime import datetime

import pyro
import pyro.distributions as dist
import torch
import torch.nn as nn
from pyro.infer import SVI, TraceEnum_ELBO
from pyro.optim import Adam


class SSVAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.accept_ratio = .5
        ones = torch.ones((), device='cuda')
        ratio = 1 / (self.accept_ratio + 1)
        self.alpha_prior = torch.stack((ones - ratio, ones - self.accept_ratio * ratio)).T
        self.cuda()

    def model(self):
        pyro.sample("y", dist.OneHotCategorical(self.alpha_prior))

    def guide(self):
        pass


def sizeof_fmt(num, suffix="B"):
    for unit in ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi"]:
        if abs(num) < 1024.0:
            return f"{num:3.1f}{unit}{suffix}"
        num /= 1024.0
    return f"{num:.1f}Yi{suffix}"


if __name__ == '__main__':
    torch.cuda.empty_cache()
    ssvae = SSVAE()

    optimizer = Adam({"lr": 1e-4})

    svi = SVI(ssvae.model, ssvae.guide, optimizer, loss=TraceEnum_ELBO())

    for i in range(10):
        print(f"{datetime.now().strftime('%H:%M:%S')} - Starting epoch #{i + 1}")
        print(
            f"Current GPU memory allocation:"
            f" {sizeof_fmt(torch.cuda.memory_allocated())}/{sizeof_fmt(torch.cuda.max_memory_allocated())}")

        # get the losses for an epoch
        for _ in range(100):
            svi.step()
```



Discussion:
@gioelelm could you check to see if #3069 fixes your issue?
Thank you for the quick attempt, but no it does not fix the problem. Neither the one without using `gc.collect()` nor the residual leak when garbage collecting.
Thanks for checking @gioelelm. I might have time in the next few weeks to dive deeper.  If you have time I can recommend some strategies (what I'd try):
- get an idea of which tensors are leaking using [this trick](https://forum.pyro.ai/t/a-trick-to-debug-tensor-memory/556)
- try to determine which objects might be holding references to the leaking tensors using something like
    ```python
    elbo = TraceEnum_ELBO()
    optim = ClippedAdam(...)
    svi = SVI(model, guide, optim, elbo)
    for step in range(steps):
        svi.step()
        print("svi", len(pickle.dumps(svi)))
        print("elbo", len(pickle.dumps(elbo)))
        print("optim", len(pickle.dumps(optim)))
        print("param_store", len(pickle.dumps(pyro.get_param_store()))
   ```
- See if this is a recent PyTorch bug by trying inference with different torch versions, say 1.11, 1.10, 1.9. 1.8. I'm pretty sure the GMM tutorial should still work with older PyTorch versions.
Ok thanks! 

I will try the first two.

Regarding the last point, since the code runs successfully anyways (provided the machine has enough memory). Don't you think that the bug could have gone unnoticed? Or you have some reason to exclude that. I am thinking at the fact that on one would have had to profile the memory usage of the program to figure that there was a problem.
> Don't you think that the bug could have gone unnoticed?

It could have, but `TraceEnum_ELBO` is pretty heavily used, and we've done a lot of memory profiling in the past. After working with Pyro and PyTorch for a few years, my posterior is 40% on a recent PyTorch regression, 40% on an edge case memory leak in Pyro that has never been noticed, and 20% on a recently introduced weird interaction between Pyro and PyTorch, so 60% chance this could be narrowed down by searching through PyTorch versions.
I have noticed a major GPU memory leak as well switching from PyTorch 1.10 to 1.11. Wasn't able to debug it and decided to stick to PyTorch 1.10.0 (and Pyro 1.8.0) for now.

Edit: CUDA 11.6, Arch Linux
Hmm maybe we should relax pytorch requirements and release to allow Pyro 1.8.2 to work with PyTorch 1.10. We'd need to do the same with Funsor. I think I was a little too eager dropping PyTorch 1.10 support, especially given colab still uses 1.10.
I have noticed a GPU memory leaks too with Pyro 1.8.1+06911dc and PyTorch 1.11.0. Downgrade to Pyro 1.6.0 and PyTorch 1.8.0 works normally. 
Downgrading, as @qinqian suggests, also resolves  #3014.


Resolving patch:
diff --git a/pyro/ops/einsum/torch_log.py b/pyro/ops/einsum/torch_log.py
--- a/pyro/ops/einsum/torch_log.py
+++ b/pyro/ops/einsum/torch_log.py
@@ -23,7 +23,10 @@ def einsum(equation, *operands):
 
     inputs, output = equation.split("->")
     if inputs == output:
-        return operands[0][...]  # create a new object
+        # Originally we return `operands[0][...]` but that caused
+        # memory leak in PyTorch >= 1.11 (issue #3068). Hence we
+        # return `operands[0].clone()` here.
+        return operands[0].clone()  # create a new object
     inputs = inputs.split(",")
 
     shifts = []


Test patch:
diff --git a/tests/distributions/test_pickle.py b/tests/distributions/test_pickle.py
--- a/tests/distributions/test_pickle.py
+++ b/tests/distributions/test_pickle.py
@@ -82,7 +82,7 @@ def test_pickle(Dist):
     try:
         dist = Dist(*args)
     except Exception:
-        pytest.skip(msg="cannot construct distribution")
+        pytest.skip(reason="cannot construct distribution")
 
     buffer = io.BytesIO()
     # Note that pickling torch.Size() requires protocol >= 2
diff --git a/tests/infer/test_predictive.py b/tests/infer/test_predictive.py
--- a/tests/infer/test_predictive.py
+++ b/tests/infer/test_predictive.py
@@ -139,7 +139,7 @@ def model():
         num_samples=num_samples,
         parallel=parallel,
     )
-    actual = predictive.get_samples()
+    actual = predictive()
     assert set(actual) == set(expected)
     assert actual["x"].shape == expected["x"].shape
     assert actual["y"].shape == expected["y"].shape
diff --git a/tests/ops/einsum/test_adjoint.py b/tests/ops/einsum/test_adjoint.py
--- a/tests/ops/einsum/test_adjoint.py
+++ b/tests/ops/einsum/test_adjoint.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2017-2019 Uber Technologies, Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+import gc
 import itertools
 
 import pytest
@@ -114,3 +115,29 @@ def test_marginal(equation):
         )
         actual = operand._pyro_backward_result
         assert_equal(expected, actual)
+
+
+@pytest.mark.filterwarnings("ignore:.*reduce_op is deprecated")
+def test_require_backward_memory_leak():
+    tensors = [o for o in gc.get_objects() if torch.is_tensor(o)]
+    num_global_tensors = len(tensors)
+    del tensors
+
+    # Using clone resolves memory leak.
+    for i in range(10):
+        x = torch.tensor(0.0)
+        require_backward(x)
+        x._pyro_backward.process(x.clone())
+
+    tensors = [o for o in gc.get_objects() if torch.is_tensor(o)]
+    assert len(tensors) <= 5 + num_global_tensors
+    del tensors
+
+    # Using [...] creates memory leak.
+    for i in range(10):
+        x = torch.tensor(0.0)
+        require_backward(x)
+        x._pyro_backward.process(x[...])
+
+    tensors = [o for o in gc.get_objects() if torch.is_tensor(o)]
+    assert len(tensors) >= 15 + num_global_tensors