[論文レビュー] High-performance training and inference for deep equivariant interatomic potentials

Chuin Wei Tan, Marc Descoteaux+12 more|ArXiv.org|Apr 22, 2025

Topic Modeling被引用数 10

ひとこと要約

この論文は NequIP を TorchInductor で分散トレーニング対応に全面改良し、Ahead-of-Time Inductor (AOTI) とカスタム Triton テンソル積カーネルを導入して、MD 推論の速度を最大18倍向上させ、SPICE 2 データで大規模トレーニングを可能にした。

ABSTRACT

Machine learning interatomic potentials, particularly those based on deep equivariant neural networks, have demonstrated state-of-the-art accuracy and computational efficiency in atomistic modeling tasks like molecular dynamics and high-throughput screening. The size of datasets and demands of downstream workflows are growing rapidly, making robust and scalable software essential. This work presents a major overhaul of the NequIP framework focusing on multi-node parallelism, computational performance, and extensibility. The redesigned framework supports distributed training on large datasets and removes barriers preventing full utilization of the PyTorch 2.0 compiler at train time. We demonstrate this acceleration in a case study by training Allegro models on the SPICE 2 dataset of organic molecular systems. For inference, we introduce the first end-to-end infrastructure that uses the PyTorch Ahead-of-Time Inductor compiler for machine learning interatomic potentials. Additionally, we implement a custom kernel for the Allegro model's most expensive operation, the tensor product. Together, these advancements speed up molecular dynamics calculations on system sizes of practical relevance by up to a factor of 18.

研究の動機と目的

• 大規模データセット上でのMLIPの訓練の拡張性と性能を向上させる。
• 従来のMDコードとのエンドツーエンド高性能推論統合を可能にする。
• 深い等変原子間ポテンシャルのためのエンドツーエンドのコンパイラベース加速を提供する。
• 大規模で多様なデータセット（SPICE 2）で実用的な利得を実証する。
• 今後のMLIPアーキテクチャ（例: Allegro）に向けたフレームワークの拡張性を高める。

提案手法

• 微分を含むモデル全体を単一の PyTorch FX グラフにトレインタイムのコンパイルと推論互換性のために追跡する。
• torch.compile を介して forward/backward の最適化カーネルを生成するために TorchInductor を適用する。
• 勾配バケツ化を回避してフュージョン機会を最大化するカスタム DDP アプローチによる分散データ並列トレーニングを実装する。
• Ahead-of-Time Inductor (AOTI) を導入して、コンパイル済み PyTorch モデルを非 Python 環境（例: LAMMPS）で使用するためのネイティブコードとしてエクスポートする。
• Allegro テンソル積用のカスタム融合 Triton カーネルを開発して PyTorch 実装の非融合部分を置換する。
• メモリ使用量を削減するために Wigner 3-j 収束係数を表す圧縮疎行列形式を使用し、中間量を実体化しない。

実験結果

リサーチクエスチョン

• RQ1 TorchInductor をどのように活用して深い等変原子間ポテンシャルの訓練と推論を加速できるか？
• RQ2 高性能MDコード（例: LAMMPS）と推論のための AOTI を統合した場合の性能向上はどの程度か？
• RQ3 テンソル積のカスタム融合カーネルは Allegro でのメモリと計算ボトルネックをさらに削減できるか？
• RQ4 SPICE 2 のような大規模MLIPデータセットに対する分散マルチGPUトレーニングはどの程度スケールするか？
• RQ5 異なる電荷状態サブセットで SPICE 2 上の Allegro モデルを訓練した場合の精度と一般化への影響はどのようになるか？

主な発見

指標	小モデル	中モデル	大モデル
障壁 MAE [meV]	22.75	15.42	11.37
障壁 RMSE [meV]	32.36	21.77	15.38

• torch.compile を用いた訓練は、テストされたGPU上で TorchScript より 2.4–5.0 倍の速度向上をもたらす。
• 分散トレーニングは 128ランクまで良好にスケールし、256ランクでは合理的な効率（報告ベースラインで AMD 40%、NVIDIA 24%）を示す。
• AOTI と最適化されたテンソル積カーネルによる推論は、小分子および水ボックスで TorchScript より最大18倍の速度向上を達成。
• AOTI は大規模システムのメモリ効率の高いスケーラブル推論を可能にし、GPUあたりの実行可能なシステムサイズを拡張する（例: ハードウェアに応じて大モデルは4320–5184 原子まで）。
• 生体分子系（DHFR、セルロース）での強連結は Frontier と Perlmutter で最大256ノードまで優れた性能を示し、AOTI+TP の顕著なスループット向上を示す。
• SPICE 2 ベンチマークと比較して、3つの Allegro モデル（小/中/大）は、モデルサイズの増加に伴い MAE/RMSE が低下する有利なトーション障壁予測を示す。