  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.
