機器之心報道

編輯：小舟、蛋醬

AlphaFold 雖好，但耗時且成本高，現在首個用于蛋白質結構預測模型的性能優化方案來了。

蛋白質結構預測一直是結構生物學中的一個重要的研究問題。直接從蛋白質的氨基酸序列預測蛋白質的三維結構在許多領域都有廣泛的應用，包括藥物設計、蛋白質設計等。

預測蛋白質結構的方法主要有兩種：實驗方法和計算方法。實驗方法需要以高昂的時間和經濟成本獲得更準確的蛋白質結構。計算方法則能夠以低成本高吞吐量預測蛋白質結構，因此提高計算方法的預測精度至關重要。借助深度神經網絡，AlphaFold 能夠周期性地以原子精度預測蛋白質結構，但 AlphaFold 模型的訓練和推理由于其特殊的性能特點和巨大的內存消耗，既耗時又昂貴。

基于此，來自潞晨科技和上海交大的研究者提出了一種蛋白質結構預測模型的高效實現 FastFold。 FastFold 包括一系列基于對 AlphaFold 性能全面分析的 GPU 優化。同時，通過動態軸并行和對偶異步算子，FastFold 提高了模型并行擴展的效率，超越了現有的模型并行方法。

論文地址：https://arxiv.org/abs/2203.00854

實驗結果表明，FastFold 將整體訓練時間從 11 天減少到 67 小時，并實現了 7.5 9.5 倍的長序列推理加速。此外，研究者還將 FastFold 擴展到 512 個 A100 GPU 的超算集群上，聚合峰值性能達到了 6.02PetaFLOPs，擴展效率達到 90.1%。

FastFold

不同于一般的 Transformer 模型，AlphaFold 在 GPU 平臺上的計算效率較低，主要面臨兩個挑戰：1) 有限的全局批大小限制了使用數據并行性將訓練擴展到更多節點，更大的批大小會導致準確率更低。即使使用 128 個谷歌 TPUv3 訓練 AlphaFold 也需要約 11 天； 2) 巨大的內存消耗超出了當前 GPU 的處理能力。在推理過程中，較長的序列對 GPU 內存的需求要大得多，對于 AlphaFold 模型，一個長序列的推理時間甚至可以達到幾個小時。

AlphaFold 模型架構

作為首個用于蛋白質結構預測模型訓練和推理的性能優化工作，FastFold 成功引入了大型模型訓練技術，顯著降低了 AlphaFold 模型訓練和推理的時間和經濟成本。

FastFold 由 Evoformer 的高性能實現、AlphaFold 的主干結構和一種稱為動態軸并行（Dynamic Axial Parallelism，DAP）的模型并行新策略組成。

Evoformer

根據計算和存儲訪問的特點，該研究將算子分為三大類:

1) GEMM。這類包括矩陣乘法、批矩陣 - 矩陣相乘和其他密集矩陣計算。來自 NVIDIA Tesla GPU 的張量核可以顯著加速 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）在擴展效率方面優于當前的標準張量并行（Tensor Parallelism），DAP 具有以下幾個優勢：

支持 Evoformer 中的所有計算模塊；所需的通信量比張量并行小得多；顯存消耗比張量并行低；給通信優化提供了更多的空間，如計算通信重疊。

通信優化

在通信方面，該研究提出了由一對通信算子組成的對偶異步算子（Duality Async Operation）。

這種新方法在模型前向傳播的過程中，前一個通信算子觸發異步通信，然后在計算流上進行一些沒有依賴性的計算，然后后一個通信算子阻塞，直到通信完成；在反向傳播的過程中，后一個算子將觸發異步通信，前一個算子阻塞通信。

利用對偶異步算子可以很容易地在 PyTorch 這樣的動態框架上實現前向傳播和反向傳播中的計算和通信遮疊。

評估

研究者首先評估了 Evoformer 內 核的性能改進，然后對端到端訓練和推理性能進行了評估。所有的實驗都在 NVIDIA Tesla A100 平臺上進行。基線是 AlphaFold 的官方實現和另一個 OpenFold 開源 PyTorch 實現。官方實現的 AlphaFold 只有推理部分，而 OpenFold 則是根據原始 AlphaFold 論文復制訓練和推理的。

A. Evoformer 性能

圖 8(a)和圖 9(a)分別展示了 Fused Softmax 和 LayerNorm 的性能比較。

對于 Fused Softmax，研究者比較了 PyTorch 原生內核和 FastFold 優化內核的性能。注意力輸入序列的長度是 x，注意力的隱藏大小是 y。從圖 8(a)可以看出，FastFold 內核的性能可以提高 1.773.32 倍 。

對于 LayerNorm，研究者不僅比較了 PyTorch 原生內核，還比較了 NVIDIA Apex 中高度優化的 LayerNorm 內核。如圖 9(a) ，FastFold 的性能比 PyTorch 和 Apex 分別提高了 5.538.65 倍 和 1.201.62 倍 。由于對有限范圍的特別優化，相比于高度優化的 Apex LayerNorm，FastFold 也實現了良好的性能改進。

B. 端到端訓練性能

在端到端訓練表現的評估中，研究者使用官方 AlphaFold 文件中的訓練參數進行了盡可能多的測試。這樣可以更好地比較不同的方法或實現在實際的訓練場景中的工作方式。所有的訓練實驗都是在 128 節點的 GPU 超級計算機上進行的。在超級計算機中，每個節點包括 4 臺 NVIDIA Tesla A100，并且有 NVIDIA NVLink 用于 GPU 互連。

由于張量并行更多地依賴于設備之間的高速互連來進行通信，在訓練期間，模型并行通常用于節點以及訓練期間節點之間的數據并行。研究者分別在模型并行和數據并行兩個級別測試了模型的訓練性能，結果如圖 10 和圖 11 所示。在模型并行性方面，論文比較了張量并行和動態軸并行兩種并行方法在初始訓練和微調兩種訓練設置下的 scalability。

如圖 10 所示，在初始訓練和微調方面，DAP 的擴展性能明顯優于 TP。

擴展結果如圖 11 所示。我們可以看到，在接下來的幾個平行里，基本上是線性規模的。微調訓練的擴展效率達到 90.1% 。

根據訓練效果的評估結果，可以推算出 AlphaFold 的總體時間和經濟成本。表 IV 列出并比較了 AlphaFold、 OpenFold 和 FastFold 三種實現的時間和經濟成本。由于沒有公開的訓練代碼，所以 AlphaFold 的數據來源于官方文件。

考慮到時間和經濟成本，研究者選用了 256 個 A100 進行初始訓練，然后在微調階段擴展到 512 個 A100。

基于這種設置，FastFold 可以將訓練時間減少到 2.81 天。與需要 11 天訓練的 AlphaFold 相比，訓練的時間成本減少了 3.91 倍。與 OpenFold 相比，訓練的時間成本降低了 2.98 倍，經濟成本降低了 20% 。

在微調階段，FastFold 在 512 × A100 的設置下實現了 6.02 PetaFLOPs 的計算速度。由于時間和經濟成本的顯著降低，FastFold 使得蛋白質結構預測模型的訓練速度更快、成本更低，這將推動相關模型的研究和開發效率，并促進基于 Evoformer 的蛋白質結構預測模型的開發。

C. 端到端推理性能

針對短序列、長序列和超長序列，研究者對比了 FastFold、OpenFold、AlphaFold 的推理性能。所有推理實驗均在由 8 個 NVIDIA A100（帶有 NVLink）組成的 GPU 服務器上完成。

對于短序列，典型的氨基酸序列長度不超過 1k，單個模型推理需要幾秒鐘到大約一分鐘。在這個序列范圍內，視頻內存消耗相對較小，分布式推理的效率較低。

研究者在 1 個 GPU 上比較了三種實現的推理延遲，結果如圖 12 所示：

對于長度為 1k 到 2.5 k 的氨基酸序列的長序列推斷，直接推理會遇到內存容量問題，推理時間達到幾分鐘甚至幾十分鐘。對于 FastFold，采用分布式推理方法可以減少內存容量的需求，顯著縮短推理時間。如圖 13 所示，當使用分布式推理時，FastFold 比 OpenFold 減少推理時間 7.59.5 倍 ，比 AlphaFold 減少推理時間 9.311.6 倍 。

如表 3 所示，對于 AlphaFold 和 OpenFold，當序列長度達到 3k 之上時，都會遇到內存不足 (OOM) 問題。然而，對于 FastFold，由于分布式推理方法，它可以利用 GPU 更多的計算和內存來完成極長的序列推理。此外，對于長度高達 4k 的序列，FastFold 的推理延遲 在 10 分鐘之內。

作者介紹

本文的作者之一尤洋現為新加坡國立大學計算機系任助理教授。2020 年，尤洋在加州大學伯克利分校計算機系獲得博士學位。

尤洋的主要研究方向是高性能計算與機器學習的交叉領域，當前研究重點為大規模深度學習訓練算法的分布式優化。他曾其以一作作者的身份發表研究論文《Large Batch Optimization for Deep Learning ：Training BERT in 76 Minutes》，提出了一種 LAMB 優化器（Layer-wise Adaptive Moments optimizer for Batch training），將超大模型 BERT 的預訓練時間由 3 天縮短到了 76 分鐘，刷新世界記錄。到目前為止，LAMB 仍為機器學習領域的主流優化器，成果被 Google、Facebook、騰訊等科技巨頭在實際中使用。

2021 年 7 月，尤洋在北京中關村創辦了高性能計算公司「潞晨科技」。不久后，潞晨科技即宣布完成超千萬元種子輪融資，該公司創業目標是「最大化人工智能開發速度的同時，最小化人工智能模型部署成本」。

參考鏈接：

https://weibo.com/ttarticle/p/show?id=2309404742878697226261

標簽： 訓練時間