Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN

要約

HydraGNNは、多頭のグラフ畳み込みニューラルネットワークアーキテクチャである。HydraGNNは、学習規模とデータの多様性の両方において、グラフ・ニューラル・ネットワーク（GNN）計算の境界を拡張する。HydraGNNはメッセージパッシングアルゴリズムを抽象化し、GNNにおける最近傍畳み込みを定義するアルゴリズムの革新の再現と比較の両方を可能にする。本研究では、数億のグラフからなるデータセットにおいて、GFMの学習を数万のGPUにスケールアップすることを可能にした一連の最適化について議論する。我々のGFMはマルチタスク学習(MTL)を用いて、エネルギーや原子力といった原子構造のグラフレベルとノードレベルの特性を同時に学習する。学習用に1億5400万以上の原子論的構造を用いて、我々のアプローチの性能とその教訓を2つの最新鋭の米国エネルギー省（US-DOE）スーパーコンピュータ、すなわち国立エネルギー研究科学計算センターのPerlmutterペタスケールシステムとOak Ridge Leadership Computing FacilityのFrontierエクサスケールシステムで説明する。HydraGNNアーキテクチャにより、Perlmutterでは2,000個以上のGPUを、Frontierでは16,000個以上のGPUを使用して、GFMは線形に近い強力なスケーリング性能を達成することができます。

要約(オリジナル)

We present our work on developing and training scalable, trustworthy, and energy-efficient predictive graph foundation models (GFMs) using HydraGNN, a multi-headed graph convolutional neural network architecture. HydraGNN expands the boundaries of graph neural network (GNN) computations in both training scale and data diversity. It abstracts over message passing algorithms, allowing both reproduction of and comparison across algorithmic innovations that define nearest-neighbor convolution in GNNs. This work discusses a series of optimizations that have allowed scaling up the GFMs training to tens of thousands of GPUs on datasets consisting 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 energy and atomic forces. Using over 154 million atomistic structures for training, we illustrate the performance of our approach along with the lessons learned on two state-of-the-art 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 Leadership Computing Facility. 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.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL