Repository: pyro-ppl/numpyro

Issue:
Memory leak when running MCMC multiple times
This was surfaced in https://github.com/pyro-ppl/brmp/issues/67.

I can see this issue on a simple NumPyro model, where the count of `jax.Var` objects (along with some other types) keeps increasing per iteration (the code below can be run with the latest jax/jaxlib to reproduce the issue). I have not yet been able to replicate this on a purely JAX based code, so it is possible that this is something specific to NumPyro.

```
import gc
from collections import Counter

import numpy as np

from numpyro import sample
import numpyro.distributions as dist
from jax import random, vmap
from numpyro.infer import MCMC, NUTS


def model(y_obs):
    mu = sample('mu', dist.Normal(0., 1.))
    sigma = sample("sigma", dist.HalfCauchy(3.))
    y = sample("y", dist.Normal(mu, sigma), obs=y_obs)

for i in range(30):
    kernel = NUTS(model)
    mcmc = MCMC(kernel, 1, 2)
    mcmc.run(random.PRNGKey(0), np.zeros((1,)))
    print("\nGC OBJECTS:")
    cnt = Counter()
    # force collection; it is expected that count of different types
    # should not increase per iteration
    gc.collect()
    for x in gc.get_objects():
        if isinstance(x, list):
            if len(x) > 1:
                cnt[type(x[0])] += 1
    print(cnt.most_common(10))
```


Discussion:
Interesting! This is the first time I touch `gc.collect()` stuff. :D I'll dive into this issue a bit more...
@neerajprad I am seeing the same issue with the following code (which only depends on JAX), but I am not sure if that's leakage (because I am not familiar with gc). Looking like the issue happens when we use while_loop together with a body_fn method containing jax.random.

```
import gc
from collections import Counter

import jax; jax.config.update('jax_platform_name', 'cpu')
from jax import numpy as np
from jax import lax, random

class ABC:
    def body_fn(self, i, key):
        key, _ = random.split(key)
        return key

for i in range(3):
    abc = ABC()
    value = lax.fori_loop(0, 10, abc.body_fn, random.PRNGKey(0))

    print("\nGC OBJECTS:")
    cnt = Counter()
    # force collection; it is expected that count of different types
    # should not increase per iteration
    gc.collect()
    for x in gc.get_objects():
        if isinstance(x, list):
            if len(x) > 1:
                cnt[type(x[0])] += 1
    print(cnt.most_common(10))
```
Thanks for your code snippet, @fehiepsi. Interestingly, this version of your code doesn't have the same issue, and behaves as you would expect. :)

