title: How to efficiently parallelize brms::brm?
body: <h1>Problem summary</h1>
  
  <p>I am fitting a <code>brms::brm_multiple()</code> model to a large dataset where missing data has been imputed using the <code>mice</code> package. The size of the dataset makes the use of parallel processing very desirable. However, it isn't clear to me how to best use the compute resources because I am unclear about how <code>brms</code> divides sampling on the imputed dataset among cores.</p>
  
  <p><strong>How can I choose the following to maximize efficient use of compute resources?</strong></p>
  
  <ul>
  <li>number of imputations (<code>m</code>)</li>
  <li>number of chains (<code>chains</code>)</li>
  <li>number of cores (<code>cores</code>)</li>
  </ul>
  
  <h1>Conceptual example</h1>
  
  <p>Let's say that I naively (or deliberately foolishly for sake of example) choose <code>m = 5</code>, <code>chains = 10</code>, <code>cores = 24</code>. There are thus 5 x 10 = 50 chains to be allocated among 24 cores reserved on the HPC. Without parallel processing, this would take ~50 time units (excluding compiling time).</p>
  
  <p>I can imagine three parallelization strategies for <code>brms_multiple()</code>, but there may be others:</p>
  
  <p><strong>Scenario 1:</strong> Imputed datasets in parallel, associated chains in serial</p>
  
  <p>Here, each of the 5 imputations is allocated to it's own processor which runs through the 10 chains in serial. The processing time is 10 units (a 5x speed improvement vs. non-parallel processing), but poor planning has wasted 19 cores x 10 time units = 190 core time units (ctu; =80% of the reserved compute resources). The efficient solution would be to set <code>cores</code> = <code>m</code>.</p>
  
  <p><strong>Scenario 2:</strong> Imputed datasets in serial, associated chains in parallel</p>
  
  <p>Here, the sampling begins by taking the first imputed dataset and running one of the chains for that dataset on each of 10 different cores. This is then repeated for the remaining four imputed datasets. The processing takes 5 time units (a 10x speed improvement over serial processing &amp; a 2x improvement over Scenario 1). However, here too compute resources are wasted: 14 cores x 5 time units = 70 ctu. The efficient solution would be to set <code>cores</code> = <code>chains</code></p>
  
  <p><strong>Scenario 3:</strong> Free-for-all, wherein each core takes on a pending imputation/chain combination when it becomes available until all are processed.</p>
  
  <p>Here, the sampling begins by allocating all 24 cores, each one to one of the 50 pending chains. After they finish their iterations, a second batch of 24 chains is processed, bringing the total chains processed to 48. But now there are only two chains pending and 22 cores sit idle for 1 time unit. The total processing time is 3 time units, and the wasted compute resource is 22 ctu. The efficient solution would be to set <code>cores</code> to a multiple of <code>m</code> x <code>chains</code>.</p>
  
  <h1>Minimal reproducible example</h1>
  
  <p>This code compares the compute time using an example modified from a <a href="https://cran.r-project.org/web/packages/brms/vignettes/brms_missings.html" rel="nofollow noreferrer">brms vignette</a>. Here we'll set <code>m</code> = 10, <code>chains</code> = 6, and <code>cores</code> = 4. This makes for a total of 60 chains to be processed. Under these conditions, I would expect speed improvement (vs. serial processing) is as follows*:</p>
  
  <ul>
  <li>Scenario 1: 60/(6 chains x ceiling(10 m / 4 cores)) = 3.3x</li>
  <li>Scenario 2: 60/(ceiling(6 chains / 4 cores) x 10 m) = 3.0x</li>
  <li>Scenario 3: 60/ceiling((6 chains x 10 m) / 4 cores) = 4.0x</li>
  </ul>
  
  <p>*(ceiling/rounding up is used because chains cannot be subdivided among cores)</p>
  
  <pre><code>library(brms)
  library(mice)
  library(tictoc)  # convenience functions for timing
  
  # Load data
  data("nhanes", package = "mice")
  
  # There are 10 imputations x 6 chains = 60 total chains to be processed
  imp &lt;- mice(nhanes, m = 10, print = FALSE, seed = 234023)
  
  # Fit the model first to get compilation out of the way
  fit_base &lt;- brm_multiple(bmi ~ age*chl, data = imp, chains = 6,
                           iter = 10000, warmup = 2000)
  
  # Use update() function to avoid re-compiling time
  # Serial processing (127 sec on my machine)
  tic()  # start timing
  fit_serial &lt;- update(fit_base, .~., cores = 1L)
  t_serial &lt;- toc()  # stop timing
  t_serial &lt;- diff(unlist(t_serial)[1:2])  # calculate seconds elapsed
  
  # Parallel processing with 3 cores (82 sec)
  tic()
  fit_parallel &lt;- update(fit_base, .~., cores = 4L)
  t_parallel &lt;- toc()
  t_parallel &lt;- diff(unlist(t_parallel)[1:2])  # calculate seconds elapsed
  
  # Calculate speed up ratio
  t_serial/t_parallel  # 1.5x
  </code></pre>
  
  <p>Clearly I am missing something. I can't distinguish between the scenarios with this approach.</p>