ScaleFold: Reducing AlphaFold Initial Training Time to 10 Hours

AlphaFold2被誉为蛋白质折叠领域的突破，能够以实验室级别的准确度快速预测蛋白质结构。然而，其实现并不包括必要的训练代码。OpenFold是AlphaFold的第一个可训练的公共重新实现。AlphaFold的训练过程耗时极长，并且从更多计算资源的扩展中获得的收益逐渐减少。本研究基于OpenFold对AlphaFold的训练过程进行了全面分析，发现低效的通信和以开销为主导的计算是阻止AlphaFold训练有效扩展的关键因素。我们引入了ScaleFold，一种系统化的训练方法，专门针对这些因素进行了优化。ScaleFold成功将AlphaFold训练扩展到2080个NVIDIA H100 GPU，并实现了高资源利用率。在MLPerf HPC v3.0基准测试中，ScaleFold以7.51分钟完成了OpenFold基准测试，比基线快了6倍。对于从头开始训练AlphaFold模型，ScaleFold在10小时内完成了预训练，比原始AlphaFold预训练基线所需的七天时间显著缩短。

本论文旨在解决AlphaFold训练过程中通信效率低下和计算开销过大的问题，提出了一个针对这些因素进行优化的系统化训练方法ScaleFold，以加速AlphaFold的训练过程。

ScaleFold是一种系统化的训练方法，通过针对AlphaFold训练过程中的通信效率和计算开销进行优化，实现了对AlphaFold模型的快速训练。

论文通过对AlphaFold训练过程的分析，提出了一个高效的训练方法ScaleFold，并在实验中展示了它的优越性能。ScaleFold在MLPerf HPC v3.0基准测试中，相对于基准线实现了超过6倍的加速，同时也在训练AlphaFold模型时取得了显著的训练速度提升。

在这个领域中，最近的相关研究包括AlphaFold的原始论文以及OpenFold的实现。