[論文レビュー] FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours

Shenggan Cheng, Zhao, Xuanlei+8 more|arXiv (Cornell University)|Mar 2, 2022

Protein Structure and DynamicsBiochemistry, Genetics and Molecular Biology被引用数 17

ひとこと要約

FastFold は、動的軸方向並列、二重非同期演算、AutoChunk を導入することで、メモリ使用量を 80% 以上削減し、GPU スケーリング効率を向上させ、AlphaFold の学習を 11 日間から 67 時間に短縮し、長時間シーケンスの推論を 7.5–9.5× 速くします。512 個の A100 GPU で 6.02 PetaFLOP/s のスループットを達成し、並列効率は 90.1% です。

ABSTRACT

Protein structure prediction helps to understand gene translation and protein function, which is of growing interest and importance in structural biology. The AlphaFold model, which used transformer architecture to achieve atomic-level accuracy in protein structure prediction, was a significant breakthrough. However, training and inference of the AlphaFold model are challenging due to its high computation and memory cost. In this work, we present FastFold, an efficient implementation of AlphaFold for both training and inference. We propose Dynamic Axial Parallelism and Duality Async Operations to improve the scaling efficiency of model parallelism. Besides, AutoChunk is proposed to reduce memory cost by over 80% during inference by automatically determining the chunk strategy. Experimental results show that FastFold reduces overall training time from 11 days to 67 hours and achieves 7.5X - 9.5X speedup for long-sequence inference. Furthermore, we scale FastFold to 512 GPUs and achieve an aggregate throughput of 6.02 PetaFLOP/s with 90.1% parallel efficiency.

研究の動機と目的

• 大規模なタンパク質構造予測において、学習と推論の計算コストとメモリコストを低減すること。
• マルチGPU学習におけるモデル並列化のスケーラビリティと通信オーバーヘッドの制限を克服すること。
• 極めて長いタンパク質の処理を可能にするために、長時間シーケンス推論中のメモリ消費量を低減すること。
• 大規模なGPUクラスタ（例：512 個の A100 GPU）で高いスループットと並列効率を達成すること。
• 構造生物学やそれ以上の分野における他のアテンションベースのモデルに適用可能な汎用的最適化フレームワークの開発

提案手法

• 通信オーバーヘッドを低減するように最適化されたテンソル分割を実現する、動的軸方向並列というモデル並列戦略の提案。
• 計算と通信のオーバーラップを可能にする PyTorch 拡張である二重非同期演算の導入。
• 推論中のメモリ削減を目的とした最適なシーケンスチャンク化を自動で決定する AutoChunk の開発。
• Evoformer ブロックに対する低レベルのカーネル最適化を実施し、計算を高速化するとともにメモリフットプリントを削減。
• ハイブリッド並列化と通信最適化されたトレーニングパイプラインを用いて、512 個の A100 GPU でのトレーニングスケーリング。
• 推論中に速度とメモリ使用量のバランスを取るために、動的調整が可能なメモリ認識型チャンク化を採用。

実験結果

リサーチクエスチョン

• RQ1 大規模なタンパク質構造予測において、モデル並列化を最適化することで通信オーバーヘッドを低減し、スケーリング効率を向上させることは可能か？
• RQ2 アテンションベースのモデル（例：AlphaFold）において、計算と通信のオーバーラップを効果的に実現する方法は何か？
• RQ3 自動化されたチャンク戦略は、長時間シーケンス推論中のメモリ使用量削減において、熟練者が調整した設定を上回ることができるか？
• RQ4 低レベルのカーネル最適化とメモリ最適化は、アテンションベースのモデルにおける学習コストと推論コストをどの程度低減できるか？
• RQ5 提案された最適化は、AlphaFold を超える他の大規模アテンションモデルにも一般化可能か？

主な発見

• FastFold は 512 個の A100 GPU で、AlphaFold の学習時間を 11 日間から 67 時間に短縮し、10.5× の高速化を達成。
• 512 個の GPU で 6.02 PetaFLOP/s の集約スループットを達成し、並列効率は 90.1% であった。
• AutoChunk は推論中にメモリ使用量を 80% 以上削減し、熟練者が設計したチャンク戦略を 30.6%–34.4% のメモリ効率向上で上回った。
• FastFold は、ベースライン実装と比較して、長時間シーケンス推論を 7.5–9.5× 速くした。
• AutoChunk は、熟練者が調整したチャンク（64）を用いた OpenFold よりも推論速度を 12% 向上させ、非チャンク化推論と比較してわずか 4% の速度低下にとどめた。
• これらの最適化は一般化可能であり、RoseTTAFold や MSA Transformer、ビデオ変換器など、他のアテンションベースのモデルにも適用可能であった。