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

arXiv: https://arxiv.org/abs/2406.12909 · HTML: https://mlsanity.com/paper/2406.12909/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 101 references that are in the map.
- First posted on arXiv 12 Jun 2024.

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour](https://mlsanity.com/paper/1706.02677/) (2017)
- [SciBERT: A Pretrained Language Model for Scientific Text](https://mlsanity.com/paper/1903.10676/) (2019)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [ChemBERTa: Large-Scale Self-Supervised Pretraining for Molecular Property Prediction](https://mlsanity.com/paper/2010.09885/) (2020)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [Hyper-Parameter Optimization: A Review of Algorithms and Applications](https://mlsanity.com/paper/2003.05689/) (2020)
- [Batch Bayesian Optimization via Local Penalization](https://mlsanity.com/paper/1505.08052/) (2015)
- [The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts](https://mlsanity.com/paper/2206.08917/) (2022)
- [PyTorch Distributed: Experiences on Accelerating Data Parallel Training](https://mlsanity.com/paper/2006.15704/) (2020)
- [ClimateBert: A Pretrained Language Model for Climate-Related Text](https://mlsanity.com/paper/2110.12010/) (2021)
- [Position: Graph Foundation Models are Already Here](https://mlsanity.com/paper/2402.02216/) (2024)
- [MatSciBERT: A Materials Domain Language Model for Text Mining and Information Extraction](https://mlsanity.com/paper/2109.15290/) (2021)
- [Scientific Large Language Models: A Survey on Biological & Chemical Domains](https://mlsanity.com/paper/2401.14656/) (2024)
- [Uncertainty Quantification over Graph with Conformalized Graph Neural Networks](https://mlsanity.com/paper/2305.14535/) (2023)
- [By-passing the Kohn-Sham equations with machine learning](https://mlsanity.com/paper/1609.02815/) (2016)
- [Uncertainty Quantification Using Neural Networks for Molecular Property Prediction](https://mlsanity.com/paper/2005.10036/) (2020)
- [DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks](https://mlsanity.com/paper/2104.06700/) (2021)
- [polyBERT: A chemical language model to enable fully machine-driven ultrafast polymer informatics](https://mlsanity.com/paper/2209.14803/) (2022)
- [Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets](https://mlsanity.com/paper/2310.04292/) (2023)
- [DGL-LifeSci: An Open-Source Toolkit for Deep Learning on Graphs in Life Science](https://mlsanity.com/paper/2106.14232/) (2021)
- [On the Scalability of GNNs for Molecular Graphs](https://mlsanity.com/paper/2404.11568/) (2024)
- [AutoDEUQ: Automated Deep Ensemble with Uncertainty Quantification](https://mlsanity.com/paper/2110.13511/) (2021)

## Built on by (most cited first, 8 of 8)

- [MatterTune: An Integrated, User-Friendly Platform for Fine-Tuning Atomistic Foundation Models to Accelerate Materials Simulation and Discovery](https://mlsanity.com/paper/2504.10655/) (2025)
- [Optimizing Data Distribution and Kernel Performance for Efficient Training of Chemistry Foundation Models: A Case Study with MACE](https://mlsanity.com/paper/2504.10700/) (2025)
- [Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data](https://mlsanity.com/paper/2506.21788/) (2025)
- [Automated Data Readiness for Scientific AI](https://mlsanity.com/paper/2607.02771/) (2026)
- [Ensemble Knowledge Distillation for Machine Learning Interatomic Potentials](https://mlsanity.com/paper/2503.14293/) (2025)
- [When Does Global Attention Help? A Unified Empirical Study on Atomistic Graph Learning](https://mlsanity.com/paper/2510.05583/) (2025)
- [Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data](https://mlsanity.com/paper/2604.15380/) (2026)
- [SetGo: Metadata Readiness for Scientific AI Datasets](https://mlsanity.com/paper/2607.22677/) (2026)

## Abstract (arXiv)

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.