Parallel optimization using glmmTMB

Nafis Sadat

2024-03-18

A new, experimental feature of glmmTMB is the ability to parallelize the optimization process. This vignette shows an example and timing of a simple model fit with and without parallelizing across threads.

If your OS supports OpenMP parallelization and R was installed using OpenMP, glmmTMB will automatically pick up the OpenMP flags from R's Makevars and compile the C++ model with OpenMP support. If the flag is not available, then the model will be compiled with serial optimization only.

Load packages:

library(glmmTMB)
set.seed(1)
nt <- min(parallel::detectCores(),5)

Simulate a dataset with large N:

N <- 3e5
xdata <- rnorm(N, 1, 2)
ydata <- 0.3 + 0.4*xdata + rnorm(N, 0, 0.25)

First, we fit the model serially. We can pass the number of parallelizing process we want using the parallel parameter in glmmTMBcontrol:

system.time(
  model1 <- glmmTMB(formula = ydata ~ 1 + xdata,
                    control = glmmTMBControl(parallel = 1))
  )

Now, we fit the same model using five threads (or as many as possible - 5 in this case):

system.time(
  model2 <- glmmTMB(formula = ydata ~ 1 + xdata,
                    control = glmmTMBControl(parallel = nt))
  )

The speed-up is definitely more visible on models with a much larger number of observations, or in models with random effects.

Here's an example where we have an IID Gaussian random effect. We first simulate the data with 200 groups (our random effect):

xdata <- rnorm(N, 1, 2)
groups <- 200
data_use <- data.frame(obs = 1:N)
data_use <- within(data_use,
{
  
  group_var <- rep(seq(groups), times = nrow(data_use) / groups)
  group_intercept <- rnorm(groups, 0, 0.1)[group_var]
  xdata <- xdata
  ydata <- 0.3 + group_intercept + 0.5*xdata + rnorm(N, 0, 0.25)
})

We fit the random effect model, first with a single thread:

(t_serial <- system.time(
  model3 <- glmmTMB(formula = ydata ~ 1 + xdata + (1 | group_var), data = data_use, control = glmmTMBControl(parallel = 1))
 )
)

Now we fit the same model, but using 5 threads. The speed-up is more noticeable with this model.

(t_parallel <- system.time(
     update(model3,  control = glmmTMBControl(parallel = nt))
 )
)

print(sessionInfo(), locale=FALSE)