[论文解读] HelixFold: An Efficient Implementation of AlphaFold2 using PaddlePaddle

Guoxia Wang, Xiaomin Fang+7 more|arXiv (Cornell University)|Jul 12, 2022

Machine Learning in Bioinformatics被引用 30

一句话总结

HelixFold 在 PaddlePaddle 中端到端实现 AlphaFold2，通过分支并行和内存优化实现更快的训练，同时在 CASP14 和 CAMEO 上保持与 AlphaFold2 相近的准确性。

ABSTRACT

Accurate protein structure prediction can significantly accelerate the development of life science. The accuracy of AlphaFold2, a frontier end-to-end structure prediction system, is already close to that of the experimental determination techniques. Due to the complex model architecture and large memory consumption, it requires lots of computational resources and time to implement the training and inference of AlphaFold2 from scratch. The cost of running the original AlphaFold2 is expensive for most individuals and institutions. Therefore, reducing this cost could accelerate the development of life science. We implement AlphaFold2 using PaddlePaddle, namely HelixFold, to improve training and inference speed and reduce memory consumption. The performance is improved by operator fusion, tensor fusion, and hybrid parallelism computation, while the memory is optimized through Recompute, BFloat16, and memory read/write in-place. Compared with the original AlphaFold2 (implemented with Jax) and OpenFold (implemented with PyTorch), HelixFold needs only 7.5 days to complete the full end-to-end training and only 5.3 days when using hybrid parallelism, while both AlphaFold2 and OpenFold take about 11 days. HelixFold saves 1x training time. We verified that HelixFold's accuracy could be on par with AlphaFold2 on the CASP14 and CAMEO datasets. HelixFold's code is available on GitHub for free download: https://github.com/PaddlePaddle/PaddleHelix/tree/dev/apps/protein_folding/helixfold, and we also provide stable web services on https://paddlehelix.baidu.com/app/drug/protein/forecast.

研究动机与目标

• 旨在在不牺牲准确性的前提下，降低 AlphaFold2 的训练时间和内存使用量
• 在 PaddlePaddle 中为 AlphaFold2 开发完整的训练与推理流水线（HelixFold）
• 引入新颖的并行性和内存优化技术，以实现高效的端到端训练

提出的方法

• 融合多个算子和张量以减小 CPU 调度开销（Tensor Fusion 和 Fused Gated Self-Attention）
• 提出 Branch Parallelism（BP）以跨设备对 Evoformer 的 MSA 与 Pair 堆栈进行并行化
• 采用混合并行（BP-DAP-DP）以通过数据并行和内存效率提升吞吐量
• 应用内存优化：重计算、BFloat16 激活、就地读写内存，以及 Subbatch/DAP 技术
• 融合参数、梯度和优化器状态以减少内核启动次数和内存碎片
• 利用 Subbatch 和 DAP 管控极长序列的峰值内存

实验结果

研究问题

• RQ1 在降低训练时间的同时，HelixFold 能否在 CASP14 和 CAMEO 上达到 AlphaFold2 的准确性？
• RQ2 融合、分支并行和内存优化技术如何影响训练吞吐量和内存占用？
• RQ3 与现有 Evoformer 的并行策略相比，BP-DAP-DP 混合并行的相对优势是什么？
• RQ4 使用基于 PaddlePaddle 的 HelixFold，与原始 AlphaFold2 和 OpenFold 相比，端到端训练时间和资源成本的降低是多少？

主要发现

• HelixFold 在 CASP14（TM-score 0.8771 对 0.8772）和 CAMEO（TM-score 0.8885 对 0.8862）上实现了与 AlphaFold2 相当的准确性
• 端到端训练时间减少到 7.5 天（混合并行为 5.3 天），相比 AlphaFold2 与 OpenFold 的约 11 天
• 初始训练吞吐量提升：HelixFold 相对于 AlphaFold2 的吞吐量提升为 +28.42% 到 +52.13%，也优于 OpenFold；在微调阶段，HelixFold 比 AlphaFold2 高出约 84.49%
• BP-DAP-DP 混合并行比仅 DAP 的效率更高，通信开销更低
• 内存优化（Recompute、BFloat16、就地读写）和张量融合显著降低峰值内存和内核启动次数
• 融合的 Gated Self-Attention 在保持准确性的同时降低 CPU/GPU 开销