從氨基酸序列預測 3D 蛋白質結構一直是生物信息學領域由來已久的重要問題。近年來，基于深度學習的計算方法不斷涌現，并已顯示出有希望的結果。
在這些工作中，AlphaFold2 是第一種與較慢的基于物理性質的計算方法相美的方法。它被 Nature 命名為 2021 年度最佳方法。該模型基于其他當代深度學習模型廣泛采用的序列注意力機制的變體而構建。
通過使用 MML，我們能夠有效地訓練和評估大規模的生物醫學數據集。MML 是一個高度可擴展和可定制化的平臺，它支持多種機器學習算法和模型，并提供了豐富的工具和接口，以便于用戶能夠輕松地集成、擴展和定制化 MML 平臺以滿足其特定的需求和應用場景。
遺憾的是，AlphaFold2 訓練所需的時間仍然是一個重大瓶頸。
AlphaFold2 訓練使用 1000 萬個樣本和 128 個 TPU 進行，收時間超過 11 天(初始訓練 7 天，微調 4 天)。如此長的訓練時間會降低研究的迭代速度。
相比之下，使用 128 個 NVIDIA A100 GPU 的 OpenFold 訓練過程需要 8 天以上，比 AlphaFold2 快 27%，但仍有很大的改進空間。
選擇 OpenFold 的初始訓練階段是為了 MLCommons MLPerf HPC v3.0?在2023 年初推出的 AI 基準測試套件。在此基準測試套件中選擇 OpenFold 凸顯了 OpenFold 作為蛋白質結構預測的強大工具的潛力，并強調了其對 高性能計算(HPC) 社區。
識別 OpenFold 訓練的瓶頸
通常，大型模型使用多個節點和每個節點多個 GPU 進行數據并行 (DP) 訓練。這種技術可以加快訓練速度，同時在實現最佳 GPU 利用率和可擴展性方面也帶來了挑戰。
分布式并行訓練的可擴展性受到全局批量大小的限制，因為批量更大會導致準確性降低。為應對這一挑戰，FastFold 引入了動態軸向并行 (DAP)，作為通過模型并行提高可擴展性的解決方案。DAP 將單個訓練樣本的中間激活和相關計算沿非歸約軸分割。
最右上角的“Optimized OF+DAP”數據源自 ID 3.0-8009，方法是將訓練時間除以 OpenFold (OF) 的 8 GPU 訓練時間。MLPerf 名稱和徽標均為 MLCommons Association 在美國和其他國家 地區的商標。保留所有權利。嚴禁未經授權使用。有關更多信息，請參閱 MLCommons。
在 NVIDIA 提交的 MLPerf HPC 中，DAP 被重新實施到 OpenFold 中。 NVIDIA MLPerf HPC 基準測試團隊還對 OpenFold 訓練過程進行了全面分析。分析顯示了阻礙訓練擴展到更多計算資源的主要因素：
- 通信用度：分布式訓練期間的通信密集，但效率低下，這在很大程度上是由于通信用度造成的。DAP 引入了額外的 all-gather 和 all-to-all 開銷。另一方面，在訓練過程中，一些通信涉及少量數據，無法充分利用通信帶寬。
- 不平衡：在分布式訓練中，掉隊是一個常見的問題。在完成分配的任務方面落后的工作者會降低整體訓練處理速度。OpenFold 的批量準備時間取決于原始殘差序列長度，這會導致不平衡。占用訓練進程 CPU 核心的后臺進程還會導致計算負載，從而導致更多的不平衡。
- 不可并行的工作負載：Amdahl 定律有助于根據可并行運行的系統的比例量化潛在的并行加速收益。DAP 無法并行處理數據加載階段和結構模組階段，因此減少了可并行計算負載，進一步限制了可擴展性。
- 內核可擴展性不佳：DAP 可將內核工作負載減少高達 8 倍。小型工作負載主要由啟動和完成任務組成，因此 GPU 帶寬不飽和，這使得內核可擴展性更加糟糕。
- CPU 使用度：OpenFold 中有超過 15 萬次運算。許多小型連續內核會導致 PyTorch 運算啟動產生重大開銷。
在圖 1 中，灰色數據僅顯示 OpenFold 和 DAP 增強 OpenFold (灰色虛線)。提供額外速度的 GPU 最大數量為 512，速度約為配備 8 個 GPU 的一個節點的 40 倍。
使用 NVIDIA GPU 優化 OpenFold 訓練
為應對這些挑戰， NVIDIA MLPerf HPC 基準測試團隊進行了一系列系統優化：
- 引入了新的非阻塞數據管道，以解決慢速工作問題
- 識別關鍵計算模式，例如多頭注意力、層數范數、隨機權重平均 (SWA) 和 OpenFold 訓練中的梯度裁剪。然后，該團隊為每種模式設計了專用的內核，在整個 OpenFold 模型中融合了碎片計算，并針對各種工作負載大小和目標硬件架構仔細調整了內核配置。
- 添加并應用了完整的 bfloat16 支持 CUDA 圖形 因為它無需與 CPU 進行交互，從而更大限度地減少了啟動操作的開銷。
OpenFold 的訓練過程
批量大小設置為 128，以便在搭載 DAP8 的 1056 NVIDIA H100 GPU 上訓練前 5K 步長，其中 1024 個 GPU 用于訓練，32 個 GPU 用于評估。在每種情況下，在執行前 5K 訓練步驟之前，訓練指標必須超過 0.8 lDDT-C。
全局批量大小已切換到 256 ( NVIDIA Triton mha 內核已禁用)，以便在搭載 DAP4 的 1056 NVIDIA H100 GPU 上訓練其余步驟(1024 GPU 用于訓練，32 GPU 用于評估)。初始訓練過程總共需要 12 小時。
或者，您可以使用搭載 DAP8 的 2080 NVIDIA H100 GPU (2048 GPU 用于訓練，32 GPU 用于評估)來訓練其余步驟，初始訓練時間可縮短至 10 小時。
使用 OpenFold 加速蛋白質結構
通過將其與一系列細粒度優化相結合，整體通信效率得到大幅提升。可擴展性問題得到成功解決，并將 OpenFold 訓練擴展到 2080 NVIDIA H100 GPU，而之前的工作僅擴展到 512 NVIDIA A100 GPU.在 MLPerf HPC v3.0 基準測試中，NVIDIA 優化的 OpenFold 在 7.51 分鐘內完成了部分訓練任務 比基準測試快 6 倍以上。
在圖 1 中，綠色數據顯示了相同的訓練，使用 DAP 和這些優化實現了 4 倍的可擴展性提升到 2048 GPU，速度提高了 4.5 倍，比運行通用 OpenFold 的一臺 8 GPU 服務器的速度快 180 倍。
經過訓練優化的 OpenFold 需要 5 萬到 6 萬步才能達到 0.9 lDDT-C，使用 1056 NVIDIA H100 GPU 需要 12.4 小時，使用 2080 NVIDIA H100 GPU 需要 10 小時。與原始 AlphaFold2 所需的 7 天相比，這是一項重大改進，與以前的工作相比，創下了新紀錄。
適用于生命科學的生成式 AI 工具包
我們希望這項工作能夠通過提供有效的機制來擴展基于深度學習的計算方法來解決蛋白質折疊問題，從而使 HPC 和生物信息學研究社區受益。
我們還希望這項工作中使用的工作負載分析和優化方法能夠說明機器學習系統的設計和實現。這項工作中的優化可以應用到其他類似 AlphaFold2 的模型。
OpenFold 模型的版本和 NVIDIA 用于訓練該模型的軟件將很快在 NVIDIA BioNeMo 以及相關的 framework 用于藥物研發的生成式 AI 平臺。
有關更多信息，請參閱以下資源：
?
