助力藥物研發,低成本加速AlphaFold訓練從11天到67小時,11倍推理加速——開源解決方案FastFold

語言: CN / TW / HK

AlphaFold被Science和Nature評選為2021年十大科學突破之首。潞晨科技與華深智藥聯合開源的AlphaFold訓練推理加速方案FastFold,將GPU優化和大模型訓練技術引入AlphaFold的訓練和推理,成功將AlphaFold總體訓練時間從11天減少到67小時,且總成本更低,在長序列推理中也實現9.3 ∼ 11.6倍提升

FastFold開源地址:https://github.com/hpcaitech/FastFold

蛋白質與AlphaFold

蛋白質是生命的物質基礎,幾乎支持着生命的所有功能。它們是由氨基酸鏈組成的大型複雜分子,而蛋白質的作用主要取決於其獨特的三維結構。弄清楚蛋白質摺疊成什麼形狀被稱為 "蛋白質摺疊問題",在過去的50年裏一直是生物學的一個巨大挑戰。

AlphaFold預測與實驗測量的蛋白質結構對比

實驗和計算方法都可用於預測蛋白質結構。實驗方法可以獲得更準確的蛋白質結構,但需要昂貴的時間和經濟成本。計算方法可以用低成本進行高通量的蛋白質結構預測,所以準確的計算方法一直是學術界和工業界努力的方向。2020年,Google DeepMind 推出了最新一代AlphaFold,成功將Transformer模型引入蛋白質結構預測,並取得了巨大的精度提升。AlphaFold使用端到端的模型架構,可以獲得原子級別精度的結構預測結果。

AlphaFold架構

然而,基於Transformer的AlphaFold模型也帶來了諸多計算方面的挑戰:

1)有限的全局批處理量,使得訓練無法利用數據並行性擴展到更多的節點。據DeepMind官方披露,即便使用128個TPU去訓練AlphaFold,也需要11天以上才能完成;

2)在訓練過程中巨大的內存消耗,遠超當前GPU硬件所能提供的顯存容量。

3)在推理方面,長序列推理對GPU顯存的需求更大。而且一個長序列的推理時間對於AlphaFold模型來説甚至長達到幾個小時。

FastFold

為了解決上述難題,我們提出了FastFold,這是一個用於蛋白質結構預測的模型訓練推理加速方案。FastFold是世界上首個系統化的針對蛋白質結構預測模型進行訓練和推理的性能優化工作。FastFold成功地引入了大模型訓練技術,極大降低了AlphaFold模型訓練和推理的時間和經濟成本。

FastFold包括了一系列基於針對AlphaFold性能特徵的GPU優化。同時,通過動態軸並行對偶異步算子,FastFold實現了很高的模型並行化擴展效率,超越了目前主流的模型並行方法。

動態軸並行。FastFold首次嘗試將模型並行技術引入AlphaFold中,並根據AlphaFold的計算特徵創新性地提出了動態軸並行技術。不同於傳統的張量並行,動態軸並行選擇在AlphaFold的特徵的序列方向上進行數據劃分,並使用All_to_All進行通信。動態軸並行對比張量並行有幾個優勢

1)動態軸並行支持Evoformer中的所有計算模塊;

2)所需的通信量比張量並行小得多;

3)動態軸並行顯存消耗比張量並行低;

4)DAP給通信優化提供了更多的空間,如計算通信重疊。

對偶異步算子。對偶異步算子由一對通信算子組成。在模型的前向傳播過程中,前一個通信算子觸發異步通信,然後在計算流上進行一些沒有依賴性的計算,然後後一個通信算子阻塞,直到通信完成。當模型反向傳播時,後一個算子將觸發異步通信,前一個算子阻塞通信。利用對偶異步算子可以讓我們在PyTorch這樣的動態圖框架上很方便的實現前向傳播和反向傳播中的計算和通信遮疊。

卓越性能

FastFold對比目前所有的AlphaFold的實現版本都有顯著的性能優勢:

在訓練方面,FastFold可以將訓練時間減少到2.81天。與AlphaFold需要11天的訓練相比,訓練提升3.91倍。與OpenFold(來自劍橋大學的AlphaFold復現版本)相比,訓練提升2.98倍,經濟成本降低20%。我們將FastFold擴展到擁有512個A100 GPU的超算集羣上,聚合峯值性能達到了6.02PetaFLOPs,擴展效率達到90.1%

在推理方面,FastFold在短序列,長序列,超長序列均有明顯的性能優勢:

1)在不超過1K的短序列推理的情況下,FastFold與AlphaFold和OpenFold相比,單GPU推理性能分別提升2.01 ∼ 4.05倍和1.25 ∼ 2.11倍。

2)在長度為1K ~ 3K的長序列推理中,FastFold可以使用分佈式推理大幅度降低推理時間,對比OpenFold提升7.5 ∼ 9.5倍,對比AlphaFold相提升9.3 ∼ 11.6倍

3)對於長度超過3K的超長序列推理,OpenFold和AlphaFold都因為爆顯存而無法推理。FastFold因為支持分佈式推理,可以利用多個GPU的計算和顯存資源完成超長序列推理的任務。

解決蛋白質結構預測模型訓練和推理的在計算上的挑戰,對其在結構生物學的廣泛應用有重要意義。FastFold大幅度降低了AlphaFold模型訓練和推理的時間成本和經濟成本,將極大促進新一代藥物研發,蛋白質設計,抗體設計等應用場景的革命創新和發展。

背後功臣

FastFold項目的成功,得到了開源AI基礎設施Colossal-AI的大力支持。大規模並行AI系統Colossal-AI ,通過高效多維並行、大規模優化庫、自適應任務調度、消除宂餘內存等方式,旨在打造一個高效的分佈式AI系統,作為深度學習框架的內核,幫助用户便捷實現最大化提升AI部署效率,同時最小化部署成本。

Colossal-AI開源地址:https://github.com/hpcaitech/ColossalAI

傳送門

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

FastFold項目地址:https://github.com/hpcaitech/FastFold

Colossal-AI項目地址:https://github.com/hpcaitech/ColossalAI

 

參考鏈接

https://arxiv.org/abs/2203.00854

https://www.sohu.com/a/478835644_354973