ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours

Feiwen Zhu,Arkadiusz Nowaczynski,Rundong Li,Jie Xin,Yifei Song,Michal Marcinkiewicz,Sukru Burc Eryilmaz,Jun Yang,Michael Andersch

2024-04-17

Abstract:AlphaFold2 has been hailed as a breakthrough in protein folding. It can rapidly predict protein structures with lab-grade accuracy. However, its implementation does not include the necessary training code. OpenFold is the first trainable public reimplementation of AlphaFold. AlphaFold training procedure is prohibitively time-consuming, and gets diminishing benefits from scaling to more compute resources. In this work, we conducted a comprehensive analysis on the AlphaFold training procedure based on Openfold, identified that inefficient communications and overhead-dominated computations were the key factors that prevented the AlphaFold training from effective scaling. We introduced ScaleFold, a systematic training method that incorporated optimizations specifically for these factors. ScaleFold successfully scaled the AlphaFold training to 2080 NVIDIA H100 GPUs with high resource utilization. In the MLPerf HPC v3.0 benchmark, ScaleFold finished the OpenFold benchmark in 7.51 minutes, shown over $6\times$ speedup than the baseline. For training the AlphaFold model from scratch, ScaleFold completed the pretraining in 10 hours, a significant improvement over the seven days required by the original AlphaFold pretraining baseline.

Machine Learning,Artificial Intelligence,Distributed, Parallel, and Cluster Computing,Quantitative Methods

What problem does this paper attempt to address?

This paper focuses on reducing the initial training time of the AlphaFold protein structure prediction model. AlphaFold is a model that can quickly and accurately predict protein structures, but its training process is time-consuming and typically takes more than 7 days. The research team analyzed the training process of the open-source implementation OpenFold and found that low communication efficiency and underutilization of computing resources were the main bottlenecks. To address these issues, they proposed ScaleFold, a systematic training method that optimized communication and computing efficiency. ScaleFold successfully scaled AlphaFold's training to 2,080 NVIDIA H100 GPUs and reduced the training time of OpenFold to 7.51 minutes in the MLPerf HPC v3.0 benchmark test, more than 6 times faster than the baseline. For training AlphaFold from scratch, ScaleFold reduced the pre-training time to 10 hours, significantly outperforming the original 7-day training time. In summary, the main contributions of the paper are: 1. Identifying key factors hindering the scalability of AlphaFold training, namely inefficient communication and underutilization of computing resources. 2. Proposing ScaleFold, a scalable and systematic training method for AlphaFold, to address the above issues. 3. Experimental results demonstrate the effectiveness of ScaleFold, setting new records for AlphaFold pre-training and the MLPerf HPC benchmark. Through these optimizations, ScaleFold improves training efficiency and accelerates research progress in protein structure prediction.

ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours

FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours

ParaFold: Paralleling AlphaFold for Large-Scale Predictions

OpenFold provides insights into AlphaFold2's learning behavior

Uni-Fold: An Open-Source Platform for Developing Protein Folding Models beyond AlphaFold

HelixFold: An Efficient Implementation of AlphaFold2 using PaddlePaddle

Proteome-scale Deployment of Protein Structure Prediction Workflows on the Summit Supercomputer

PiFold: Toward effective and efficient protein inverse folding

ColabFold: making protein folding accessible to all

OPUS-Fold3: a Gradient-Based Protein All-Atom Folding and Docking Framework on TensorFlow

Easy and accurate protein structure prediction using ColabFold

DeepFold: Enhancing Protein Structure Prediction through Optimized Loss Functions, Improved Template Features, and Re-optimized Energy Function

Distance-AF: Modifying Predicted Protein Structure Models by Alphafold2 with User-Specified Distance Constraints

Highly accurate protein structure prediction with AlphaFold

GDFold2: a fast and parallelizable protein folding environment with freely defined objective functions

IMPROVING ALPHAFOLD2 PERFORMANCE WITH A GLOBAL METAGENOMIC & BIOLOGICAL DATA SUPPLY CHAIN

Applying and improving AlphaFold at CASP14

ManyFold: an efficient and flexible library for training and validating protein folding models

DeepFoldit -- A Deep Reinforcement Learning Neural Network Folding Proteins

AlphaFold Distillation for Protein Design

Enhancing Protein Function Prediction Performance by Utilizing AlphaFold-Predicted Protein Structures