[Paper Review] FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours

Shenggan Cheng, Zhao, Xuanlei+8 more|arXiv (Cornell University)|Mar 2, 2022

Protein Structure and Dynamics17 citations

TL;DR

FastFold accelerates AlphaFold training from 11 days to 67 hours and enables 7.5–9.5× faster long-sequence inference by introducing Dynamic Axial Parallelism, Duality Async Operations, and AutoChunk—low-level optimizations that reduce memory usage by over 80% and improve GPU scaling efficiency. The system achieves 6.02 PetaFLOP/s throughput on 512 A100 GPUs with 90.1% parallel efficiency.

ABSTRACT

Protein structure prediction helps to understand gene translation and protein function, which is of growing interest and importance in structural biology. The AlphaFold model, which used transformer architecture to achieve atomic-level accuracy in protein structure prediction, was a significant breakthrough. However, training and inference of the AlphaFold model are challenging due to its high computation and memory cost. In this work, we present FastFold, an efficient implementation of AlphaFold for both training and inference. We propose Dynamic Axial Parallelism and Duality Async Operations to improve the scaling efficiency of model parallelism. Besides, AutoChunk is proposed to reduce memory cost by over 80% during inference by automatically determining the chunk strategy. Experimental results show that FastFold reduces overall training time from 11 days to 67 hours and achieves 7.5X - 9.5X speedup for long-sequence inference. Furthermore, we scale FastFold to 512 GPUs and achieve an aggregate throughput of 6.02 PetaFLOP/s with 90.1% parallel efficiency.

Motivation & Objective

• Address the high computational and memory costs of training and inferring the AlphaFold model at scale.
• Overcome limitations in model parallelism scalability and communication overhead in multi-GPU training.
• Reduce memory consumption during long-sequence inference to enable processing of extremely long proteins.
• Achieve high throughput and parallel efficiency on large-scale GPU clusters (e.g., 512 A100 GPUs).
• Develop a general-purpose optimization framework applicable to other attention-based models in structural biology and beyond.

Proposed method

• Proposes Dynamic Axial Parallelism, a model parallelism strategy that optimizes tensor partitioning across GPUs with reduced communication overhead.
• Introduces Duality Async Operations, a PyTorch extension enabling computation-communication overlap to hide communication latency.
• Develops AutoChunk, a technique that automatically determines optimal sequence chunking for memory reduction during inference.
• Applies low-level kernel optimizations to the Evoformer block to accelerate computation and reduce memory footprint.
• Scales training to 512 A100 GPUs using hybrid parallelism and communication-optimized training pipelines.
• Employs memory-aware chunking with dynamic adjustment to balance speed and memory usage during inference.

Experimental results

Research questions

• RQ1 Can model parallelism be optimized to reduce communication overhead and improve scaling efficiency in large-scale protein structure prediction?
• RQ2 How can computation-communication overlap be effectively achieved in attention-based models like AlphaFold?
• RQ3 Can automated chunking strategies outperform expert-tuned configurations in reducing memory usage during long-sequence inference?
• RQ4 To what extent can low-level kernel and memory optimizations reduce training and inference costs in attention-based models?
• RQ5 Can the proposed optimizations be generalized to other large attention models beyond AlphaFold?

Key findings

• FastFold reduces AlphaFold training time from 11 days to 67 hours on 512 A100 GPUs, achieving a 10.5× speedup.
• The system achieves 6.02 PetaFLOP/s of aggregate throughput with 90.1% parallel efficiency on 512 GPUs.
• AutoChunk reduces memory usage by over 80% during inference, surpassing expert-designed chunking strategies by 30.6%–34.4% in memory efficiency.
• FastFold achieves 7.5–9.5× speedup for long-sequence inference compared to baseline implementations.
• AutoChunk improves inference speed by 12% compared to OpenFold with expert-tuned chunking (64), while losing only 4% compared to unchunked inference.
• The optimizations are generalizable and applicable to other attention-based models such as RoseTTAFold, MSA Transformer, and video transformers.