```python
import gc
from collections import Counter

import jax; jax.config.update('jax_platform_name', 'cpu')
from jax import random, lax


def model(key):
    return random.normal(key)

def cond_fn(state):
    return state[0] < 100

def body_fn(state):
    i, key, _ = state
    key, subkey = random.split(key)
    return i + 1, key, model(subkey)


init_state = (0, random.PRNGKey(0), model(random.PRNGKey(0)))


for i in range(10):
    value = lax.while_loop(cond_fn, body_fn, init_state)

    print("\nGC OBJECTS:")
    cnt = Counter()
    # force collection; it is expected that count of different types
    # should not increase per iteration
    gc.collect()
    for x in gc.get_objects():
        if isinstance(x, list):
            if len(x) > 1:
                cnt[type(x[0])] += 1
    print(cnt.most_common(10))
```
Yeah, should we report this issue to jax devs?
Thanks, @fehiepsi. I just shared it on the same issue. I tried debugging this further last night, but I think it might be best to hear from the JAX devs before spending more time debugging this. 
I am closing this issue because this is due to the JAX cache working as expected. The right solution to https://github.com/pyro-ppl/brmp/issues/67 would be for us to not miss the cache and recompile each time, which should be solvable via #441. 
@fehiepsi - thanks for helping debug this! 
I think with all the refactoring we have done (reg. caching), this issue has resurfaced. 😄 We should aim to fix this asap. 
Just a context for future reference: the topic issue can be solved with
```
import gc
from collections import Counter

import numpy as np

from numpyro import sample
import numpyro.distributions as dist
from jax import random, vmap
from numpyro.infer import MCMC, NUTS


def model(y_obs):
    mu = sample('mu', dist.Normal(0., 1.))
    sigma = sample("sigma", dist.HalfCauchy(3.))
    y = sample("y", dist.Normal(mu, sigma), obs=y_obs)

kernel = NUTS(model)
mcmc = MCMC(kernel, 1, 2)
mcmc._compile(random.PRNGKey(0), np.zeros((1,)))
for i in range(10):
    mcmc.run(random.PRNGKey(0), np.zeros((1,)))
    print("\nGC OBJECTS:")
    cnt = Counter()
    # force collection; it is expected that count of different types
    # should not increase per iteration
    gc.collect()
    for x in gc.get_objects():
        if isinstance(x, list):
            if len(x) > 1:
                cnt[type(x[0])] += 1
    print(cnt.most_common(10))
```
and use `jax.jit(lax.map, static_argnums=(0,))(self.constrain_fn, states['z'])` to transform mcmc params back to constrained domain.
The leak is due to caching mechanism of JAX and the current caching mechanism has resolved a big part of this.

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -271,7 +271,7 @@ def __init__(self,
         self._collection_params = {}
         self._set_collection_params()
 
-    def _get_cached_fn(self):
+    def _get_cached_fns(self):
         if self._jit_model_args:
             args, kwargs = (None,), (None,)
         else:
@@ -279,20 +279,36 @@ def _get_cached_fn(self):
             kwargs = tree_map(lambda x: _hashable(x), tuple(sorted(self._kwargs.items())))
         key = args + kwargs
         try:
-            fn = self._cache.get(key, None)
+            fns = self._cache.get(key, None)
         # If unhashable arguments are provided, proceed normally
         # without caching
         except TypeError:
-            fn, key = None, None
-        if fn is None:
+            fns, key = None, None
+        if fns is None:
+
+            def laxmap_postprocess_fn(states, args, kwargs):
+                if self.postprocess_fn is None:
+                    body_fn = self.sampler.postprocess_fn(args, kwargs)
+                else:
+                    body_fn = self.postprocess_fn
+                if self.chain_method == "vectorized" and self.num_chains > 1:
+                    body_fn = vmap(body_fn)
+
+                return lax.map(body_fn, states)
+
             if self._jit_model_args:
-                fn = partial(_sample_fn_jit_args, sampler=self.sampler)
+                sample_fn = partial(_sample_fn_jit_args, sampler=self.sampler)
+                postprocess_fn = jit(laxmap_postprocess_fn)
             else:
-                fn = partial(_sample_fn_nojit_args, sampler=self.sampler,
-                             args=self._args, kwargs=self._kwargs)
+                sample_fn = partial(_sample_fn_nojit_args, sampler=self.sampler,
+                                    args=self._args, kwargs=self._kwargs)
+                postprocess_fn = jit(partial(laxmap_postprocess_fn,
+                                             args=self._args, kwargs=self._kwargs))
+
+            fns = sample_fn, postprocess_fn
             if key is not None:
-                self._cache[key] = fn
-        return fn
+                self._cache[key] = fns
+        return fns
 
     def _get_cached_init_state(self, rng_key, args, kwargs):
         rng_key = (_hashable(rng_key),)
@@ -310,10 +326,7 @@ def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
         if init_state is None:
             init_state = self.sampler.init(rng_key, self.num_warmup, init_params,
                                            model_args=args, model_kwargs=kwargs)
-        if self.postprocess_fn is None:
-            postprocess_fn = self.sampler.postprocess_fn(args, kwargs)
-        else:
-            postprocess_fn = self.postprocess_fn
+        sample_fn, postprocess_fn = self._get_cached_fns()
         diagnostics = lambda x: self.sampler.get_diagnostics_str(x[0]) if rng_key.ndim == 1 else ''   # noqa: E731
         init_val = (init_state, args, kwargs) if self._jit_model_args else (init_state,)
         lower_idx = self._collection_params["lower"]
@@ -323,7 +336,7 @@ def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
         collection_size = collection_size if collection_size is None else collection_size // self.thinning
         collect_vals = fori_collect(lower_idx,
                                     upper_idx,
-                                    self._get_cached_fn(),
+                                    sample_fn,
                                     init_val,
                                     transform=_collect_fn(collect_fields),
                                     progbar=self.progress_bar,
@@ -344,9 +357,10 @@ def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
         # so we only need to filter out the case site_value.shape[0] == 0
         # (which happens when lower_idx==upper_idx)
         if len(site_values) > 0 and jnp.shape(site_values[0])[0] > 0:
-            if self.chain_method == "vectorized" and self.num_chains > 1:
-                postprocess_fn = vmap(postprocess_fn)
-            states[self._sample_field] = lax.map(postprocess_fn, states[self._sample_field])
+            if self._jit_model_args:
+                states[self._sample_field] = postprocess_fn(states[self._sample_field], args, kwargs)
+            else:
+                states[self._sample_field] = postprocess_fn(states[self._sample_field])
         return states, last_state
 
     def _set_collection_params(self, lower=None, upper=None, collection_size=None, phase=None):


Test patch: