[論文レビュー] Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning

Weile Jia, Handong Wang+6 more|arXiv (Cornell University)|May 1, 2020

Machine Learning in Materials Science参考文献 72被引用数 35

ひとこと要約

機械学習ベースのMDプロトコル（Deep Potential Molecular Dynamics）を示し、100百万原子を超える系に対して第一原理精度を達成する。Summit上で高度に最適化されたDeePMD-kitにより最大91 PFLOPSの倍精度を達成し、1日あたりナノ秒スケールのシミュレーションを可能にする。

ABSTRACT

For 35 years, {\it ab initio} molecular dynamics (AIMD) has been the method of choice for modeling complex atomistic phenomena from first principles. However, most AIMD applications are limited by computational cost to systems with thousands of atoms at most. We report that a machine learning-based simulation protocol (Deep Potential Molecular Dynamics), while retaining {\it ab initio} accuracy, can simulate more than 1 nanosecond-long trajectory of over 100 million atoms per day, using a highly optimized code (GPU DeePMD-kit) on the Summit supercomputer. Our code can efficiently scale up to the entire Summit supercomputer, attaining $91$ PFLOPS in double precision ($45.5\%$ of the peak) and {$162$/$275$ PFLOPS in mixed-single/half precision}. The great accomplishment of this work is that it opens the door to simulating unprecedented size and time scales with {\it ab initio} accuracy. It also poses new challenges to the next-generation supercomputer for a better integration of machine learning and physical modeling.

研究の動機と目的

• 精度を損なうことなく、第一原理分子動力学（AIMD）を大規模系（数千万〜数億原子）および長い時間スケールへと拡張する。
• 高価な電子構造計算を保持する第一原理忠実性を保った代理モデルへと置換するため、機械学習を活用する。
• 最新のスパコンにおける異種GPUアーキテクチャを最大限に活用するよう、HPC指向の実装を開発・最適化する。

提案手法

• Deep Potential (DP)ニューラルネットワークを用いて、原子間ポテンシャルを原子エネルギーの和として表現する。
• DPモデルを第一原理データで訓練し、LAMMPSと統合されたDeePMD-kitに実装してMD積分を行う。
• データレイアウト、近傍リスト、カスタムTensorFlowオペレータを最適化して、Summit上でGPUのスループットを最大化する。
• 混合精度スキーム（MIX-32とMIX-16）を用いて計算を加速しつつ精度を維持する。
• TensorFlowオペレータの再構成（GEMMベースの置換、CUDAカーネル）とMPI通信の削減により、スケーラビリティを向上させる。

実験結果

リサーチクエスチョン

• RQ1 深層学習ベースの原子間ポテンシャルは、非常に大規模な系（ tens〜hundreds of millions 原子）および長いMD軌道に対して第一原理精度を再現できるか？
• RQ2 DPベースMDを異種系スパコン（Summitのような）でスケールさせるために必要なHPC最適化と混合精度戦略は何か？
• RQ3 DP-MDを極端な系サイズへ推進した場合の実用的な性能向上（FLOPS、解決時間）と精度のトレードオフはどうなるか？
• RQ4 DP-MDは材料・液体系のスケーラビリティと忠実度の面で、従来のAIMDや経験的力場と比較してどうなるか？

主な発見

• DP-MDは Summit 上で4,560ノードを用いて、91 PFLOPS（倍精度）および混合単一/半精度で162 PFLOPS / 275 PFLOPSの軌道を持つ1日あたりのAIMD様軌道を可能にする。
• 127百万原子の銅系では、1ステップあたりの解決時間が倍精度で8.1e-10 s/step/atom、日あたり約0.8 ns、混合半精度では日あたり約2.5 ns。
• 近傍リストデータレイアウトの最適化と近傍エントリの64ビット整数圧縮により分岐を排除しGPU効率を改善、単一GPUでのカスタムTensorFlowオペレータの速度向上は基準CPU実装に対して64.6xの速度アップを達成。
• MATMULとSUMをGEMMに置換し、TANHとTANHGradを統合し、CUDAカーネルを統合することでTensorFlowオペレータの効率が大幅に向上し、MDループ全体のスピードアップに寄与。
• 混合精度スキームMIX-32とMIX-16は、試験された水の構成で倍精度に匹敵するエネルギー/力の精度を達成し、エネルギーと力の誤差においてMIX-32は倍精度に匹敵する。