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

Prasanna Balaprakash, Oak Ridge National Laboratory
Seminar
TPC Seminar Graphic featuring title of event with the date and speaker photo.

We present our work on developing and training scalable, trustworthy, and 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. We will discuss 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 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.

BIO:  Prasanna Balaprakash is the Director of AI Programs and a Distinguished R&D Scientist at Oak Ridge National Laboratory, where he directs research, development, and application of artificial intelligence and machine learning (AI/ML) to solve problems of national importance. Balaprakash’s research interests span AI, machine learning, optimization, and high-performance computing. He serves as the AI lead for several significant DOE-funded projects. He received the U.S. Department of Energy’s Early Career Award in 2018. Prior to his current role at ORNL, Balaprakash held several positions at Argonne National Laboratory, evolving from a postdoctoral researcher to an R&D Group Leader within the Mathematics and Computer Science Division, with a joint appointment at the Leadership Computing Facility. Previously, he served as the Chief Technology Officer at Mentis SA in Brussels, Belgium. Balaprakash received his PhD in 2010 from CoDE-IRIDIA (AI Lab), Université Libre de Bruxelles, Brussels, Belgium, where he received Marie Skłodowska-Curie and F.R.S-FNRS Aspirant fellowships from the European Commission and the Belgian-French Community’s National Fund for Scientific Research, respectively.

For more information about upcoming speakers please visit the TPC Seminar Series Webpage:

https://tpc.dev/tpc-seminar-series/