I am able to finish the sampling process on a large model using the JAX backend with a GPU, but it fails during the “transforming variables” step because it requires too much RAM. I set postprocessing_backend="cpu" which has a lot more RAM than my GPU, but it still runs out. Looking at the code in …

in fact this seems to work: from jax.experimental.maps import SerialLoop, xmap num_chunks = 10 # must be a multiple of number of samples mapper = xmap(jax_fn, in_axes=['chain', 'samples', ...], out_axes=["chain", 'samples', ...], axis_resources={"samples…

Update: from my testing, it is happening during the vmap step: result = jax.vmap(jax.vmap(jax_fn))( *jax.device_put(raw_mcmc_samples, jax.devices(postprocessing_backend)[0]) ) where postprocessing_backend = "cpu". Is there some way to make this process require less RAM?

[image] jhrcook: "cpu". Is there some way to make this process require less RAM? CC @ferrine

The only way to use less RAM as I know so far is to reduce the amount of deterministic variables

I’ll definitely try that! Thank you!

Also, disabling the loglikelihood computation via idata_kwargs if you are not planning to do model comparison can help if your model has many observations

Thank you for the recommendation. I don’t think it would fix this problem, though, because those parameters aren’t used at the step that is failing. (Of course, you know the code base far better than me, so please correct me if I’m wrong.) But I have seen this suggested for other RAM problems, so I’…

@ferrine made very little difference for me. i’m limited to running very small chains. an option to use jax.lax.map — JAX documentation should save memory usage and stop the memory spike

This is really nice!! How is the speed compare to vmap? Would you like to send a PR to PyMC?

As far as I could tell the speed was the same. Sure :+1:

Out of memory when "transforming variables" in Numpyro & JAX

Questions v5

wnorcbrown December 12, 2022, 10:02am 11

As far as I could tell the speed was the same.

Sure

Pymc, numpyro GPs and Transform RVs memory behaviour

Topic		Replies	Views
Pymc, numpyro GPs and Transform RVs memory behaviour v5 gpu , gaussian_process , sampling	9	1002	December 21, 2022
Out of Memory during Transforming Variables Step using Numpyro. Increasing postprocessing_chunks doesn't seems to help. What should I do? v5 modeling	4	573	October 9, 2023
Reduce memory requirements on the GPU when sampling with pm.sampling_jax.sample_numpyro_nuts() v5 gpu	3	1392	March 15, 2023
Batch process capability for pymc.sampling_jax.sample_numpyro_nuts() with GPU? v5 modeling	3	645	September 12, 2022
Out of Memory when using pm.sampling.jax.sample_blackjax_nuts v5 gpu , modeling , jax	2	981	March 23, 2023

Out of memory when "transforming variables" in Numpyro & JAX

Related topics