Title:ParaFold: Paralleling AlphaFold for Large-Scale Predictions

Abstract:AlphaFold developed by DeepMind predicts protein structures from the amino acid sequence at or near experimental resolution, solving the 50-year-old protein folding challenge, leading to progress by transforming large-scale genomics data into protein structures. The AlphaFold framework is a mixture of two types of workloads: 1) MSA construction based on CPUs and 2) model inference on GPUs. The first CPU stage dominates the overall runtime, taking up to hours for a single protein due to the large database sizes and I/O bottlenecks. However, GPUs in this CPU stage remain idle, resulting in low GPU utilization and restricting the capacity of large-scale structure predictions. Therefore, we proposed ParaFold, an open-source parallel version of AlphaFold for high throughput protein structure predictions. ParaFold separates the CPU and GPU parts to enable large-scale structure predictions and to improve GPU utilization. ParaFold also effectively reduces the CPU and GPU runtime with two optimizations without compromising the quality of prediction results: using multi-threaded parallelism on CPUs and using optimized JAX compilation on GPUs. We evaluated ParaFold with three datasets of different size and protein lengths. For the small dataset, we evaluated the accuracy and efficiency of optimizations on CPUs and GPUs; For the medium dataset, we demonstrated a typical usage of structure predictions of proteins of different lengths ranging from 77 to 734 residues; For the large dataset, we showed the large-scale prediction capability by running model 1 inferences of $\sim$20,000 small proteins in five hours on one NVIDIA DGX-2. ParaFold offers a rapid and effective approach for high-throughput structure predictions, leveraging the predictive power by running on supercomputers, with shorter time, and at a lower cost.