# Pymc, numpyro GPs and Transform RVs memory behaviour

**URL:** <https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996>\
**Category:** v5\
**Tags:** gpu, gaussian\_process, sampling\
**Created:** [December 13, 2022, 2:17pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996 "2022-12-13T14:17:29Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![Thops](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/thops/32/5286_2.png) [@Thops](https://discourse.pymc.io/u/Thops)\
**Post date:** [December 13, 2022, 2:17pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/1 "2022-12-13T14:17:29Z")

</div>

While developing GP models with numpyro to sample via the GPU I’ve noticed somewhat strange memory behavior. After `pymc.sampling.jax.sample_numpyro_nuts` finishes running (very fast) it enters a ‘Transforming RVs’ phase witch triggers a massive memory spike. I’m not sufficiently familiar with the internals of `pymc` to know which transformation is being done and how (I think it might be the same one mentioned by the `reparameterize=True` parameter of `pymc.gp.Latent().prior` however this step seems to have massive memory requirements. I’ve empirically determined the amount of memory is roughly:

64(b) \times n\_{chains}\times n\_{draws}\times d(number\ of\ GPs)\times N^2\times 10^{-9}(GB\ b^{-1})

Is there a way to avoid this giant tensor? Perhaps do the transform in minibatches? If we want to build a model with many target variables, the memory requirements increase unreasonably i.e.

64(b) \times 2(chains)\times 1000(draws)\times 23(number\ of\ GPs)\times (500)^2\times 10^{-9}(GB\ b^{-1})\approxeq 507 GB\ !!

Even if it impossible to avoid this, it would be usefull (and pretty easy I think) to add a check for sufficient memory. It certainly was pretty unclear to me where all the XLA OOM errors were coming from. Here’s the model I’m working with for reference:

```auto
with pymc.Model() as binary_model:
    gps=[]
    gppriors=[]
    _,M=chem_indicators.shape
    η=pymc.Normal('η', mu=3.0, sigma=1.0)
    for target in binary_targets_tidy.columns:
        mean_func=pymc.gp.mean.Constant(c=0)
        σv = pymc.LogNormal(f'σv_{target}', mu=.0, sigma=.25)+.5
        λ = pymc.MvNormal(f'λ_{target}', mu=3.0*np.ones(M),cov=σv*np.eye(M),)
        σ = pymc.Normal(f'σ_{target}', mu=3.0,sigma=1.0)
        k = η**2*pymc.gp.cov.ExpQuad(M, ls=λ**2)
        noise = pymc.gp.cov.WhiteNoise(σ)
        κ = k+noise
        gp=pymc.gp.Latent(mean_func=mean_func, 
                          cov_func=κ)
        _f=gp.prior(f'_f_{target}', X_train.values)
        gps.append(gp)
        gppriors.append(_f)
    f = pymc.Deterministic('f', at.math.sigmoid(at.stack(*gppriors).T))
    for idx,target in enumerate(binary_targets_tidy.columns):
        y_obs = pymc.Bernoulli(f'y_obs_{target}',p=f[:,[idx]], observed=Y_train.loc[:,target] )

```

---

<div class="post-metadata">

**Author:** ![bwengals](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/bwengals/32/1237_2.png) [@bwengals](https://discourse.pymc.io/u/bwengals)\
**Post date:** [December 14, 2022, 3:13am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/2 "2022-12-14T03:13:55Z")

</div>

Could you maybe provide a bit more detail on how fast it samples? It looks like a pretty large model, so I’m a bit surprised if it samples very quickly. The reparameterization you note here shouldn’t be responsible for the issue you’re seeing, because the reparameterization is really about whether you use a centered or non-centered parameterization. It shouldn’t trigger anything after sampling is done.

One speedup I see is using a Normal vector instead of an MvNormal with a diagonal covariance matrix for the lengthscale prior. If you need a non-diagonal covariance I’d try the non-centered version there too. [Here’s more info](https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Computational_methods).

Is your error coming from [here?](https://github.com/pymc-devs/pymc/blob/23a163663c52a11cf9985c698d4112291a0a7acb/pymc/sampling/jax.py#L624) It looks like this happens when Jax moves everything from the GPU to the CPU.

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 14, 2022, 5:56am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/3 "2022-12-14T05:56:04Z")

</div>

Maybe this is relevant? [Out of memory when "transforming variables" in Numpyro & JAX - #11 by wnorcbrown](https://discourse.pymc.io/t/out-of-memory-when-transforming-variables-in-numpyro-jax/9904/11)

---

<div class="post-metadata">

**Author:** ![Thops](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/thops/32/5286_2.png) [@Thops](https://discourse.pymc.io/u/Thops)\
**Post date:** [December 14, 2022, 10:40pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/4 "2022-12-14T22:40:53Z")

</div>

Sampling time varies slightly weirdly enough. If I downsample from the original \sim 415 points to about 30%, it samples in about 20mins/chain. Increasing the number of points past some unclear point causes the sampling time to jump to about 2 hours per chain. Sampling time doesn’t seem to monotonically increase sampling time. Rather, past some unclear number of points, sampling time jumps abruptly. If it isn’t then it would be vary strange as I’ve repeatedly noticed:

1. GPU memory consumption is very low during sampling, and as soon as ‘Transforming RVs’ appears it jumps to near max
2. Every time `reparameterize=True` I get hit with an OOM error. Every time `repameterize=False` no OOM errors appear. Specifically from `postprocessin_backend='cpu'` there are no error messages - the kernel just dies. Setting `postprocessing_backend='gpu'` raises an XLA OOM error (`nvidia-smi` seems to agree. All 11GBs of vram are used up at this point)
3. The reported error mentioned shows roughly `... f64(2,N,N,1)` (it might also have a sample dimention, I’ll have to check). I note 64 are the columns of the input data
4. I’ve no idea if `vmap` has anything to do with the OOM error. The CPU transfer explanation doesn’t seem plausible to me, that transfer should happen regardless of whether `reparameterize` is True or not
5. If by normal vector you mean something like `pymc.Normal(..., shape=M)` I’ve considered this. However I’d like each weight to have it’s own mean (there’s no reason in my case they should be the same) and in my understanding if I do `pymc.Normal(..., shape=M)` that gives the same expectation for all the weights.
6. [This post](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996) Does not seem applicable here. Splitting it into two deterministics for `stack` and `sigmoid` seems to have no effect in computation time or memory as far as I can tell

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 15, 2022, 7:39am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/5 "2022-12-15T07:39:56Z")

</div>

> [@Thops](#):
>
> If by normal vector you mean something like `pymc.Normal(..., shape=M)` I’ve considered this. However I’d like each weight to have it’s own mean (there’s no reason in my case they should be the same) and in my understanding if I do `pymc.Normal(..., shape=M)` that gives the same expectation for all the weights.

No it doesn’t. `Normal(”name”, shape=(10,))` will create a vector of 10 independent normal variates. You can even give them different prior parameters `Normal(”name”, mu=np.arange(10), shape=(10,))`

This results in a much more efficient computational graph because we only need to represent one node instead of 10.

This is explained here: [Distribution Dimensionality — PyMC 5.10.0 documentation](https://www.pymc.io/projects/docs/en/stable/learn/core_notebooks/dimensionality.html#explicit-batch-dimensions)

---

<div class="post-metadata">

**Author:** ![Thops](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/thops/32/5286_2.png) [@Thops](https://discourse.pymc.io/u/Thops)\
**Post date:** [December 15, 2022, 7:43am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/6 "2022-12-15T07:43:10Z")

</div>

Good to know, I’ll look into implementing it. The number of variables doesn’t seem related to the memory error though, as swapping them out with a single scalar (or `pymc.Normal` RV) doesn’t seem to change anything about the memory errors

---

<div class="post-metadata">

**Author:** ![bwengals](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/bwengals/32/1237_2.png) [@bwengals](https://discourse.pymc.io/u/bwengals)\
**Post date:** [December 21, 2022, 8:06am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/7 "2022-12-21T08:06:48Z")

</div>

It’s a bit hard to figure out what might exactly be happening since your model isn’t simple and we can’t run it on our end and see what you’re seeing. Is it possible to boil it down to something that we could run?

I can say though regarding point 2 and repameterize. Reparameterize should basically always be True. It’s very strange that that’s what gives you the errors. As far as I know, at least with NUTS, the only reason to turn reparameterize to False is to show that it doesn’t work! Trying to figure out what’s up with this would be where I’d start.

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 21, 2022, 8:28am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/8 "2022-12-21T08:28:02Z")

</div>

You have many variables created like this in nested loops. Did you try vectorizing with batched shapes everywhere you could or just in one place?

---

<div class="post-metadata">

**Author:** ![Thops](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/thops/32/5286_2.png) [@Thops](https://discourse.pymc.io/u/Thops)\
**Post date:** [December 21, 2022, 8:34am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/9 "2022-12-21T08:34:35Z")

</div>

@ricardoV94 I’m not sure how vectorizing would be done in this instance. The idea is I want multiple GPs, each with it’s own kernel and I can’t think of any way to get this behavior other than a loop.

@bwengals I will try, although the multiple GP’s seem key to getting the error, so aside from fixing all the kernel parameters I can’t think of another way to trim this down

---

<div class="post-metadata">

**Author:** ![ricardoV94](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/ricardov94/32/5775_2.png) [@ricardoV94](https://discourse.pymc.io/u/ricardoV94)\
**Post date:** [December 21, 2022, 10:05am UTC](https://discourse.pymc.io/t/pymc-numpyro-gps-and-transform-rvs-memory-behaviour/10996/10 "2022-12-21T10:05:04Z")

</div>

> [@Thops](#):
>
> @ricardoV94 I’m not sure how vectorizing would be done in this instance. The idea is I want multiple GPs, each with it’s own kernel and I can’t think of any way to get this behavior other than a loop.

All the variables that are not GPs can probably be created outside of the loop and indexed where needed, and the same goes for the likelihood (maybe?).

Also I didn’t notice, but you can consider not wrapping the contents of the Determinstic `f` (just write `f = at.math.sigmoid(...)`. This way it won’t be computed and stored in the trace after sampling is over. If you need these values you can always try to recreate them later. This might go a long way in reducing the memory spike you’re observing