Exascale Multi-Task Graph Foundation Models for Imbalanced, Multi-Fidelity Atomistic Data
Authors: Massimiliano Lupo Pasini, Jong Youl Choi, Kshitij Mehta, Richard Messerly, Rylie Weaver, Linda Ungerboeck, Isaac Lyngaas, Benajmin Stump, Ashwin M. Aji, Karl W. Schulz, Jorda Polo
Organizations: Computational Sciences and Engineering Division, Oak Ridge National Laboratory · Computer Science and Mathematics Division, Oak Ridge National Laboratory · National Center for Computational Sciences Division, Oak Ridge National Laboratory · Bredesen Center Data Science and Engineering, University of Tennessee, Knoxville · AMD, Advanced Micro Devices
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.
Figures & tables
Attribute
Value
Category of achievement
scalability / time-to-solution / peak performance
Type of method used
n/a
Results reported on the basis of
whole application including I/O
Precision reported
mixed precision
System scale
results measured on full-scale system
Measurement mechanism
timers (GPTL); other (Omnistat telemetry)
TABLE I: Performance Attributes Required by ACM Gordon Bell Nomination
Dataset
Type of compounds
Approximation theory
Samples
Files
QM7-X [ 8 ]
Organic molecules
DFT (PBE0+MBD)
8,390,474
8
QCML-DFT [ 9 ]
Organic molecules
DFT (PBE0)
33,496,167
10
Transition1x [ 10 ]
Organic molecules
DFT ( ω b97x/6-31G(d))
9,644,740
1
ANI-1x [ 11 ]
Organic molecules
DFT ( ω b97x/6-31G(d))
4,956,005
1
∇2 -DFT [ 12 ]
Organic molecules
DFT ( ω b97x-D/def2-SVP)
10,392,954
3
Materials Project Trajectories (MPTrj) [ 13 ]
3D bulk alloys
DFT (PBE0)
1,580,395
1
TABLE II: Multi-source, multi-fidelity first-principles datasets used in this work.
Fig. 1: Multi-task learning training for MLIPs on multi-source, multi-fidelity data, with shared message-passing representations and dataset-specific supervision heads. Message passing layers learn fundamental features that describe basic physics of interatomic interactions common to all datasets. Output heads customize features to the physics captured by the approximation theory used to generate a specific dataset
TABLE III: Hyperparameters tuned in all single-model HPO campaigns. Convolution depth and hidden dimensions were sampled as integers, force weight and learning rate were sampled as a continuous real-valued variable.
TABLE IV: Additional HPO dimensions exposed for each MPNN family in the local SC26 HPO driver.
Fig. 2: Validation-loss trajectories for all trials from the six Frontier HPO campaigns, with one campaign per MPNN backbone. It shows 181 successful DeepHyper trials out of 574 explored configurations across MPNN backbones.
Parameter
Value
Convolution layers
2
Hidden dimension
337
Interaction filters
126
Radial basis functions
4 (Bessel)
Cutoff radius
5 Å
Max neighbors
20
TABLE V: Lead PaiNN model configuration selected from the HPO campaign and used for all downstream analysis.
Fig. 3: Strong and weak scaling results, along with time decomposition, are presented for Perlmutter, Aurora, and Frontier. Speedup (y-axis) is normalized to the 32-node performance on each system, with the number of GPUs shown on the x-axis. The decomposed timing corresponds to the weak scaling results, detailing time spent on data loading, forward pass, backward pass, and optimization.
Dataset
Size
Domain
Target(s)
QM9 [ 52 ]
∼ 130k
Small organic molecules (up to 9 heavy atoms; C, H, O, N, F)
Atomization energy
MD17 [ 53 ]
∼ 134k
Uracil ab-initio MD trajectory
Total energy and forces
Wiggle150 [ 54 ]
150
Strained conformers of adenosine, benzylpenicillin, efavirenz
Relative conformer energy
MS25 - MgO-2 × 2 [ 55 ]
10k
MgO(100) oxide surface
Per-atom energy
MS25 - MgO-4 × 4 [ 55 ]
10k
MgO(100) oxide surface
Per-atom energy
MS25 -Reaction [ 55 ]
10k
CH 4 dissociation on Pt(111)
Per-atom energy
TABLE VI: Downstream fine-tuning task suite. All tasks predict graph-level (per-structure) properties. Size indicates the number of structures available for fine-tuning.
Dataset
Split
Epochs
Batch Size
Learning Rate
QM9
70/15/15
100
32
10−3
MD17
70/15/15
100
32
10−4
Wiggle150
70/15/15
500
32
10−4 a
MS25
70/15/15
10
2
10−4
OQMD
70/15/15
100
32
10−4
ABX 3
70/15/15
100
32
10−4
TABLE VII: Training configuration for each benchmark. All runs use AdamW and FP64 unless otherwise noted.
Method
MAE (eV/atom)
Training from scratch
0.047
Frozen backbone
0.050
Unfrozen backbone
0.004
TABLE VIII: QM9 validation MAE (eV/atom) after 100 epochs for atomization energy at 0 K ( U0 , property index 12).
Energy MAE
Forces MAE
Method
kcal/mol
kcal/(mol Å)
Training from scratch
3.421
8.438
Frozen backbone
4.273
13.663
Unfrozen backbone
1.481
6.543
TABLE IX: MD17 uracil MLIP benchmark: best validation over 100 epochs.
Method
MAE (kcal/mol)
RMSE (kcal/mol)
Training from scratch
19.914
29.203
Frozen backbone
24.715
34.212
Unfrozen backbone
18.732
23.953
TABLE X: Validation performance on Wiggle150 relative conformer energies (kcal/mol).
System
Scratch
Frozen Backbone
Unfrozen Backbone
MgO-2 × 2
0.030
0.030
0.017
MgO-4 × 4
0.042
0.046
0.003
HEA
0.032
0.032
0.031
Reaction
0.047
0.046
0.019
Zr-O
0.293
0.293
0.152
TABLE XI: MS25 Validation MAE (eV/atom) after 10 epochs. Energy-only supervision; no force labels used.
Method
OQMD
ABX 3
FP64 training from scratch
0.596
0.444
BF16 unfrozen backbone
0.684
0.352
FP32 unfrozen backbone
0.245
0.341
FP64 unfrozen backbone
0.227
0.310
TABLE XII: Validation MAE (eV/atom) after 100 epochs for OQMD and ABX 3 .
Fig. 4: Validation-loss trajectories for all fine-tuning trials on the OQMD dataset. A smoothing window of 5 epochs was used.
Method
Avg. MAE
Std. MAE
Avg. RMSE
Max. Error
Scratch
47.763
15.215
116.756
1,561.365
Frozen
50.714
15.835
116.951
1,525.047
Unfrozen
48.346
14.418
112.694
1,536.159
TABLE XIII: MB-jdft2d (Exfoliation Energy in meV/Atom) Five-Fold Cross Validation Statistics after 100 epochs of fine-tuning.
ROCAUC
F1
Accuracy
Method
Avg.
Std.
Avg.
Avg.
Training from scratch
0.911
0.002
0.809
0.835
Frozen backbone
0.924
0.003
0.823
0.847
Unfrozen backbone
0.924
0.003
0.825
0.848
TABLE XIV: MB-is_metal Five-Fold Cross Validation Statistics after 10 epochs of fine-tuning.
Fig. 5: Fused Gradient Inference Pipeline. The encoder executes once and its activations are cached. All 16 decoder heads run forward-only, their energies are combined into a weighted sum, and a single backward pass computes forces.
Graph Foundation Models are pretrained graph-learning systems designed to transfer across datasets, tasks, and domains with different structures and features.
General-purpose large language models (LLMs) that rely on in-context learning do not reliably deliver the scientific understanding and performance required for drug discovery tasks. Simply increasing model size or introducing reasoning tokens does not yield significant performance gains. To address this gap, we introduce the MMAI Gym for Science, a one-stop shop molecular data formats and modalities as well as task-specific reasoning, training, and benchmarking recipes designed to teach foundation models the 'language of molecules' in order to solve practical drug discovery problems. We use MMAI Gym to train an efficient Liquid Foundation Model (LFM) for these applications, demonstrating that smaller, purpose-trained foundation models can outperform substantially larger general-purpose or specialist models on molecular benchmarks. Across essential drug discovery tasks - including molecular optimization, ADMET property prediction, retrosynthesis, drug-target activity prediction, and functional group reasoning - the resulting model achieves near specialist-level performance and, in the majority of settings, surpasses larger models, while remaining more efficient and broadly applicable in the domain.
Maksim Kuznetsov, Zulfat Miftahutdinov, Rim Shayakhmetov +17
Universal Machine Learning Interatomic Potentials (uMLIPs), pre-trained on massively diverse datasets encompassing inorganic materials and organic molecules across the entire periodic table, serve as foundational models for quantum-accurate physical simulations. However, uMLIP training requires second-order derivatives, which lack corresponding parallel training frameworks; moreover, scaling to the billion-parameter regime causes explosive growth in computation and communication overhead, making its training a tremendous challenge. We introduce MatRIS-MoE, a billion-parameter Mixture-of-Experts model built upon invariant architecture, and {Janus}, a pioneering high-dimensional distributed training framework for uMLIPs with hardware-aware optimizations. Deployed across two Exascale supercomputers, our code attains a peak performance of 1.2/1.0 EFLOPS (24%/{35.5%} of theoretical peak) in single precision at over 90% parallel efficiency, compressing the training of billion-parameter uMLIPs from weeks to hours. This work establishes a new high-water mark for AI-for-Science (AI4S) foundation models at Exascale and provides essential infrastructure for rapid scientific discovery.
Yuanchang Zhou, Hongyu Wang, Yiming Du +12
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · Independent Researcher · Sun Yat-Sen University
Energy systems face converging pressures for security, affordability, resilience, and sustainability, creating a need for faster discovery of deployable energy materials. Here we introduce SCALE (Simulation-Calibrated Amortized Learning for Energy Materials), a physics-grounded, real-world-data-calibrated learning architecture that connects deterministic scientific operators, experimental calibration, expanded calibrated label generation, and transformer-scale inference. SCALE converts selected high-cost mechanistic computation and measured evidence into reusable models for rapid screening, ranking, inverse design, and active learning. We formulate the framework, identify ten method-based application regimes, and demonstrate SCALE for solid-state metal-hydride hydrogen-storage capacity prediction. In this implementation, a hydride phase-equilibrium capacity operator is calibrated against 381 measured ML-HydPARK capacity anchors and used to generate 5,000 candidate-condition-prototype teacher labels. A crystallographically anchored periodic-graph representation preserves atomic sites, periodic neighbor relationships, and local metal environments absent from formula-only encodings. An edge-biased graph transformer with 2.90 million parameters reproduces calibrated teacher labels with five-fold surrogate fidelity of MAE 0.0582 wt% H2, RMSE 0.0833 wt% H2, R2 = 0.9927, and Pearson r = 0.9963. Post hoc attention analysis suggests that SCALE learns chemically organized element groupings and metal-metal relationships consistent with established hydride chemistry, without chemistry-group labels as supervision. Once trained, SCALE shifts million-candidate evaluation from repeated deterministic workflow execution to batched learned inference, reducing per-candidate screening cost by approximately 10^7-10^8 while retaining links to simulation and experimental evidence.