コンテンツにスキップ

マルチGPU対応の汎用GNN力場SevenNetによる分子動力学計算ベンチマーク#

SevenNet(Scalable EquiVariance Enabled Neural Network)12は、グラフニューラルネットワーク(GNN)を使って原子間の相互作用を計算し、様々な組成の系の分子動力学計算を行うことができる汎用の力場です。GNN力場の中でも、SevenNetはCUDA-aware OpenMPIを使うことで、複数のGPUを使った並列実行による高速化ができることが特長となっています。

SevenNetはLAMMPSに実装する形で公開されていますが、ナノ材料解析統合GUI Advance/NanoLaboを使うことで、GUI上の操作でSevenNetを使った計算を行うことができます。

本事例では、硫化物リチウムイオン伝導体Li10GeP2S12の大規模系での分子動力学計算を、実際に複数のGPUを使って並列実行し、ベンチマークを行いました。

計算環境#

本事例で使用した計算環境のスペックを以下に示します。

  • CPU：Intel Xeon Gold 6338（32コア）×2
  • GPU：NVIDIA A100 (VRAM 40GB)×8
  • メモリ：1024GB
  • CUDA：12.2

計算環境は富士通株式会社にご協力いただき、Fujitsu Computing as a Service(CaaS)の「Fujitsu クラウドサービス HPC」を利用して作成しました。

計算条件#

Materials Projectより取得したLi10GeP2S12の構造ファイル(mp-696128)をもとにスーパーセルを生成し、21600原子および98000原子の大規模系を用意しました。これらの系に対し、学習済みモデルparallel SevenNet-0を使って、NVTアンサンブル(500K)・時間刻み幅0.5fsで100ステップの分子動力学計算を実行しました。同じ計算をGPU使用数1から8まで変えながら実行し、計算時間を計測しました。GPUあたりのMPIプロセス数は1、プロセスあたりのスレッド数は1としました。

ベンチマーク結果#

各GPU使用数での計算時間と、同じ条件で1日実行した場合のシミュレーション時間を表で示します。21600原子系ではGPU使用数1、98000原子系ではGPU使用数4以下の場合にGPUメモリ不足(OOM)となりました。この程度の規模の系では複数GPUが必須となるようです。

100ステップの計算時間(s) 1日実行時のシミュレーション時間(ns/day)
21600原子系 98000原子系 21600原子系 98000原子系
GPU×1 OOM OOM OOM OOM
GPU×2 74 OOM 0.060 OOM
GPU×3 51 OOM 0.088 OOM
GPU×4 40 OOM 0.112 OOM
GPU×5 34 204 0.132 0.022
GPU×6 28 115 0.163 0.039
GPU×7 25 101 0.184 0.044
GPU×8 23 87 0.204 0.052

なお、SevenNetの論文1ではA100 80GBを8枚使用してSi3N4の計算を行い、10万原子系でおよそ0.1ns/day3という計算速度を報告しており、それと概ね同程度の計算ができています。

また、GPU使用数に対する並列効率を見るため、相対計算速度をプロットします。通常はGPU使用数1を基準としますが、今回はそれが使えないため、21600原子系・GPU使用数2の計算時間の半分を基準とします。98000原子系についても、計算時間が原子数に比例するとして同じ基準を使います。

GPU使用数に対する相対計算時間のプロット

理想的には相対計算速度はGPU使用数に比例し、これに近いほど並列効率が良いということになります。21600原子系ではGPU使用数が多くなるほど並列効率が徐々に下がっていますが、98000原子系では6～8であまり変化しておらず、原子数が多いと複数GPUの性能を効率的に使えるようです。

また、GPU使用数5の時に計算速度が低くなっています。当初、各GPUへの原子の割り当てに偏りがあるため非効率になっているのではと疑いましたが、確認したところ原子数は等分されていました。おそらくセル分割の形状により、GPU間の通信量が多くなっているなどの理由によると考えられます。

このように、SevenNetによる分子動力学計算を複数のGPUによる並列実行で高速化できることが分かりました。

関連ページ#

  1. Y. Park, J. Kim, S. Hwang, and S. Han, "Scalable Parallel Algorithm for Graph Neural Network Interatomic Potentials in Molecular Dynamics Simulations". arXiv, arXiv:2402.03789. (2024) (https://arxiv.org/abs/2402.03789) ↩↩

  2. S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt, and B. Kozinsky, "E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials". Nat. Commun., 13, 2453. (2022) (https://www.nature.com/articles/s41467-022-29939-5) ↩

  3. シミュレーション時間は時間刻み幅に依存するため、そのままの比較はできません。 ↩