LOO with brms::add_criterion crashes the R session

Hi folks,

I’m trying to compare 8 brms models (6 as of now, still waiting for two to finish) using loo with moment_match = TRUE but keep constantly running into crashes where the OS kills the R process. These models are:

M2 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|taxa)) + set_rescor(TRUE), data = df, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
M4 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|q|taxa)) + set_rescor(TRUE), data = df, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
M1 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|gr(binominal, cov = corrmat)) + (1|taxa)) + set_rescor(TRUE), data = df, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
M3 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ state + (1|p|gr(binominal, cov = corrmat)) + (1|q|taxa)) + set_rescor(TRUE), data = df, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))

N2 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|taxa)) + set_rescor(TRUE), data = fred4, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
N4 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|q|taxa)) + set_rescor(TRUE), data = fred4, chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
N1 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|gr(binominal, cov = corrmat)) + (1|taxa)) + set_rescor(TRUE), data = fred4, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))
N3 <- brms::brm(brms::brmsformula(mvbind(F00727, F00679, F00709) ~ 1 + (1|p|gr(binominal, cov = corrmat)) + (1|q|taxa)) + set_rescor(TRUE), data = fred4, data2 = list(corrmat = corrmat), chains = 8, cores = 8, threads = 4, iter = 5000, warmup = 2500, backend = "cmdstanr", save_pars = save_pars(all = TRUE))

The problem happens when adding the loo criterion to the fitted models (reloaded from serialized .Rds files).

M1 <- brms::add_criterion(readRDS("./ScratchData/brms_phylo.Rds"), criterion = "loo", moment_match = TRUE, save_psis = TRUE, cores = 8, pointwise = TRUE)

Every time, the R process gets killed by the OS because of it maxxes out the RAM (the server I’m using has 32 virtual cores and 126 GiB of RAM). This happens shortly after I see the output

Recompiling the model with 'rstan'
Recompilation done

Setting pointwise = TRUE didn’t help either, unfortunately!

Please see the screenshot for reference:

The models I’m trying to add loo to are relatively big (1.2 GiB to 2.3 GiBs each on disk), so I partly understand the memory issues. However, brms::add_criterion works just fine when moment_match is set to FALSE. In fact I was able to add loo to all the six models and run loo::loo_compare in that case - but got the following warnings about pareto K estimates advising me to set moment_match to TRUE.

> M1 <- brms::add_criterion(M1, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 267 observations with a pareto_k > 0.7 in model 'M1'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> M2 <- brms::add_criterion(M2, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 696 observations with a pareto_k > 0.7 in model 'M2'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> M4 <- brms::add_criterion(M4, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 625 observations with a pareto_k > 0.7 in model 'M4'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> N1 <- brms::add_criterion(N1, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 257 observations with a pareto_k > 0.7 in model 'N1'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> N2 <- brms::add_criterion(N2, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 720 observations with a pareto_k > 0.7 in model 'N2'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> N4 <- brms::add_criterion(N4, criterion = "loo", moment_match = FALSE, save_psis = TRUE, cores = 8, pointwise = FALSE)
Warning messages:
1: Ignoring relative efficiencies as some were NA. See argument 'r_eff' in ?loo::loo for more details. 
2: Found 662 observations with a pareto_k > 0.7 in model 'N4'. We recommend to set 'moment_match = TRUE' in order to perform moment matching for problematic observations.  

> compres <- loo::loo_compare(M1, M2, M4, N1, N2, N4) 
Warning message:
Not all models have the same y variable. ('yhash' attributes do not match) 
> print(compres)
 model elpd_diff se_diff p_worse diag_diff        diag_elpd
    M4       0.0     0.0      NA           625 k_psis > 0.7
    N4     -21.6   431.0    0.52           662 k_psis > 0.7
    M1    -142.2    81.2    0.96           267 k_psis > 0.7
    N1    -144.5   416.2    0.64           257 k_psis > 0.7
    M2    -844.7    81.2    1.00           696 k_psis > 0.7
    N2    -865.7   409.7    0.98           720 k_psis > 0.7

This is why I decided to proceed with moment_match = TRUE but the process ends up being extremely demanding on the RAM. With options(mc.cores = 1) options(loo.cores = 1) I was able to do brms::add_criterion for one model but with more cores specified, the RAM use goes over the roof (each core ends up using about 22 GiBs of RAM) and the process gets killed by the OS. Doesn’t help having to load in these huge .Rds models into the process’s memory as well.

Unfortunately, I cannot afford to re-run the brms fits (I’ve already completed, serialized and backed up 6 fits and have two models at 20% needing 20 more days to complete). I’d appreciate if anyone can help me find a workaround to do a moment matched loo comparison? I don’t even mind the non moment matched loo if I ignoring the warnings is okay. Also don’t understand why it says Not all models have the same y variable. ('yhash' attributes do not match) even though all the Y variables are identical across all the brms models. Please feel free to ask for more info if needed. Thanks!

System specs:

Ubuntu 24.04.4 LTS (Noble Numbat)
R version - 4.6.1
brms version - 2.23.0
cmdstanr version - 0.9.0
CmdStan version - 2.39.0
rstan version - 2.32.7
Stan version - 2.32.2
StanHeaders version - 2.32.10