Making Room for AI: Multi-GPU MD with Deep Potentials in GROMACS
Pennati, Hu, Peng, Müllender, Markidis (KTH) — arXiv:2604.07276 [cs.DC] 2026 — 最大32 GPU・DeePMD-kit統合
🎯 GROMACS + DeePMD-kitでマルチGPU量子精度MD。古典MD速度 × ab initio精度を大規模系で実現
① 背景と課題

古典MD（AMBER/CHARMM力場）はO(N)で高速だが量子効果（電荷移動・結合生成）を捉えられない。ab initio MDは量子精度だがO(N³)コストで大規模系に適用困難。MLIPがこのギャップを埋めるが、GROMACS既存実装のDeePMD連携はシングルGPUに制限されていた。

NNPotの旧実装: MPI all-to-all→プロセス0のみ推論→力再配布 = シングルGPU相当
GROMACSのeighth-shell DD（半リスト）はDeePMDのフルネイバーリスト要件と非互換

→ 独立した仮想DDレイヤーをNNPot内に導入し、GROMACSの既存DDを改変せずに解決

② 仮想DD実装フロー
各タイムステップ:
📡 MPI allgather（全NN原子座標をbroadcast）
↓ 🗺️ 均一カルテシアンGrid → Virtual Subdomain分割
↓ 👻 ローカル原子 + 2rc halo のゴースト原子抽出
↓ 🤖 各ランク独立にdeepmd::compute()（MPI非対応API）
↓ 📊 MPI allreduce で力を集約・再分配

Virtual DDはGROMACSのメインDDと完全独立。コード侵襲性を最小化しつつLAMMPS相当のスケーリングを実現。

③ DPA-1モデル選択の理由
  • ローカル記述子（単一カットオフ）→ ゴーストhalo 2rcで完全
  • DPA-2/3（メッセージパッシング）はhalo∝(l+1)×rc — スケーリング非効率
  • 1.6Mパラメータ、単精度FP32 → GPU効率最大化
  • 溶媒化タンパク質フラグメント2.6Mフレームで訓練済み
2M epoch
RTX 4080×1枚で約19時間で訓練完了
④ 主要結果 (a) Strong-scaling
ideal 100% 16 GPUs 32 GPUs 66% 40% Strong-scaling Efficiency
④ 主要結果 (b) Weak-scaling
ideal ≤16 GPU MI250x×32 A100×32 ~80% 48% 40% Weak-scaling Efficiency
④ 主要結果 (c) ウォールタイム内訳
コンポーネント割合
DeePMD推論>90%
MPIコレクティブ≤10%
その他~数%
主ボトルネック: カットオフ半径起因のゴーストアトム不可避コスト
第2ボトルネック: ランク間ロード不均衡（GROMACSのDLBが全原子対象）
④ 主要結果 (d) ソフトウェアスタック
コンポーネントバージョン
GROMACSv2025.2
DeePMD-kitv3.1.0
PyTorchv2.7.1
CUDA12.6.0 (A100)
ROCm6.3.3 (MI250x)
Benchmark1HCI 15,668原子
DPA-1 params1.6M (FP32)
⑤ テイクホームメッセージ
🏗️ 侵襲的改変なしで実現
GROMACSのメインDDを変更せず、NNPot内部の仮想DDで解決。既存コードへの影響を最小化しながらスケーラビリティを獲得。
⚡ 32 GPU・量子精度MD
古典力場には存在しないデノボ非共有結合・電荷移動を含む近ab initio精度MDを大規模タンパク質系（15,668原子）で実証。
📊 ボトルネック定量化
ゴーストアトムコストをスループットモデルで理論的に確認。ROCm Profilerで実測値と一致。将来の最適化指針を明確に示した。
🔮 LAM統合は今後の課題
DPA-2/3（メッセージパッシング型大規模原子モデル）のDD対応はLAMMPS相当のインフラが必要。より広い化学空間への拡張に向けた明確なロードマップ。
ケムインフォへの応用
適用先ユースケース
lib/mdDeePMD-kit連携ラッパーをRMSDAnalyzerに統合
lib/mdMLIPトラジェクトリのHBond解析パイプライン
lib/fep量子精度MDからMMGBSAへの接続（enhanced FEP）

GROMACS 2025.2 + DeePMD-kit v3.1で再現可能。DeepPot APIでlib/md/mlip/に実装予定。

本研究のインパクト
  • GROMACSでのMLIP完全分散メモリ推論を初実証
  • A100/MI250xの両GPUアーキテクチャに対応
  • ゴーストアトムコストとロード不均衡を定量化し将来改善指針を提示
Pennati, Hu, Peng, Müllender, Markidis (KTH) · arXiv:2604.07276 [cs.DC] · 2026年4月 · Category: computational_chemistry