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.

基于 HydraGNN 的可扩展图基础模型 (GFM) 的开发和训练，具有高度的训练规模和数据多样性，以及并行训练算法的创新；我们使用超级计算资源展示了多个优化策略和性能，开发出适用于材料发现和设计的 AI 加速的 GFM 模型。