Numpyro: Excessive memory usage at the beginning of NUTS inference

[d-caviedes]

Hello everyone,

asked by @fehiepsi I create the following issue regarding what seems "excessive" memory consumption at the beginning of NUTS inference (originally addressed in the discussion forum https://forum.pyro.ai/t/numpyro-memory-usage-at-the-beginning-of-nuts-inference/).

I am running inferences on a linear regression model such as

def model_sbl(data):
    obs = data["obs"]
    Nw = data["Nw"]
    Nt = data["Nt"]
    X = data["X"]
    b = data["b"]
    tau0 = data["tau0"]
    sigma = numpyro.sample("sigma", dist.LogNormal(0.0, 1.0))
    dist_var = dist.InverseGamma(1, 10 ** (-b) * jnp.ones(Nw * Nt))
    with numpyro.plate("coefficients", Nw * Nt):
        w_var = jnp.sqrt(numpyro.sample("w_var", dist_var))
        dist_w = dist.Normal(jnp.zeros((Nw * Nt)), tau0 * w_var)
        w = numpyro.sample("w", dist_w)
    mean = jnp.matmul(X, w)
    y = dist.Normal(mean, sigma)
    with numpyro.plate("observations", size=len(obs)):
        numpyro.sample("obs", y, obs=obs)

The parameter space is quite large, with _X_ being a dictionary with dimensions (5236, 123200) in this case, and _w_ 123200 unknown coefficients. To run the inferences

data = dict(
      obs=jnp.array(h_fit.T.flatten()),
      X=jnp.array(D_fit),
      b=0,
      tau0=1e-1,
      sigma=0.1,
      w_var=0.1,
      Nw=Nw,
      Nt=Nt,
  )
rng_key = jrandom.PRNGKey(0)
rng_key, rng_key_ = jrandom.split(rng_key)
my_kernel = NUTS(model_sbl, max_tree_depth=7)
posterior = MCMC(my_kernel, num_samples=100, num_warmup=50, num_chains=1)
pdb.set_trace()
# run_posterior(posterior, rng_key_, data)
posterior.run(rng_key_, data=data)

The memory consumption tracked with memory_profiler is detailed in the following Figure. I annotated what happened while debugging. The first and second samples when numpyro.sample("obs", y, obs=obs) is called consume quite a lot of memory. But when these two first samples are computed, it stabilizes to a constant RAM, and keeps like that during the entire inference

In order to alleviate it, I tried creating the dictionary _X_ with float16 instead of float32 (as I assumed it could be related). Memory profile detailed in the next figure. The consumption of the first two samples, while still high, is well below the values before. However, the consumption of the 3rd and subsequent is almost identical to the float32 case.

It seems that the two first samples are problematic, adding up the in the RAM usage.

I tried to come with a workaround. I have seen here https://github.com/google/jax/issues/2072 that one can flush the cache from XLA. I tried xla._xla_callable.cache_clear() at the beginning of my model, such that every time it goes to sample, the cache is cleared. It can be seen that the pattern is the same as in the first figure for the first sample, repeated for each sample. Now the first two samples don't add up, but there is no stable RAM consumption after the 2nd, as it happened by default.

What I found quite interesting at this point is that if one clears the cache only after the first few samples, and doesn't for the rest, then it doesn't grow but stabilizes, allowing for the computation of the samples (see last figure). This way, it avoids the first two samples to add up in memory.

So, what is it going on? Should this be fixed? Am I removing any info that is vital for the inferences by clearing the cache? I can post results about the posteriors for the different cases and compare, but so far it seems to be fine.

I think that the option of alleviating the memory consumption in NumPyro would be very useful, instead of going to SVI when the model is rather large. As an example, there is no way I can even try to solve this inference problem in a machine with 32GB of RAM with other ppl software such as STAN :).

[fehiepsi]

@d-caviedes Could you try using the master branch and update jax/jaxlib?

pip install --upgrade jax jaxlib
pip install https://github.com/pyro-ppl/numpyro/archive/master.zip

I didn't see the memory issue with the following settings

data = dict(
      obs=jnp.ones(5236),
      X=jnp.ones((5236, 123200)),
      b=0,
      tau0=1e-1,
      sigma=0.1,
      w_var=0.1,
      Nw=1232,
      Nt=100,
  )

Probably it is fixed in https://github.com/pyro-ppl/numpyro/pull/924?

