# Batch process capability for pymc.sampling\_jax.sample\_numpyro\_nuts() with GPU?

**URL:** <https://discourse.pymc.io/t/batch-process-capability-for-pymc-sampling-jax-sample-numpyro-nuts-with-gpu/10337>\
**Category:** v5\
**Tags:** modeling\
**Created:** [September 8, 2022, 5:45pm UTC](https://discourse.pymc.io/t/batch-process-capability-for-pymc-sampling-jax-sample-numpyro-nuts-with-gpu/10337 "2022-09-08T17:45:37Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![Yixian\_Li](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/yixian_li/32/5631_2.png) [@Yixian\_Li](https://discourse.pymc.io/u/Yixian_Li)\
**Post date:** [September 8, 2022, 5:45pm UTC](https://discourse.pymc.io/t/batch-process-capability-for-pymc-sampling-jax-sample-numpyro-nuts-with-gpu/10337/1 "2022-09-08T17:45:37Z")

</div>

hi,

I am trying to run a large hierarchical model with pymc + jax on GPU but I keep running into out-of-memory issues. I tried various solutions that I found online but couldn’t really solve the problem so far. Is there anyway to go around this, for example by using batches? I’m very new to pymc and gpu-accelerated computing in general and any help is much appreciated!

The model that I’m trying to estimate has 10,000 respondents and 280 latent variables. It’s a hierarchical model where each respondent completed 12 tasks, from which I was hoping to estimate their individual values for the 280 latent variables. Ideally, I’d like to run 1000 tunes + 1000 draws, but the best that I could do so far before running into the OOM issue was 300 tunes + 300 warms on a GPU with 24G GPU memory.

Here’s my pm model:

```auto
basic_model = pm.Model()

with basic_model:

    alphas = pm.Normal("alphas", mu=0, sigma=1, shape=num_parameters)
    sd_distrib = pm.Normal.dist(mu=1, sigma=10, shape=num_parameters)
    L_sigma, corr, stds = pm.LKJCholeskyCov('L_sigma', eta=1, n=num_parameters, sd_dist=sd_distrib, compute_corr=True, store_in_trace=False)
    cov = pm.Deterministic("cov", L_sigma.dot(L_sigma.T))
    betas_init = pm.MvNormal('betas_init', mu=np.zeros((num_parameters)), cov=np.eye(num_parameters), shape=(num_respondents, num_parameters))
    
    
    def logp(choicedata, alphas, betas_init, L_sigma):

        #get transfomred beta from alpha and beta-init
        betas = alphas.reshape(num_respondents,1) + at.dot(betas_init,L_sigma)
        #get respondent-level log likelihood
        resp_log = CHOICE_PROBABILITIES(choicedata, aggregation = True)
      
        return resp_log
    
    
    likelihood = pm.DensityDist(
        "likelihood",
        alphas, betas_init, L_sigma,
        logp=logp,
        observed = ts_choicedata
        )

with basic_model:
    idata = pm.sampling_jax.sample_numpyro_nuts(draws = draws, 
                                                tune=tune, 
                                                target_accept=target_accept, 
                                                chains=chains,
                                                postprocessing_backend = 'cpu',
                                                idata_kwargs = {'log_likelihood': False},
                                                progress_bar=True,
                                               )

```

The CHOICE\_PROBABILITIES function is a customized likelihood function to calculate the loglikelihood at the individual respondent level.

Here are the things that I’m already doing:

- Add the jax env setting before calling jax:  
%env XLA\_PYTHON\_CLIENT\_PREALLOCATE = false  
%env XLA\_PYTHON\_CLIENT\_ALLOCATOR= platform

- Use idata\_kwargs = {‘log\_likelihood’: False} when sampling

- I also noticed that if I do more tunes and fewer draws, the GPU use (reading from nvidia-smi) seems to be smaller than when I run the same total number of iterations but fewer tunes and more draws.

Thanks a lot!

```auto

```

---

<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:** [September 9, 2022, 9:10am UTC](https://discourse.pymc.io/t/batch-process-capability-for-pymc-sampling-jax-sample-numpyro-nuts-with-gpu/10337/2 "2022-09-09T09:10:44Z")

</div>

Yout can try setting `chain_method` to `"sequential"` or `"vectorized"` and see if it does better.

It is normal for less RAM to be needed for tuning than drawing, because tuning samples can be discarded over time.

---

<div class="post-metadata">

**Author:** ![Yixian\_Li](https://yyz2.discourse-cdn.com/flex036/user_avatar/discourse.pymc.io/yixian_li/32/5631_2.png) [@Yixian\_Li](https://discourse.pymc.io/u/Yixian_Li)\
**Post date:** [September 12, 2022, 7:14pm UTC](https://discourse.pymc.io/t/batch-process-capability-for-pymc-sampling-jax-sample-numpyro-nuts-with-gpu/10337/3 "2022-09-12T19:14:28Z")

</div>

Thanks a lot for your reply!

I’ve tried that but it didn’t seem to improve much. I’m running more tunes in an attempt to reach the stable posterior before drawing and it seems to improve the accuracy when modelling with my synthetic data.

Are there any other ways to deal with the OOM issue with GPU?

---

<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:** [September 12, 2022, 7:29pm UTC](https://discourse.pymc.io/t/batch-process-capability-for-pymc-sampling-jax-sample-numpyro-nuts-with-gpu/10337/4 "2022-09-12T19:29:29Z")

</div>

That’s more of a JAX/Numpyro question so you may have more luck asking in their forums.