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

TL;DR

This paper introduces HydraGNN, a scalable and energy-efficient graph neural network framework for atomistic materials modeling, capable of training on tens of thousands of GPUs with high performance and data diversity.

Contribution

It presents HydraGNN, a multi-headed GNN architecture that scales training of large GFMs to hundreds of millions of graphs across supercomputers, enabling advanced materials modeling.

Findings

Scales training to over 154 million graphs using 16,000 GPUs.

Achieves near-linear strong scaling on US-DOE supercomputers.

Demonstrates effective multi-task learning for atomistic property prediction.

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.