English

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

Materials Science 2026-04-20 v1 Artificial Intelligence

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.

Keywords

Cite

@article{arxiv.2604.15380,
  title  = {Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data},
  author = {Massimiliano Lupo Pasini and Jong Youl Choi and Kshitij Mehta and Richard Messerly and Rylie Weaver and Linda Ungerboeck and Isaac Lyngaas and Benajmin Stump and Ashwin M. Aji and Karl W. Schulz and Jorda Polo},
  journal= {arXiv preprint arXiv:2604.15380},
  year   = {2026}
}

Comments

12 pages; 5 figures; 15 tables

R2 v1 2026-07-01T12:13:19.641Z