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

追記（2025-12-13）
本シリーズは方針を更新し、まず LLZO粒界におけるLi偏析・集積（短絡の前駆現象） を優先して解析します。最新ロードマップは 研究ノート（ハブ） をご覧ください。
本稿は「QE → extxyz → NequIP/Allegro 学習」パイプラインと GPU 学習スケーリング（A100 → H200 NVL 移行）に焦点を当てます。

Quantum ESPRESSO GPU 実行と ML-IAP (Allegro) 学習を統合し、AI-加速分子動力学の最適化を行うシリーズの第 2 回 (Revised) です。
本稿では、まず A100x4 ノード上で QE → extxyz → NequIP/Allegro 学習 までのパイプラインを実際に動かし、1 GPU / 4 GPU での学習ベンチマークを示します。
なお Ep2-R 時点では「環境確認＋ベンチマーク」のため、手元にある LLZO 系 AIMD のサブセット（LLZO/Li 界面データ） を題材にしていますが、方針更新により今後は LLZO 粒界（grain boundary） を優先してDFTデータセットを拡張し、H200 NVL (4 GPU / 8 GPU) 上でのスループットとスケーリング特性を評価していきます。
最終的なターゲットは、このフローを H200 NVL (4 GPU / 8 GPU) サーバーへそのまま移行し、従来の DeepMD ベースのワークロードに対して 2× 前後のスループット向上を目指すことです。

1. GPU 最適化 QE の設定

• NVIDIA HPC SDK 25.x ＋ Quantum ESPRESSO 7.4.1（CUDA 12 系）
• MPI ＋ OpenACC によるノード内並列（1 ノード内の複数 GPU を利用）
• LLZO 系（粒界・界面を含む）を対象に DFT AIMD を実行し、energy / forces / stress を含むトラジェクトリを取得
  • Ep2-R 時点では、まず LLZO/Li 界面の AIMD サブセットを用いてパイプライン動作を確認
  • 以降は LLZO 粒界（grain boundary） を優先してサンプリング対象を拡張
• QE 出力を .extxyz 形式へ変換し、そのまま NequIP / Allegro が読める形に整形
  • 今回のテストでは、llzo_li_aimd_allegro_clean.extxyz（約 8 MB）のサブセットを使用

2. Allegro 学習の流れ（NequIP 0.16.1 + Allegro 0.8.0）

• フレームワーク構成
  • PyTorch 2.6.0 + CUDA 12.4（公式 cu124 ホイール）
  • NequIP 0.16.1 / Allegro 0.8.0 / e3nn 0.5.8
• データ・前処理
  • nequip.data.datamodule.ASEDataModule
  • 入力: llzo_li_aimd_allegro_clean.extxyz
  • 元素: Li / O / Zr / La（type_names = [Li, O, Zr, La]）
  • 近接リスト: r_max = 6.0 Å の NeighborListTransform
• モデル
  • allegro.model.AllegroModel
  • カットオフ半径: r_max = 6.0 Å
  • スカラー特徴: num_scalar_features = 64
  • 損失関数: EnergyForceLoss（per-atom energy + forces）
  • EMA 付き LightningModule で学習
• 最適化と学習条件
  • Optimizer: Adam（学習率 1.0e-3）
  • バッチサイズ: 2（今回のベンチマーク条件）
  • Ep2-R では「環境確認＋ベンチマーク」のため 1 epoch だけ学習して時間を計測
  • 本番モデルでは、epoch 数やデータ量を増やし精度とスケールを両立する予定
• 出力
  • チェックポイント: epoch=*.ckpt / last.ckpt
  • nequip-compile により .nequip.pth（TorchScript）を生成し、LAMMPS Allegro ペアスタイルから利用可能

3. Ep2-R 時点の Allegro 学習ベンチマーク（A100x4）

まずは H200 NVL 実機ではなく、手元の A100x4 ノード を用いて Ep2-R 時点のベンチマークを取得しました。条件は次の通りです。

• ノード構成: NVIDIA A100 80GB PCIe × 4
• フレームワーク: PyTorch 2.6.0+cu124 / NequIP 0.16.1 / Allegro 0.8.0
• データ: llzo_li_aimd_allegro_clean.extxyz（LLZO/Li AIMD サブセット, 約 8 MB）
• 実験内容: 「1 epoch 学習 + テスト」を 1 GPU / 4 GPU × compile / eager の 4 パターンで計測
• real は、nequip-train プロセスの起動からテスト終了までの壁時計時間（データ統計の計算を含む）

今回用いた extxyz は 8 MB 程度の小規模データであり、「環境構築と分散学習が正しく動くか」の確認が主目的です。
そのため、4 GPU では通信や初期化のオーバーヘッドが支配的で、1 GPU に対する性能スケールは得られていません。今後は、方針更新に沿って LLZO 粒界（grain boundary）を優先して DFT データセットのサイズを増やし、バッチサイズや epoch 数を伸ばした「本番条件」で再計測することで、H200 NVL (4 GPU / 8 GPU) 上でのスループットとスケーリング特性を評価していきます。

4. 次章予告

Ep3-R では、ここで学習した NequIP / Allegro ベースの ML-IAP モデルを LAMMPS (Kokkos) に投入し、4 GPU / 8 GPU 構成での大規模 AI-MD を実行します。
まずは LLZO 粒界×Li偏析（前駆現象） を主題として、Li の偏析・拡散経路や局在（クラスター）を 3D で可視化しながら、H200 NVL サーバーでの実効 atom・step/s と既存 DeepMD 版との比較結果を報告する予定です。
その後、粒界で得た知見を界面シナリオへ接続し、短絡（フィラメント）に至る理解へ展開します。

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

現在、GPU 版 Quantum ESPRESSO 7.4.1 を用いた LLZO 系 DFT AIMD から、energy・forces・stress を含むラベル付きデータを extxyz／Allegro 形式へ変換するパイプラインが動作しており、llzo_li_aimd_allegro_clean.extxyz を使った NequIP / Allegro 学習が A100x4 ノード上で安定して回る状態になりました。
PyTorch Lightning ベースの分散学習フローについても、1 GPU / 4 GPU での eager / compile モードの実測結果を取得済みで、上記のように「小規模データでは DDP のオーバーヘッドが目立つ」ことを確認しています。

今後は、方針更新に従い、DFT データセットを まず LLZO 粒界（grain boundary）へ拡張（優先）し、その上でバルク LLZO・バルク Li、温度や Li カバー率の異なる界面、軽いひずみを加えた構造へと順次広げつつ、その都度 Allegro の学習・評価を行います。
その上で、本稿で示した A100x4 上のフローを H200 NVL (4 GPU / 8 GPU) 構成へそのまま持ち込み、「Ep0〜Ep2-R で構築した ML-IAP 学習パイプラインが、そのまま H200 NVL サーバーの性能を引き出せる」 ことを、実測スループットとスケーリングカーブで示していきます。
最終的には、Ep3-R における LAMMPS (Kokkos) 上での大規模 AI-MD と 3D 可視化を通じて、「H200 NVL 4GPU・8GPU サーバーにそのまま載せられる LLZO（粒界優先）向け Allegro ワークフロー」としてまとめる予定です。

• 最新ロードマップ（ハブ）｜LLZO短絡を「粒界×Li偏析」から
• 対応GPUサーバー一覧（見積相談OK）
• Founding 5（無料PoC）