[论文解读] 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 DynamicsBiochemistry, Genetics and Molecular Biology被引用 17

一句话总结

FastFold 将 AlphaFold 训练时间从 11 天缩短至 67 小时，并通过引入动态轴向并行、对偶异步操作和自动分块技术，使长序列推理速度提升 7.5–9.5 倍。这些底层优化使内存使用量减少 80% 以上，并显著提升 GPU 扩展效率。该系统在 512 张 A100 GPU 上实现了 6.02 PetaFLOP/s 的吞吐量，平行效率达 90.1%。

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.

研究动机与目标

• 解决大规模训练和推理 AlphaFold 模型时计算与内存成本过高的问题。
• 克服多 GPU 训练中模型并行扩展性不足与通信开销过大的限制。
• 在长序列推理过程中降低内存消耗，以支持超长蛋白质的处理。
• 在大规模 GPU 集群（如 512 张 A100 GPU）上实现高吞吐量与高并行效率。
• 开发一种适用于其他基于注意力机制模型的通用优化框架，涵盖结构生物学及其他领域。

提出的方法

• 提出动态轴向并行，一种模型并行策略，通过优化张量在 GPU 间的划分方式，显著降低通信开销。
• 引入对偶异步操作，作为 PyTorch 扩展，支持计算与通信重叠，以隐藏通信延迟。
• 开发自动分块技术，可自动确定最优序列分块策略，以减少推理过程中的内存占用。
• 对 Evoformer 模块实施底层内核优化，以加速计算并减少内存占用。
• 采用混合并行策略与通信优化的训练流水线，将训练扩展至 512 张 A100 GPU。
• 在推理过程中采用感知内存的分块策略，并结合动态调整机制，以在速度与内存使用之间实现平衡。

实验结果

研究问题

• RQ1 模型并行是否可通过优化降低通信开销，并提升大规模蛋白质结构预测中的扩展效率？
• RQ2 在类似 AlphaFold 的基于注意力机制的模型中，如何有效实现计算与通信的重叠？
• RQ3 自动化分块策略是否能在长序列推理中比专家调优的配置更有效地降低内存使用？
• RQ4 底层内核与内存优化在多大程度上可降低基于注意力机制模型的训练与推理成本？
• RQ5 所提出的优化方法是否可推广至 AlphaFold 之外的其他大型注意力模型？

主要发现

• FastFold 在 512 张 A100 GPU 上将 AlphaFold 训练时间从 11 天缩短至 67 小时，实现 10.5 倍加速。
• 该系统在 512 张 GPU 上实现 6.02 PetaFLOP/s 的聚合吞吐量，平行效率达 90.1%。
• AutoChunk 在推理过程中使内存使用量减少 80% 以上，其内存效率比专家设计的分块策略高出 30.6%–34.4%。
• 与基线实现相比，FastFold 在长序列推理中实现 7.5–9.5 倍加速。
• 与使用专家调优分块（64）的 OpenFold 相比，AutoChunk 将推理速度提升 12%，同时仅比未分块推理损失 4% 的性能。
• 所提出的优化方法具有可扩展性，可应用于其他基于注意力机制的模型，如 RoseTTAFold、MSA Transformer 和视频变换器。