Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data

TL;DR

This paper introduces an exascale workflow using graph foundation models for materials discovery, enabling billion-scale screening of atomistic structures with high efficiency and transferability across diverse tasks.

Contribution

The work develops a scalable, multi-task graph model trained on extensive datasets, demonstrating fast screening and transfer learning for atomistic data at exascale.

Findings

Evaluated 1.1 billion structures in 50 seconds

Achieved transfer across twelve diverse chemical tasks

Demonstrated strong and weak scaling on multiple supercomputers

Abstract

We present an exascale workflow for materials discovery using atomistic graph foundation models built on HydraGNN. We jointly train on 16 open first-principles datasets (544+ million structures covering 85+ elements) using a multi-task architecture with per-dataset heads and a scalable ADIOS2/DDStore data pipeline. On Frontier, we execute six large-scale DeepHyper hyperparameter optimization campaigns in FP64 and promote the top-performing message-passing models to sustained 2,048-node training, yielding a PaiNN-based lead model. The resulting model enables billion-scale screening, evaluating 1.1 billion atomistic structures in 50 seconds, compressing a workload that would require years of first-principles computation, and supports data-scarce fine-tuning across diverse downstream tasks. We quantify precision-performance tradeoffs (BF16/FP32/FP64), demonstrate transfer across twelve chemically diverse downstream tasks, and establish seamless strong- and weak-scaling across Frontier, Aurora, and Perlmutter. This work allows fast and reliable exploration of vast chemical design spaces that are otherwise inaccessible to first-principles methods.