[d-caviedes]

Yes, updating jaxlib from 0.1.59 to 0.1.61 did the trick. But now there is another issue, and I am still confused.

My baseline model should be Nw=700 Nt=308, s.t. X = (5236, 215600) and w (215600). Now there is no problem with the RAM, but it just terminates abruptly at some point without exceptions or any error message.

I tried this with anothe computer (HPC with 128Gb of ram), and same story

also with precision float32

No idea why, is it jax?

[fehiepsi]

Thank you, @d-caviedes! It is very cool to learn those experiments from your clear report. :)

In my system, with X.shape == (5236, 215600), Nw=700, Nt=308, MEM is stable at 28.9% (I have a 32GB ram machine) this is the current progress (almost 11 minutes)

warmup:  13%|█▎        | 20/150 [10:56<1:58:22, 54.63s/it, 127 steps of size 1.83e-07. acc. prob=0.50]

I guess we should report your issue to JAX team, but I don't have a replicable code. Probably they will understand the situation better based on your hardware.

[d-caviedes]

Thanks again. After creating a completely fresh environment and following all the indications, it does work for your example. I can see that the problem comes when I create the dictionary from a function like this one

def sinc_dictionary(s, r, t, fs, c, precision="float32"):
    """Creates dictionary with sinc function as basis

    Args:
        s ([type]): [description]
        r ([type]): [description]
        t ([type]): [description]
    """
    Nw = s.shape[0]
    Nt = len(t)
    Nm = r.shape[0]
    d_ = distance_between(s.T, r.T)
    Dt = t[..., None, None] - (d_ / c)[None]
    #%% Dictionary
    D_ = np.empty((Nt * Nm, Nt * Nw), dtype=precision)
    for m in tqdm(range(Nm)):
        for l in range(Nw):
            column = 1 / (4 * np.pi * d_[m, l]) * np.sinc(fs * Dt[:, m, l])
            # row = 1 / (4 * np.pi * d_[m, l]) * np.sinc(fs_resample * Dt[:, m, l])
            row = np.zeros(column.shape)
            row[0] = column[0]
            D_[Nt * m : Nt * (m + 1), Nt * l : Nt * (l + 1)] = sc.linalg.toeplitz(
                column, row
            )
    return D_

X = sinc_dictionary(s, r_fit, t, fs_new, c, precision=precision)

Does it have to be with jax trying to have derivatives of all those functions? Because I don't need them :S.

[fehiepsi]

Oh, I think I faced this issue some time ago, when jax converts a _numpy array with some "special" stride (or something else) properties_ to a DeviceArray. I guess the issue will go away if you apply .copy(...) or .tolist(...) before returning D_. Sorry, I can't recall exactly how it was resolved... :(

[fehiepsi]

@d-caviedes Would return D_.copy() solve the issue? I don't think that gradient or something like that plays any role here.

[d-caviedes]

@fehiepsi hi again :). I couldn't resolve the issue just returning such copy. I have created an issue in JAX repo, and I hope they can help me with it (https://github.com/google/jax/issues/5944). I have written there a small example, trying to see where is it that the problem comes from. It appears that if ones or zeros are used for the dictionary, it is not a problem. However, if the dictionary combines zeros and ones, or it is created with random numbers, it appears to be a problem. Let's see...

[fehiepsi]

Could you make a gist file with a full reproducible code? I would like to play with this example to see where the issue comes from.

[d-caviedes]

Here it goes.

https://gist.github.com/d-caviedes/b13b5a02a0452204a3b3bb7ea8b02661

[fehiepsi]

Interesting! I can replicate the issue in my system. For now, I have no idea why memory issues depend on input values (rather than input shapes). Will look into this more...

[fehiepsi]

It seems that the issue is not relevant to sinc_dictionary. If I set X = np.zeros((Nt * Nm, Nt * Nw)), the issue does not happen but if I set X = np.random.randn(Nt * Nm, Nt * Nw), the issue happens. I guess jnp.dot(X, w) performs better for constant X.

[d-caviedes]

Yes, it seems to be something about it. It would be great if it gets somehow resolved. The problem is that even if I had sufficient RAM, the inference fails at some point (as shown in the graphs in the report). I had to go back to pyro at this point :S.

[fehiepsi]

Good idea on using Pyro! For those very large tensor computations, it is unnecessary to use JAX backend.