机器之心报道
蛋白质结构预测一直是结构生物学中的一个重要的研究问题。直接从蛋白质的氨基酸序列预测蛋白质的三维结构在许多领域都有广泛的应用，包括药物设计、蛋白质设计等。
预测蛋白质结构的方法主要有两种：实验方法和计算方法。实验方法需要以高昂的时间和经济成本获得更准确的蛋白质结构。计算方法则能够以低成本高吞吐量预测蛋白质结构，因此提高计算方法的预测精度至关重要。借助深度神经网络，AlphaFold能够周期性地以原子精度预测蛋白质结构，但AlphaFold模型的训练和推理由于其特殊的性能特点和巨大的内存消耗，既耗时又昂贵。
基于此，来自潞晨科技和上海交大的研究者提出了一种蛋白质结构预测模型的高效实现FastFold。FastFold包括一系列基于对AlphaFold性能全面分析的GPU优化。同时，通过动态轴并行和对偶异步算子，FastFold提高了模型并行扩展的效率，超越了现有的模型并行方法。
论文地址：
实验结果表明，FastFold将整体训练时间从11天减少到67小时，并实现了7.5∼9.5倍的长序列推理加速。此外，研究者还将FastFold扩展到512个A100GPU的超算集群上，聚合峰值性能达到了6.02PetaFLOPs，扩展效率达到90.1%。
FastFold
不同于一般的Transformer模型，AlphaFold在GPU平台上的计算效率较低，主要面临两个挑战：1)有限的全局批大小限制了使用数据并行性将训练扩展到更多节点，更大的批大小会导致准确率更低。即使使用128个谷歌TPUv3训练AlphaFold也需要约11天；2)巨大的内存消耗超出了当前GPU的处理能力。在推理过程中，较长的序列对GPU内存的需求要大得多，对于AlphaFold模型，一个长序列的推理时间甚至可以达到几个小时。
AlphaFold模型架构
作为首个用于蛋白质结构预测模型训练和推理的性能优化工作，FastFold成功引入了大型模型训练技术，显著降低了AlphaFold模型训练和推理的时间和经济成本。
FastFold由Evoformer的高性能实现、AlphaFold的主干结构和一种称为动态轴并行（DynamicAxialParallelism，DAP）的模型并行新策略组成。
Evoformer
根据计算和存储访问的特点，该研究将算子分为三大类:
1)GEMM。这类包括矩阵乘法、批矩阵-矩阵相乘和其他密集矩阵计算。来自NVIDIATeslaGPU的张量核可以显著加速GEMM算子;
2)批递减算子。这类包括LayerNorm、Softmax等，计算强度低于GEMM算子，并且更容易出现访问瓶颈。
3)Element-wise算子。这类包括element-wise的相加、相乘、dropout和激活，是计算密集度最低的一类。
Evoformer的注意力机制如下图所示
值得注意的是，Evoformer和原版Transformer之间有几个关键区别：
该研究分析了Evoformer的复杂结构并进行了内核融合，并针对Evoformer中的特定操作和基于性能特征的特定内核进行了优化，例如Softmax和LayerNorm。高性能的Evoformer实现大大降低了训练和推理的经济成本。
动态轴并行
在并行策略方面，该研究根据AlphaFold的计算特征创新性地提出了动态轴并行策略，在AlphaFold的特征的序列方向上进行数据划分，并使用All_to_All进行通信。动态轴并行（DAP）在扩展效率方面优于当前的标准张量并行（TensorParallelism），DAP具有以下几个优势：
支持Evoformer中的所有计算模块；
所需的通信量比张量并行小得多；
显存消耗比张量并行低；
给通信优化提供了更多的空间，如计算通信重叠。
通信优化
在通信方面，该研究提出了由一对通信算子组成的对偶异步算子（DualityAsyncOperation）。
这种新方法在模型前向传播的过程中，前一个通信算子触发异步通信，然后在计算流上进行一些没有依赖性的计算，然后后一个通信算子阻塞，直到通信完成；在反向传播的过程中，后一个算子将触发异步通信，前一个算子阻塞通信。
利用对偶异步算子可以很容易地在PyTorch这样的动态框架上实现前向传播和反向传播中的计算和通信遮叠。
评估
研究者首先评估了Evoformer内核的性能改进，然后对端到端训练和推理性能进行了评估。所有的实验都在NVIDIATeslaA100平台上进行。基线是AlphaFold的官方实现和另一个OpenFold开源PyTorch实现。官方实现的AlphaFold只有推理部分，而OpenFold则是根据原始AlphaFold论文复制训练和推理的。
性能
图8(a)和图9(a)分别展示了FusedSoftmax和LayerNorm的性能比较。
对于FusedSoftmax，研究者比较了PyTorch原生内核和FastFold优化内核的性能。注意力输入序列的长度是x，注意力的隐藏大小是y。从图8(a)可以看出，FastFold内核的性能可以提高1.77∼3.32倍。
对于LayerNorm，研究者不仅比较了PyTorch原生内核，还比较了NVIDIAApex中高度优化的LayerNorm内核。如图9(a)，FastFold的性能比PyTorch和Apex分别提高了5.53∼8.65倍和1.20∼1.62倍。由于对有限范围的特别优化，相比于高度优化的ApexLayerNorm，FastFold也实现了良好的性能改进。
B.端到端训练性能
在端到端训练表现的评估中，研究者使用官方AlphaFold文件中的训练参数进行了尽可能多的测试。这样可以更好地比较不同的方法或实现在实际的训练场景中的工作方式。所有的训练实验都是在128节点的GPU超级计算机上进行的。在超级计算机中，每个节点包括4台NVIDIATeslaA100，并且有NVIDIANVLink用于GPU互连。
由于张量并行更多地依赖于设备之间的高速互连来进行通信，在训练期间，模型并行通常用于节点以及训练期间节点之间的数据并行。研究者分别在模型并行和数据并行两个级别测试了模型的训练性能，结果如图10和图11所示。在模型并行性方面，论文比较了张量并行和动态轴并行两种并行方法在初始训练和微调两种训练设置下的scalability。
如图10所示，在初始训练和微调方面，DAP的扩展性能明显优于TP。
扩展结果如图11所示。我们可以看到，在接下来的几个平行里，基本上是线性规模的。微调训练的扩展效率达到90.1%。
考虑到时间和经济成本，研究者选用了256个A100进行初始训练，然后在微调阶段扩展到512个A100。
基于这种设置，FastFold可以将训练时间减少到2.81天。与需要11天训练的AlphaFold相比，训练的时间成本减少了3.91倍。与OpenFold相比，训练的时间成本降低了2.98倍，经济成本降低了20%。
在微调阶段，FastFold在512×A100的设置下实现了6.02PetaFLOPs的计算速度。由于时间和经济成本的显著降低，FastFold使得蛋白质结构预测模型的训练速度更快、成本更低，这将推动相关模型的研究和开发效率，并促进基于Evoformer的蛋白质结构预测模型的开发。
C.端到端推理性能
针对短序列、长序列和超长序列，研究者对比了FastFold、OpenFold、AlphaFold的推理性能。所有推理实验均在由8个NVIDIAA100（带有NVLink）组成的GPU服务器上完成。
研究者在1个GPU上比较了三种实现的推理延迟，结果如图12所示：
对于长度为1k到2.5k的氨基酸序列的长序列推断，直接推理会遇到内存容量问题，推理时间达到几分钟甚至几十分钟。对于FastFold，采用分布式推理方法可以减少内存容量的需求，显著缩短推理时间。如图13所示，当使用分布式推理时，FastFold比OpenFold减少推理时间7.5∼9.5倍，比AlphaFold减少推理时间9.3∼11.6倍。
如表3所示，对于AlphaFold和OpenFold，当序列长度达到3k之上时，都会遇到内存不足(OOM)问题。然而，对于FastFold，由于分布式推理方法，它可以利用GPU更多的计算和内存来完成极长的序列推理。此外，对于长度高达4k的序列，FastFold的推理延迟在10分钟之内。
本文的作者之一尤洋现为新加坡国立大学计算机系任助理教授。2020年，尤洋在加州大学伯克利分校计算机系获得博士学位。
尤洋的主要研究方向是高性能计算与机器学习的交叉领域，当前研究重点为大规模深度学习训练算法的分布式优化。他曾其以一作作者的身份发表研究论文《LargeBatchOptimizationforDeepLearning：TrainingBERTin76Minutes》，提出了一种LAMB优化器（Layer-wiseAdaptiveMomentsoptimizerforBatchtraining），将超大模型BERT的预训练时间由3天缩短到了76分钟，刷新世界记录。到目前为止，LAMB仍为机器学习领域的主流优化器，成果被Google、Facebook、腾讯等科技巨头在实际中使用。
2021年7月，尤洋在北京中关村创办了高性能计算公司「潞晨科技」。不久后，潞晨科技即宣布完成超千万元种子轮融资，该公司创业目标是「最大化人工智能开发速度的同时，最小化人工智能模型部署成本」。
