Accelerating Genome- and Phenome-Wide Association Studies using GPUs - A case study using data from the Million Veteran Program

Abstract

The expansion of biobanks has significantly propelled genomic discoveries yet the sheer scale of data within these repositories poses formidable computational hurdles, particularly in handling extensive matrix operations required by prevailing statistical frameworks. In this work, we introduce computational optimizations to the SAIGE (Scalable and Accurate Implementation of Generalized Mixed Model) algorithm, notably employing a GPU-based distributed computing approach to tackle these challenges. We applied these optimizations to conduct a large-scale genome-wide association study (GWAS) across 2,068 phenotypes derived from electronic health records of 635,969 diverse participants from the Veterans Affairs (VA) Million Veteran Program (MVP). Our strategies enabled scaling up the analysis to over 6,000 nodes on the Department of Energy (DOE) Oak Ridge Leadership Computing Facility (OLCF) Summit High-Performance Computer (HPC), resulting in a 20-fold acceleration compared to the baseline model. We also provide a Docker container with our optimizations that was successfully used on multiple cloud infrastructures on UK Biobank and All of Us datasets where we showed significant time and cost benefits over the baseline SAIGE model.

Fig. 1

Overview of genomic analysis in multiple population groups. a) Schematic representation illustrating the diverse set of GIA population groups. The analysis covers a deep catalog of traits extracted from electronic health records, clinical laboratory tests, vital signs, and survey responses. b) Chart categorizing traits into binary or quantitative types across different population groups. The height of each bar corresponds to the number of traits in each category, providing an overview of the trait composition for subsequent genomic analyses.

Fig. 2

Distributed BLAS gemv(), matrix-vector multiplication, using GPUs on the cluster. The columns of matrix $A$ are distributed and preloaded on GPUs, with node $i$ having columns with indices from $s_i$ to $e_i$, and these columns are distributed on GPUs on that node. To compute $p=A{A}^{T}v$, we first broadcast $\mathrm{pv}$ to GPUs, and each node computes a partial solution on GPUs. These partial solutions are aggregated to compute a solution $p$.

Fig. 3

Comparative Performance of GPU and CPU Implementations in SAIGE Step one - This figure compares the execution time for each iteration of matrix operations in SAIGE Step one for the European population group. a) Demonstration of the time required for a single PCG iteration on a GPU, showcasing the efficient parallelization within the GPU. b) Contrast with the OpenMP implementation on CPUs, emphasizing the significant speed improvement achieved with GPU acceleration. As the genotype matrix size increases, the advantage of using the GPU version becomes more pronounced, as highlighted by the diminishing execution time on the GPU compared to the CPU.

Fig. 4

GPU Node Requirements and Memory Impact – GPU node requirement highlight the linear relationship between genotype matrix size and the required number of nodes, offering insights into efficient GPU utilization. The GPU node requirement factored in the GPU memory, the byte size of a single precision floating-point number, and the conversion between bytes and gigabytes. A) Impact of changing the memory available in the GPU. B) Impact of changing number of genotype variants in the input matrix and fixing the GPU memory to 16 gigabytes per GPU, emphasizing considerations for diverse biobank cohorts and computational environments.

Fig. 5

SAIGE step one run time for All of Us data. The figure shows the time comparison of running SAIGE step one for the T2D phenotype on the Google Cloud Platform for the 5 population groups (EUR, AFR, AMR, EAS, SAS). The analysis was executed on 4 NVIDIA T4 GPUs for the SAIGE-GPU version and a 64-CPU VM for the SAIGE-CPU version.