Accelerating AlphaFold2 Inference of Protein Three-Dimensional Structure on the Supercomputer Fugaku

The three-dimensional (3D) structure of proteins provides essential information for understanding life phenomena and drug discovery. Recently, the AI model AlphaFold2 for predicting 3D structures from amino acid sequences has been proposed and has shown very high prediction accuracy. AlphaFold2 is also expected to have high throughput by using accelerators such as GPUs, but its computational performance in a massively parallel computing environment has not been fully discussed. This paper proposes optimization methods to perform inference with AlphaFold2 on the supercomputer Fugaku, the largest CPU-based supercomputer at present. For the multiple sequence alignment (MSA) and template search, a dynamic resource assignment method is proposed and 8.5-fold throughput is achieved for processing 10,000 input sequences. For the computation of the network, a tensor rearrangement for efficient batch matrix multiplication and a memory-efficient attention modules for CPU system are implemented, achieving a 1.3-fold speedup. 6.3-fold throughput is achieved for the entire inference pipeline, demonstrating that Fugaku is a practical CPU system for performing large-scale protein structure inference.