# Pymc/numpyro GPU memory allocation

**URL:** <https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623>\
**Category:** General\
**Created:** [June 12, 2022, 3:45pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623 "2022-06-12T15:45:59Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![Erlebach](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/erlebach/32/6605_2.png) [@Erlebach](https://discourse.pymc.io/u/Erlebach)\
**Post date:** [June 12, 2022, 3:45pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/1 "2022-06-12T15:45:59Z")

</div>

I am solving a coin flipping problem with `pymc` (version 4). Here is the (very simple) model:

```python
def create_coin(data):
    with pm.Model() as coin_model:
        p_coin = pm.Beta("p_coin", 2, 2)
        heads = pm.Bernoulli("flips", p_coin, observed=data)
    return coin_model

p = 0.55
nb_flips = 100
nb_chains = 4
nb_draws = 10000
data = np.random.choice([0,1], size=nb_flips, p=[1-p, p])
model = create_coin(data)

with model:
   idata = jx.sample_numpyro_nuts(target_accept=0.9, draws=nb_draws, tune=nb_draws, chains=nb_chains, chain_method='parallel')

```

When I run 20,000 draws with 100,000 coin flips and 4 chains, I am exceeding 10GB of memory on the GPU. I found that pymc preallocates memory in the amount `nb_chains * nb_coin_flips * nb_draws`.

My question is why pymc must store this amount of information given that I am only computing the trace of `p_coin`? Why must it store the equivalent of all observables at every draw?  
Is there a way to reduce memory usage?  
Are there routines to track memory usage, not only in JAX, but using pymc’s standard `sample` method?  
Thanks for any insight you might provide!

---

<div class="post-metadata">

**Author:** ![twiecki](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/twiecki/32/6930_2.png) [@twiecki](https://discourse.pymc.io/u/twiecki)\
**Post date:** [June 12, 2022, 4:52pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/2 "2022-06-12T16:52:38Z")

</div>

You can set `postprocessing_backend='cpu'` to get around that problem.

---

<div class="post-metadata">

**Author:** ![Erlebach](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/erlebach/32/6605_2.png) [@Erlebach](https://discourse.pymc.io/u/Erlebach)\
**Post date:** [June 12, 2022, 5:06pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/3 "2022-06-12T17:06:31Z")

</div>

Thanks. I know imcan operate on the cpu, which is much faster. I am still interested in the rationale for how memory is preallocated on the GPU. thanks.

---

<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:** [June 12, 2022, 6:32pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/4 "2022-06-12T18:32:37Z")

</div>

That’s beyond the scope of PyMC. We simple rely on Numpyro to create a JAX graph from the model logp and sample it. You would have to ask the Numpyro folks for more details, but even they might only forward you to the JAX folks.

---

<div class="post-metadata">

**Author:** ![Erlebach](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/erlebach/32/6605_2.png) [@Erlebach](https://discourse.pymc.io/u/Erlebach)\
**Post date:** [June 12, 2022, 6:54pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/5 "2022-06-12T18:54:13Z")

</div>

Thank you, Ricardo. You are correct of course, but hope springs eternal.

---

<div class="post-metadata">

**Author:** ![lucianopaz](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/lucianopaz/32/2186_2.png) [@lucianopaz](https://discourse.pymc.io/u/lucianopaz)\
**Post date:** [June 13, 2022, 6:43am UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/6 "2022-06-13T06:43:14Z")

</div>

Pymc doesn’t actually need to allocate this memory for inference. It does the allocation when creating the inference data object  
By default, pymc will compute the point wise log likelihood, and it does this to enable arviz to run model comparison without needing to have a reference to the model that generated the inference results (it’s model agnostic and can compare results drawn from different PPLs).  
Again, this isn’t a hard requirement and you can ask pymc not to compute the point wise log likelihood, which will let you avoid the out of memory error. To do this, you need to pass the following key word argument to `sample` (or which ever jax variant you choose):  
`idata_kwargs={“log_likelihood”: False}`

---

<div class="post-metadata">

**Author:** ![Erlebach](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/erlebach/32/6605_2.png) [@Erlebach](https://discourse.pymc.io/u/Erlebach)\
**Post date:** [June 13, 2022, 10:25am UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/7 "2022-06-13T10:25:04Z")

</div>

Thanks. If I only want the posterior, what else should I cut out through proper argument selection?

---

<div class="post-metadata">

**Author:** ![lucianopaz](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/lucianopaz/32/2186_2.png) [@lucianopaz](https://discourse.pymc.io/u/lucianopaz)\
**Post date:** [June 13, 2022, 8:11pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/8 "2022-06-13T20:11:47Z")

</div>

Just discarding the log likelihood should be fine. The other sampling stats are important to check your result’s quality, and don’t eat up that much memory

---

<div class="post-metadata">

**Author:** ![Erlebach](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/erlebach/32/6605_2.png) [@Erlebach](https://discourse.pymc.io/u/Erlebach)\
**Post date:** [June 13, 2022, 8:32pm UTC](https://discourse.pymc.io/t/pymc-numpyro-gpu-memory-allocation/9623/9 "2022-06-13T20:32:32Z")

</div>

Thank you!