Scalable Training of Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN

要約

多頭グラフ畳み込みニューラル ネットワーク アーキテクチャである HydraGNN を使用したスケーラブルなグラフ基盤モデル (GFM) の開発とトレーニングに関する研究を紹介します。
HydraGNN は、トレーニング規模とデータ多様性の両方においてグラフ ニューラル ネットワーク (GNN) の境界を拡張します。
メッセージ パッシング アルゴリズムを抽象化し、GNN で畳み込みを定義するアルゴリズムの革新の再現と比較の両方を可能にします。
この研究では、数億のグラフで構成されるデータセット上で GFM トレーニングを数万の GPU にスケールアップできる一連の最適化について説明します。
当社の GFM は、マルチタスク学習 (MTL) を使用して、総エネルギーや原子間力などの原子構造のグラフ レベルとノード レベルの特性を同時に学習します。
トレーニングに 1 億 5,000 万を超える原子構造を使用して、米国エネルギー省 (US-DOE) の 2 つのスーパーコンピューター、つまり国立エネルギー研究科学計算センターのパールマッター ペタスケール システムと、
オークリッジ国立研究所のフロンティアエクサスケールシステム。
HydraGNN アーキテクチャにより、GFM は、Perlmutter では 2,000 個を超える GPU、Frontier では 16,000 個を超える GPU を使用して、ほぼ線形の強力なスケーリング パフォーマンスを実現できます。
Frontier 上の 64,000 を超える GPU でハイパーパラメータ最適化 (HPO) が実行され、GFM アーキテクチャが高精度で選択されました。
このような極端な規模のタスクを実行する際のエネルギー認識のために、各 GFM アーキテクチャに早期停止が適用されました。
最高ランクの GFM アーキテクチャのアンサンブルのトレーニングは、アンサンブル学習による不確実性定量化 (UQ) 機能を確立するために、収束するまで継続されました。
私たちの貢献により、大規模な計算リソースを使用した GFM の迅速な開発、トレーニング、展開への扉が開かれ、AI によって加速される材料の発見と設計が可能になります。

要約(オリジナル)

We present our work on developing and training scalable graph foundation models (GFM) using HydraGNN, a multi-headed graph convolutional neural network architecture. HydraGNN expands the boundaries of graph neural network (GNN) in both training scale and data diversity. It abstracts over message passing algorithms, allowing both reproduction of and comparison across algorithmic innovations that define convolution in GNNs. This work discusses a series of optimizations that have allowed scaling up the GFM training to tens of thousands of GPUs on datasets that consist of hundreds of millions of graphs. Our GFMs use multi-task learning (MTL) to simultaneously learn graph-level and node-level properties of atomistic structures, such as the total energy and atomic forces. Using over 150 million atomistic structures for training, we illustrate the performance of our approach along with the lessons learned on two United States Department of Energy (US-DOE) supercomputers, namely the Perlmutter petascale system at the National Energy Research Scientific Computing Center and the Frontier exascale system at Oak Ridge National Laboratory. The HydraGNN architecture enables the GFM to achieve near-linear strong scaling performance using more than 2,000 GPUs on Perlmutter and 16,000 GPUs on Frontier. Hyperparameter optimization (HPO) was performed on over 64,000 GPUs on Frontier to select GFM architectures with high accuracy. Early stopping was applied on each GFM architecture for energy awareness in performing such an extreme-scale task. The training of an ensemble of highest-ranked GFM architectures continued until convergence to establish uncertainty quantification (UQ) capabilities with ensemble learning. Our contribution opens the door for rapidly developing, training, and deploying GFMs using large-scale computational resources to enable AI-accelerated materials discovery and design.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google