Threading in Stan — threading • brms

Use threads for within-chain parallelization in Stan via the brms interface. Within-chain parallelization is experimental! We recommend its use only if you are experienced with Stan's reduce_sum function and have a slow running model that cannot be sped up by any other means.

Usage

threading(threads = NULL, grainsize = NULL, static = FALSE, force = FALSE)

Arguments

threads: Number of threads to use in within-chain parallelization.
grainsize: Number of observations evaluated together in one chunk on one of the CPUs used for threading. If NULL (the default), grainsize is currently chosen as max(100, N / (2 * threads)), where N is the number of observations in the data. This default is experimental and may change in the future without prior notice.
static: Logical. Apply the static (non-adaptive) version of reduce_sum? Defaults to FALSE. Setting it to TRUE is required to achieve exact reproducibility of the model results (if the random seed is set as well).
force: Logical. Defaults to FALSE. If TRUE, this will force the Stan model to compile with threading enabled without altering the Stan code generated by brms. This can be useful if your own custom Stan functions use threading internally.

Value

A brmsthreads object which can be passed to the threads argument of brm and related functions.

Details

The adaptive scheduling procedure used by reduce_sum will prevent the results to be exactly reproducible even if you set the random seed. If you need exact reproducibility, you have to set argument static = TRUE which may reduce efficiency a bit.

To ensure that chunks (whose size is defined by grainsize) require roughly the same amount of computing time, we recommend storing observations in random order in the data. At least, please avoid sorting observations after the response values. This is because the latter often cause variations in the computing time of the pointwise log-likelihood, which makes up a big part of the parallelized code.

Examples

# \dontrun{
# this model just serves as an illustration
# threading may not actually speed things up here
fit <- brm(count ~ zAge + zBase * Trt + (1|patient),
           data = epilepsy, family = negbinomial(),
           chains = 1, threads = threading(2, grainsize = 100),
           backend = "cmdstanr")
#> Start sampling
#> Running MCMC with 1 chain, with 2 thread(s) per chain...
#> 
#> Chain 1 Iteration:    1 / 2000 [  0%]  (Warmup) 
#> Chain 1 Iteration:  100 / 2000 [  5%]  (Warmup) 
#> Chain 1 Iteration:  200 / 2000 [ 10%]  (Warmup) 
#> Chain 1 Iteration:  300 / 2000 [ 15%]  (Warmup) 
#> Chain 1 Iteration:  400 / 2000 [ 20%]  (Warmup) 
#> Chain 1 Iteration:  500 / 2000 [ 25%]  (Warmup) 
#> Chain 1 Iteration:  600 / 2000 [ 30%]  (Warmup) 
#> Chain 1 Iteration:  700 / 2000 [ 35%]  (Warmup) 
#> Chain 1 Iteration:  800 / 2000 [ 40%]  (Warmup) 
#> Chain 1 Iteration:  900 / 2000 [ 45%]  (Warmup) 
#> Chain 1 Iteration: 1000 / 2000 [ 50%]  (Warmup) 
#> Chain 1 Iteration: 1001 / 2000 [ 50%]  (Sampling) 
#> Chain 1 Iteration: 1100 / 2000 [ 55%]  (Sampling) 
#> Chain 1 Iteration: 1200 / 2000 [ 60%]  (Sampling) 
#> Chain 1 Iteration: 1300 / 2000 [ 65%]  (Sampling) 
#> Chain 1 Iteration: 1400 / 2000 [ 70%]  (Sampling) 
#> Chain 1 Iteration: 1500 / 2000 [ 75%]  (Sampling) 
#> Chain 1 Iteration: 1600 / 2000 [ 80%]  (Sampling) 
#> Chain 1 Iteration: 1700 / 2000 [ 85%]  (Sampling) 
#> Chain 1 Iteration: 1800 / 2000 [ 90%]  (Sampling) 
#> Chain 1 Iteration: 1900 / 2000 [ 95%]  (Sampling) 
#> Chain 1 Iteration: 2000 / 2000 [100%]  (Sampling) 
#> Chain 1 finished in 1.4 seconds.
summary(fit)
#>  Family: negbinomial 
#>   Links: mu = log 
#> Formula: count ~ zAge + zBase * Trt + (1 | patient) 
#>    Data: epilepsy (Number of observations: 236) 
#>   Draws: 1 chains, each with iter = 2000; warmup = 1000; thin = 1;
#>          total post-warmup draws = 1000
#> 
#> Multilevel Hyperparameters:
#> ~patient (Number of levels: 59) 
#>               Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
#> sd(Intercept)     0.55      0.07     0.42     0.69 1.00      270      534
#> 
#> Regression Coefficients:
#>            Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
#> Intercept      1.78      0.12     1.53     2.01 1.00      326      438
#> zAge           0.10      0.08    -0.07     0.25 1.00      320      411
#> zBase          0.70      0.11     0.49     0.91 1.00      320      392
#> Trt1          -0.26      0.16    -0.57     0.02 1.01      281      540
#> zBase:Trt1     0.05      0.16    -0.25     0.35 1.00      361      547
#> 
#> Further Distributional Parameters:
#>       Estimate Est.Error l-95% CI u-95% CI Rhat Bulk_ESS Tail_ESS
#> shape     7.30      1.75     4.53    11.37 1.00      609      605
#> 
#> Draws were sampled using sample(hmc). For each parameter, Bulk_ESS
#> and Tail_ESS are effective sample size measures, and Rhat is the potential
#> scale reduction factor on split chains (at convergence, Rhat = 1).
# }