GPRat: Gaussian Process Regression with Asynchronous Tasks

TL;DR

GPRat introduces a fully asynchronous, task-based Gaussian process regression library implemented in C++ with HPX and exposed to Python via pybind11. By tiling covariance computations and leveraging HPX futures, it achieves strong and problem-size scaling on multi-core CPUs, surpassing GPflow and GPyTorch in several metrics while incurring negligible binding overhead. The approach demonstrates substantial prediction speedups and scalable training on up to 64 cores, highlighting the potential of asynchronous runtimes for Python-based AI workloads. Future work targets backpropagation-style training, GPU acceleration, distributed execution, and broader kernel support to improve portability and performance across architectures.

Abstract

Python is the de-facto language for software development in artificial intelligence (AI). Commonly used libraries, such as PyTorch and TensorFlow, rely on parallelization built into their BLAS backends to achieve speedup on CPUs. However, only applying parallelization in a low-level backend can lead to performance and scaling degradation. In this work, we present a novel way of binding task-based C++ code built on the asynchronous runtime model HPX to a high-level Python API using pybind11. We develop a parallel Gaussian process (GP) li- brary as an application. The resulting Python library GPRat combines the ease of use of commonly available GP libraries with the performance and scalability of asynchronous runtime systems. We evaluate the per- formance on a mass-spring-damper system, a standard benchmark from control theory, for varying numbers of regressors (features). The results show almost no binding overhead when binding the asynchronous HPX code using pybind11. Compared to GPyTorch and GPflow, GPRat shows superior scaling on up to 64 cores on an AMD EPYC 7742 CPU for train- ing. Furthermore, our library achieves a prediction speedup of 7.63 over GPyTorch and 25.25 over GPflow. If we increase the number of features from eight to 128, we observe speedups of 29.62 and 21.19, respectively. These results showcase the potential of using asynchronous tasks within Python-based AI applications.

Figures (7)

• Figure 1: Illustration of the first step of a tiled Cholesky decompositionButtari2008 using BLAS operations: The left panel shows a single tile, the middle panel shows 4 tiles per dimension, and the right panel shows 16 tiles per dimension, with the BLAS operations POTRF, TRSM, GEMM, and SYRK.
• Figure 2: Strong scaling runtimes for hyperparameter optimization on up to $64$ cores. The problem size was set to $N$$=$$2^{13}$ training samples. GPRat uses 16 tiles per dimension (see Figure \ref{['fig:tiled_alg']}) and yield identical runtimes.
• Figure 3: Strong scaling runtimes for prediction with full covariance matrix on up to $64$ cores. The problem size was set to $N$$=$$M$$=$$2^{13}$ training samples and test samples. GPRat uses 16 tiles per dimension (see Figure \ref{['fig:tiled_alg']}) and yield identical runtimes.
• Figure 4: Strong scaling runtimes for prediction with full covariance matrix versus variance only on up to $64$ cores for GPRat and GPyTorch with LOVE. The problem size was set to $N$$=$$M$$=$$2^{13}$ training samples and test samples. GPRat uses 16 tiles per dimension (see Figure \ref{['fig:tiled_alg']}).
• Figure 5: Problem size scaling runtimes for prediction with full covariance matrix versus variance only on $64$ cores for GPRat, GPyTorch, and GPyTorch with LOVE. GPRat uses 16 tiles per dimension (see Figure \ref{['fig:tiled_alg']}).