Hierarchical model with copula on GPU: NumPyro or nutpie using PyMC 5.19?

Questions v5
,
1

My compute cluster only supports PyMC 5.19 using Python 3.12.4, and I need to fit a bivariate joint model with copulas and two quite complicated marginal models (~2k parameters).

It took ~5 hours for 5k warmups + 5k iterations running on an H100 GPU, which is 3x faster than running the same on a CPU. However, it’s still quite slow, and I still want to speed up the sampling time so my Bayesian workflow/model iteration can fail faster.

Any suggestions on how to speed up sampling? Even though Nutpie is what I use in CPU sampling and the latest v6 uses it as default, is it also better than NumPyro in lower versions in v5?

The following is my current sampling function:

def _sample_model(

model: Any,

draws: int,

tune: int,

seed: int,

chains: int = 4,

cores: int = 4,

max_treedepth: int = 12,

target_accept: float = 0.99,

nuts_sampler: str | None = None,

nutpie_backend: str | None = None,

sampling_device: str = "gpu",

chain_method: str | None = None,

postprocessing_backend: str | None = None,

jitter: bool = True,

thin: int = 1,

    ) -> Any:

if draws <= 0:

raise ValueError("--fit-draws must be positive when no posterior file is supplied.")

    _configure_runtime_cache_dirs()

import pymc as pm

    sampling = _resolve_sampling_config(

sampling_device=sampling_device,

chains=chains,

cores=cores,

target_accept=target_accept,

nuts_sampler=nuts_sampler,

nutpie_backend=nutpie_backend,

chain_method=chain_method,

postprocessing_backend=postprocessing_backend,

thin=thin,

max_treedepth=max_treedepth,

jitter=jitter,

    )

print(f"Sampling configuration: {json.dumps(sampling, indent=2)}")

if sampling["sampling_device"] == "gpu":

        _jax_runtime_info(require_gpu=True)

    init_method = "jitter+adapt_diag" if jitter else "adapt_diag"

with model:

        idata = pm.sample(

draws=draws,

tune=tune,

chains=sampling["chains"],

cores=sampling["cores"],

target_accept=sampling["target_accept"],

nuts_sampler=sampling["nuts_sampler"],

nuts_sampler_kwargs=sampling["nuts_sampler_kwargs"],

init=init_method,

progressbar=False,

random_seed=seed,

compute_convergence_checks=True,

return_inferencedata=True,

        )

if sampling["thin"] > 1:

        idata = idata.isel(draw=slice(None, None, sampling["thin"]))

return idata

Thanks!

2

The usual suspect in slow fitting is the model itself and how it’s coded or its match to the data. Andrew Gelman calls this the “folk theorem of statistical computing.”

I’d also be suspicious that the GPU code on an H100 is only 3 times as fast as CPU code. Is it having to bounce in-and-out of kernel or are you using a Blackjax-based sampler that’s itself coded in JAX?

3

Hi Bob,

That makes sense. I am using the JAX backend with NumPyro, so I do not think the main issue is CPU–GPU transfer. The data are also relatively small, roughly (70 \times 200 \times 2), so my current interpretation is that the H100 is probably underutilized: the computation consists of many small sequential HMC/NUTS operations rather than large matrix operations that would fully exploit the GPU.

The harder issue seems to be posterior geometry. I need to take the marginal hierarchical models as given and then build the bivariate copula model on top of them. The samplers used in the original marginal-model papers are not HMC-based, but hybrid MCMC schemes. Even after reparameterizing some marginal priors, I still see a surprisingly large number of divergences, around 25–30% out of 10k iterations. However, the predictive evaluation metrics from my PyMC/NumPyro replication are almost identical to those from the original sampler, usually matching up to two decimal places.

I am happy with the roughly 3x speedup on the H100, but I would certainly be even happier if there were further ways to speed up the sampling:)

4

A better sampler like nutpie (even better with low-rank adaptation) may help. Your data is small, that I would invite you to try the numba backend (default if you go with nutpie). It’s usually the best performing on CPU, although not always. Also suggest you update to newer pymc > 6.x, we did a lot of improvements on the numba backend: PyMC 6.0 & PyTensor 3.0: ecosystem updates — PyMC project website

If you don’t see any improvements, we may need to look at the model like bob hinted at