利用并行 Evoformer 与分支并行的高效 AlphaFold2 训练
分布式、并行与集群计算
2022-11-02 v1
摘要
AlphaFold2 作为前沿的端到端结构预测系统，其准确率已接近实验测定技术。由于其复杂的模型架构与较大的内存消耗，从零训练 AlphaFold2 需要大量的计算资源与时间。高效的 AlphaFold2 训练可加速生命科学的发展。本文中，我们提出一种并行 Evoformer 与分支并行以加速 AlphaFold2 的训练。我们在 PyTorch 实现的 UniFold 与 PaddlePaddle 实现的 HelixFold 上进行了充分实验，分支并行分别将训练性能提升了 38.67% 与 36.93%。我们还证明，在 CASP14 与 CAMEO 数据集上，并行 Evoformer 的准确率可与 AlphaFold2 相当。源代码可在 https://github.com/PaddlePaddle/PaddleFleetX 获取。
引用
@article{arxiv.2211.00235,
title = {Efficient AlphaFold2 Training using Parallel Evoformer and Branch Parallelism},
author = {Guoxia Wang and Zhihua Wu and Xiaomin Fang and Yingfei Xiang and Yiqun Liu and Dianhai Yu and Yanjun Ma},
journal= {arXiv preprint arXiv:2211.00235},
year = {2022}
}
备注
arXiv admin note: substantial text overlap with arXiv:2207.05477