[論文レビュー] Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size

Albert Musaelian, Anders Johansson+2 more|arXiv (Cornell University)|Apr 20, 2023

Machine Learning in Materials Science被引用数 14

ひとこと要約

Allegro は、スケーラブルな等変ニューラルポテンシャルと巨大GPU並列性を組み合わせることにより、スケーラブルで量子精度の高い生体分子シミュレーションを実現し、ナノ秒のタンパク質動態と Perlmutter 上での 44-million-atom HIV capsid シミュレーションを強力および弱いスケーリングで可能にする。

ABSTRACT

This work brings the leading accuracy, sample efficiency, and robustness of deep equivariant neural networks to the extreme computational scale. This is achieved through a combination of innovative model architecture, massive parallelization, and models and implementations optimized for efficient GPU utilization. The resulting Allegro architecture bridges the accuracy-speed tradeoff of atomistic simulations and enables description of dynamics in structures of unprecedented complexity at quantum fidelity. To illustrate the scalability of Allegro, we perform nanoseconds-long stable simulations of protein dynamics and scale up to a 44-million atom structure of a complete, all-atom, explicitly solvated HIV capsid on the Perlmutter supercomputer. We demonstrate excellent strong scaling up to 100 million atoms and 70% weak scaling to 5120 A100 GPUs.

研究の動機と目的

• 大規模な生体分子システムのための、スケーラブルで正確な等変ディープラーニングポテンシャルを実証する。
• 最先端の等変モデルの精度とエクサスケールの計算能力を橋渡しする。
• 前例のないシステムサイズでの安定した長時間スケールの生体分子シミュレーション（タンパク質動態）を実証する。

提案手法

• 系エネルギーを原子成分と順序付きペア成分に分解する、スケーラブルな等変MLIPである Allegro を導入する。
• スカラー軌道と等変テンソル軌道の二重トラック構成を用い、計算コストを制御しつつ容量を高める。
• ストライドメモリ配置と最適化されたテンソル積を活用して、GPU 実行と単一パスの縮約を効率化する。
• 混合精度計算（TF32 を FP32 と時々 FP64）を適用し、速度と数値安定性のバランスを取る。
• 近傍計算を正確さを損なわずに削減するため、種別ごとの対ペアカットオフを組み込む。
• LAMMPS への GPU レジデントパイプラインを介して Allegro を統合し、確立済みの MD 空間分解を活用する。

実験結果

リサーチクエスチョン

• RQ1 Allegro は最先端の等変精度を達成し、数億原子規模へスケールさせることができるのか？
• RQ2 最新世代の GPU 上での大規模生体分子システムに対する Allegro の性能（timesteps/s, ns/day）はどれほどか？
• RQ3 現実的な生体分子構造に対して、長時間のシミュレーションにおける安定性と精度はどうなるのか？
• RQ4 MD のための等変学習を、精度を犠牲にすることなくスケーラブルにするための建築的革新は何か？
• RQ5 Allegro は生体分子系の精度とスケーラビリティの点で、既存の MLIP および古典的力場とどう比較されるのか？

主な発見

• Allegro はベンチマークで高い精度を達成し、perlmutter-規模のハードウェアで 126.4 百万原子にスケールし、100 timesteps/s の性能を発揮する。
• 100 百万原子への強いスケーリングと、5120 台の A100 GPU（1280 ノード）での 70% 弱いスケーリングを実現。
• 最先端の ML 相互原子ポテンシャルの大規模生体分子シミュレーションへの初の適用。
• タンパク質動態のナノ秒長の安定シミュレーション（例：DHFR、Factor IX）および 44M-atom HIV_capsid 溶媒化系を生成。
• スカラー軌道とテンソル軌道を、ストライドメモリ配置と最適化されたテンソル積と組み合わせることで、精度を保ちつつ計算コストを削減。
• 種別ごとの対ペアカットオフと混合精度技術を組み込み、GPU 上でのスループットを最大化。