Extracting draw summaries prohibitively slow for massive models

I’m working with a model with ~ 2 million parameters (~4 million together with transformed parameters). The workflow is:

library(cmdstanr)

model <- cmdstan_model(...)
fit <- model$sample(...)
fit$cmdstan_diagnose()
s <- fit$summary(.cores=6)

Currently the script takes days on the last line, using only 1 core. I’m assuming that it’s working on “repairing draws” somewhere around here: posterior/summarise_draws.R at c58ea9fb2328d58227121b01e2a0469516cdc7a7 · stan-dev/posterior · GitHub, but I’m unsure what it is actually doing and if it is necessary. Note that the cmdstan_diagnose call terminates in reasonable time, and it also needs to touch every draw to do so.

Please let me know if this is better asked on cmdstanr or posterior github. I’m using cmdstanr 2.31 and posterior 1.4.1.

What it’s doing is computing an FFT on each chain for each parameter in order to compute ESS. And R-hat is also slow. Just reading the data in of that size is going to be slow because 4 million params * 1000 draws each is 4G floating point values, or about 32GB of memory just for that.

You are probably best served by (a) cutting down the number of parameters that are actually saved, or if you need them all, (b) cutting down the number of parameters you summarize.

I also don’t know how many posterior draws you are taking, but thinning the draws can also help. We usually recommend only running enough iterations to get an ESS of 100 or so. If there is a very long autocorrelation times, thinning will help.

Thanks for your reply, Bob. Unfortunately I need all the parameters (or rather: need only the transformed ones, but the primitive ones are always saved.) It is unlikely that computing diagnostics is the problem, because that should be happening in parallel, and my code hangs (for days) with only one core used.

I opened a GitHub issue and will continue the discussion there.