Skip to contents

Running emeans on a GPU with torch

Source: vignettes/gpu.qmd

CRAN’s build machines have no GPU, so none of the chunks below run when the package is built; they show the calls to make on your own hardware.

emeans runs the same algorithm on a GPU through the optional torch package, which needs its LibTorch backend installed before anything interesting can happen:

What changes on a device

Where the work happens On the CPU, a C++ kernel assigns points and sums them in double precision. With torch, assignment, sums and updates stay on the device in float32, and only small totals come back to R. CPUR matrixdoubleC++ kerneldirect or BLASupdate in Rk × dGPUtorch tensorfloat32, on deviceaddmm + minindex_add_update on devicek × dBack to R from the GPU: cluster counts at revival checks and the totals at the end of each pass.
  • Computation happens in float32 on the device, and on Apple MPS the counts and the objective are float32 as well.
  • Every random draw still comes from R’s generator, so set.seed() remains in charge of the fit.
  • Results will not match the CPU path bit for bit, because float32 rounds differently, and on CUDA they may not even match each other between runs, since index_add_ sums with atomics in an order that can vary.
  • The result is an ordinary "emeans" object holding R matrices and vectors, so nothing downstream needs to know a GPU was involved.

Apple GPUs (MPS) 

library(emeans)
set.seed(1)
x <- matrix(rnorm(2e6 * 16), ncol = 16)
fit <- emeans(x, 1000, device = "mps")
fit
predict(fit, x[1:10, ])   # assigned on the same device

NVIDIA GPUs (CUDA) 

fit <- emeans(x, 1000, device = "cuda")       # or "cuda:1"

Warning

The CUDA path is experimental: the torch backend has been tested on Apple MPS and on torch’s CPU device but has yet to meet an NVIDIA card. If you give it a whirl, please report what you see at https://github.com/coatless-rpkg/emeans/issues.

Torch on the CPU, and tensors as input

device = "torch:cpu" runs the torch backend on the CPU, which is handy for testing without a GPU. On any torch device x may also be a 2-D torch_tensor, which emeans casts to float32 and moves to the device but does not check for NA, NaN or infinite values, so make sure it arrives clean.

xt <- torch::torch_tensor(x, device = "mps")
fit <- emeans(xt, 1000, device = "mps")