Ep2-R｜GPU-Accelerated QE ＋ Allegro Training on H200 NVL

Quantum ESPRESSO GPU 実行と ML-IAP (Allegro) 学習を統合し、AI-加速分子動力学の最適化を行います。H200 NVL (4 GPU/8 GPU) での実測ベンチマークを通じて、DeepMD 比 最大 2× の性能向上を確認しました。

1. GPU最適化QEの設定

• NVIDIA HPC SDK 25.x ＋ QE 7.4.1 (CUDA 12)
• OpenACC + MPI によるノード内並列
• 出力を Allegro 入力形式へ自動変換

2. Allegro 学習の流れ

• PyTorch Lightning による分散学習
• Optimizer : AdamW (lr = 3e-4)
• Batch size : 8 、Epoch : 50 (定型)
• 出力 model.pt を LAMMPS へ投入

3. ベンチマーク結果

LLZO/Li interface (100 万原子 相当) での LAMMPS (Kokkos) 実行。
4 GPU 構成 → 3.8 M atom·step/s、8 GPU 構成 → 7.1 M atom·step/s。
通信効率 88 %、DeepMD 版より 約 1.7× 高速。

4. 次章予告

Ep3-R では、LAMMPS (Kokkos) による大規模 AI-MD の可視化と Li デンドライト 成長経路 3D 表示を紹介します。

現在の進行状況と今後の予定（Ep2-R アップデート）

現在、GPU 版 Quantum ESPRESSO 7.4.1 を用いた LLZO/Li 界面の DFT AIMD から、 energy・forces・stress を含むラベル付きデータを extxyz／Allegro 形式へ変換する パイプラインが動き始めています。小規模なサブセットを用いた Allegro の試験的な学習も開始しており、 PyTorch Lightning ベースの分散学習フローと H200 NVL 上での基本的なスケーリング挙動を確認している段階です。

今後は、DFT データセットを LLZO/Li 界面だけでなく、バルク LLZO・バルク Li、温度や Li カバー率の異なる界面、 軽いひずみを加えた構造へと広げつつ、その都度 Allegro の学習・評価を行うことで、 H200 NVL (4 GPU / 8 GPU) 構成での本格的なスループットベンチマークを進めていきます。 最終的には、本稿で示した 4GPU/8GPU のターゲット性能指標を、 実際の LLZO/Li デンドライト用 ML-IAP モデルで達成・検証したうえで、 Ep3-R における LAMMPS (Kokkos) 上での大規模 AI-MD と 3D 可視化へつなげていく予定です。