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.

翻译：蛋白质结构预测有助于理解基因翻译和蛋白质功能，在结构生物学领域日益受到关注且具有重要价值。采用Transformer架构、在蛋白质结构预测中实现原子级精度的AlphaFold模型是一项重大突破。然而，由于计算和内存成本高昂，AlphaFold模型的训练与推理面临巨大挑战。本文提出FastFold——一种高效的AlphaFold训练与推理实现方案。我们提出动态轴向并行（Dynamic Axial Parallelism）与双重异步操作（Duality Async Operations）以提升模型并行的扩展效率。此外，通过自动确定分块策略，我们提出的AutoChunk可在推理过程中降低80%以上的内存成本。实验结果表明，FastFold将总训练时间从11天缩短至67小时，并在长序列推理中实现7.5倍至9.5倍加速。我们还将FastFold扩展至512块GPU，在90.1%并行效率下实现6.02 PetaFLOP/s的聚合吞吐量。