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

TL;DR

This work demonstrates scalable, trustworthy predictive graph foundation models for atomistic materials by integrating HydraGNN with extreme-scale computing. It combines large-scale data aggregation (over 154 million structures), scalable data management (ADIOS and in-memory DDStore), multitask learning for energy and forces, and scalable hyperparameter optimization with asynchronous Bayesian search, plus ensemble-based epistemic UQ. The results show near-linear strong scaling up to thousands of GPUs on Frontier and Perlmutter, substantial I/O throughput, and energy-aware pre-training that yields robust, transferable GFMs suitable for DOE-inspired materials discovery. The study also outlines practical pathways for reducing training energy, improving load balance, and deploying these models to downstream tasks and active learning regimes for accelerated materials design.

Abstract

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.

Figures (31)

• Figure 1: Normalized histograms for each dataset of the number of atoms within an atomistic structure (top) and number of edges within the graph representation of each atomistic structure (bottom).
• Figure 2: Heatmap that describes the frequency of occurrence of each element of the periodic table across data sampled resulting from the aggregation of the datasets ANI1x, QM7-X, OC2020, OC2022, and MPTrj.
• Figure 3: Normalized histograms for each dataset of the total energy per atom (top) , energy per atom after removing the linear regression term from each dataset (center), and the $L^2$-norm of the force tensor (bottom) for each dataset after removing data samples with $L^2$-norm of the force tensor unreasonably higher than 100 eV/angstrom.
• Figure 4: Histograms showing the energy and force distributions across five datasets, along with the distribution breakdown for the training, validation, and test sets. The y-axis shows the probability density on a logarithmic scale.
• Figure 5: Different approaches for shuffling data during the training process. In a), data is read from the shared file system in which each graph object is stored in its own separate file (high file system metadata overhead). In b), data is read from an ADIOS file (low metadata overhead, high I/O bandwidth). In c), all data is read once into DDStore 10.1145/3624062.3624171, an in-memory data store which uses MPI one-sided RMA operations to obtain data from remote processes (best performance).