GPU Sampling vs. CPU Sampling for Bayesian MCMC: When Is the Switch Worth It?
A data science team running MCMC-based causal inference in PyMC or Stan eventually hits the same question: move sampling to GPU, or stay on CPU? Getting it wrong either way costs. Accelerating a model that never needed it burns engineering hours nobody gets back. Staying on a CPU sampler past the point where GPU sampling would help stalls model iteration and delays whatever decision depends on the result.
What the benchmark actually measured
One published benchmark tested this tradeoff directly, fitting a hierarchical Bradley-Terry model, the standard pairwise-comparison model used for ranking competitors from win/loss records, to 160,420 professional tennis matches. It compared standard PyMC, PyMC with a JAX backend on CPU, PyMC with JAX on GPU, and Stan, at dataset sizes ranging from a single recent year up to the full match history.
These are the original benchmark's figures, kept here as a planning example from that specific setup rather than a claim about current PyMC, Stan, or JAX performance:
| Sampling configuration | Runtime on the full 160,420-match dataset, in minutes |
|---|---|
| PyMC with JAX on GPU, vectorized | ~2.7 |
| PyMC with JAX on GPU, sequential chains | ~4.5 |
| PyMC with JAX on CPU, parallel chains | ~7.5 |
| Standard PyMC (CPU) | ~12 |
| Stan via cmdstanpy | ~20 |
The fastest GPU configuration was at least 4x faster in wall time than the fastest CPU-only method. The benchmark also reported effective-sample-size-per-second gains up to roughly 11x for the fastest GPU configuration versus CPU-only PyMC and Stan, meaning the GPU run produced more usable independent samples per second, not just a shorter runtime. Even without a GPU, running PyMC with a JAX backend on CPU alone delivered a roughly 2 to 3x effective-sample-size-per-second improvement over standard PyMC and Stan.
Where the crossover sits, and why it matters more than the headline number
The benchmark's runtime curves were close to flat for GPU configurations across dataset sizes, while CPU runtime grew with data volume. That produced a crossover point: below roughly 50,000 observations, GPU sampling ran behind CPU because of fixed per-run overhead; above it, GPU sampling won. For a team deciding where to spend engineering effort, that crossover point matters more than the headline 2.7-minute runtime, because it tells you whether your own dataset is even in the range where the investment pays off.
What does not generalize from this benchmark
The setup that produced these numbers was narrow, which matters for anyone using the crossover point as a planning guide:
- One model class. The benchmark used a single relatively simple hierarchical model. More complex models, including ones with additional covariates or dense covariance structures, may see a different crossover point or a different relative ordering of methods.
- One consumer GPU. Testing ran on a single RTX 2070, not a data-center GPU or a multi-GPU cluster. Results on different hardware, or at larger scale, were not measured here.
- Dated software versions. The benchmark used JAX v0.2.13 and CUDA 10.1. Current JAX and CUDA releases can produce materially different timings for the same model and hardware.
- Unvalidated precision tradeoff. The benchmark noted that switching from double to single precision could theoretically yield up to a further 32x speedup on GPU, based on the consumer GPU's rated double- versus single-precision throughput, but flagged that numerical stability for that switch had not been validated. Treat that figure as an open question, not a confirmed gain.
Before committing to GPU sampling
Check where a candidate dataset falls relative to the roughly 50,000-observation range where this benchmark found GPU sampling starts to win, and treat that boundary as a reference for a similar model class, not a universal threshold. Then reproduce the comparison with the intended model, data, precision, chain configuration, and hardware before choosing a backend. For teams building causal research into decision infrastructure, the relevant question is whether the sampling design supports the decision without changing it. This historical benchmark does not establish current platform performance or imply that Subconscious uses the same stack.