Training POWERUP models in parallel

fit_powerup_models() trains the requested targets sequentially. For analyses with many targets, you can speed up training by dividing the targets into small batches and training several batches at the same time.

In this example, each batch contains 10 targets. Models within each batch are still trained sequentially by fit_powerup_models(). The parallelization occurs across batches.

For example, with four workers:

Worker 1: targets 1-10
Worker 2: targets 11-20
Worker 3: targets 21-30
Worker 4: targets 31-40

When a worker finishes a batch, it can begin another batch until all targets have been trained.

Start from prepared POWERUP data

Prepare the data once using prepare_powerup_data(). See Preparing data for POWERUP for the complete data-preparation workflow.

The selected targets are stored in the prepared object:

targets <- prepared$perturbations$perturbation
length(targets)

For example, 1,000 targets can be divided into 100 batches of 10 targets each.

Create target batches

batch_size <- 10L
batches <- split(targets, ceiling(seq_along(targets) / batch_size))
length(batches)

Each element of batches now contains up to 10 targets.

Train batches in parallel

Choose how many batches to train at the same time based on the CPU and memory available on your computer. Here, four workers are used.

workers <- 4L
batch_dir <- "powerup_model_batches"

dir.create(batch_dir, showWarnings = FALSE)
batch_dir <- normalizePath(batch_dir, mustWork = TRUE)

cl <- parallel::makeCluster(workers)

parallel::clusterEvalQ(cl, library(powerup))

parallel::clusterExport(cl, c("prepared", "batches", "batch_dir"))

train_batch <- function(i) {
  batch_targets <- batches[[i]]
  batch_file <- file.path(batch_dir, sprintf("models_batch_%03d.rds", i))

  models <- fit_powerup_models(prepared, models_to_make = batch_targets, seed = 123L, n_threads = 1L)

  saveRDS(models, batch_file)
  batch_file
}

parallel::parLapplyLB(cl, seq_along(batches), train_batch)
parallel::stopCluster(cl)

The workers are created once and remain active throughout the training run. Each worker loads POWERUP once, then processes one 10-target batch at a time. parLapplyLB() assigns another available batch when a worker finishes its current batch.

Setting n_threads = 1L gives each worker one CPU thread per model. If additional CPU threads are available for each worker, n_threads can be increased to allow XGBoost to use multiple threads while fitting each model.

Each completed batch is saved separately:

powerup_model_batches/
  models_batch_001.rds
  models_batch_002.rds
  models_batch_003.rds
  ...

These files also provide convenient checkpoints for a long training run.

Merge the trained batches

After all batches have finished, load the saved model lists and combine them in the original batch order.

batch_files <- file.path(batch_dir, sprintf("models_batch_%03d.rds", seq_along(batches)))
stopifnot(all(file.exists(batch_files)))

model_batches <- lapply(batch_files, readRDS)
models <- do.call(c, model_batches)

stopifnot(identical(names(models), targets))

The merged models object has the same structure as the named model list returned by a single call to fit_powerup_models().

You can optionally save the combined object:

saveRDS(models, "powerup_models.rds")

Continue with downstream analysis

The merged model list can be used normally with the rest of the POWERUP workflow.

model_summary <- summarize_models(models)

models <- add_powerup_predictions(models, prepared)
predictions <- summarize_predictions(models, format = "long")

Prediction and result interpretation can then proceed as described in Interpreting POWERUP results. For SHAP-based model explanations, see Explaining POWERUP predictions.

Choosing the number of workers

A simple starting point is four workers with batches of 10 targets. Increasing the number of workers can train more batches simultaneously, but each worker is a separate R process and therefore requires additional memory.

workers <- 4L
batch_size <- 10L

If more CPU and memory are available, you can increase workers. Batch sizes around 10-20 targets are also reasonable for larger training runs.

Session information

sessionInfo()
#> R version 4.4.2 (2024-10-31)
#> Platform: aarch64-apple-darwin20
#> Running under: macOS Sequoia 15.7.3
#> 
#> Matrix products: default
#> BLAS:   /Library/Frameworks/R.framework/Versions/4.4-arm64/Resources/lib/libRblas.0.dylib 
#> LAPACK: /Library/Frameworks/R.framework/Versions/4.4-arm64/Resources/lib/libRlapack.dylib;  LAPACK version 3.12.0
#> 
#> locale:
#> [1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8
#> 
#> time zone: America/New_York
#> tzcode source: internal
#> 
#> attached base packages:
#> [1] stats     graphics  grDevices utils     datasets  methods   base     
#> 
#> other attached packages:
#> [1] powerup_1.0.95
#> 
#> loaded via a namespace (and not attached):
#>  [1] digest_0.6.37     desc_1.4.3        R6_2.6.1          fastmap_1.2.0    
#>  [5] xfun_0.51         cachem_1.1.0      knitr_1.50        htmltools_0.5.8.1
#>  [9] rmarkdown_2.29    lifecycle_1.0.4   cli_3.6.5         sass_0.4.9       
#> [13] pkgdown_2.2.0     textshaping_1.0.0 jquerylib_0.1.4   systemfonts_1.2.2
#> [17] compiler_4.4.2    tools_4.4.2       ragg_1.5.1        bslib_0.9.0      
#> [21] evaluate_1.0.3    yaml_2.3.10       jsonlite_2.0.0    rlang_1.1.6      
#> [25] fs_1.6.5          htmlwidgets_1.6.4