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

This paper presents an exascale multi-task graph foundation model built on HydraGNN and trained on over 544 million atomistic structures across 16 datasets, which achieves billion-scale materials screening in seconds and enables efficient fine-tuning for diverse downstream tasks by leveraging high-performance computing resources like Frontier.

The Gist

Imagine you are trying to design a new, super-efficient battery or a life-saving medicine. To do this, scientists usually have to simulate how atoms interact using "first-principles" methods (like Density Functional Theory). Think of this as trying to understand a car engine by building a brand-new, perfect engine from scratch for every single part you want to test. It's incredibly accurate, but it takes so long and costs so much that you can only test a few dozen ideas in a year.

This paper describes a breakthrough that changes the game entirely. The researchers built a "Super-Brain" for atoms that can test 1.1 billion potential materials in just 50 seconds. That's the difference between spending years in a library reading one book at a time versus having a robot that can read the entire Library of Congress in the time it takes to brew a cup of coffee.

Here is how they did it, broken down with simple analogies:

1. The Problem: Too Many Cooks, Too Many Recipes

Usually, AI models for atoms are trained on just one type of data (like only organic molecules or only metals). It's like teaching a chef to cook only Italian food. If you ask them to make sushi, they fail.

• The Challenge: The researchers wanted to train a model on 16 different datasets containing over 544 million different atomic structures. These datasets were messy, unbalanced (some had millions of examples, others only a few), and used different scientific "languages" (different levels of accuracy).
• The Analogy: Imagine trying to teach a student by giving them 16 different textbooks written in different languages, by different authors, with some chapters missing and others written in crayon. Most students would get confused and give up.

2. The Solution: The "Specialized Team" Approach (Multi-Task Learning)

Instead of forcing one single brain to memorize everything at once, they built a HydraGNN (a type of AI architecture).

• The Shared Brain: The core of the model learns the universal rules of physics (how atoms generally stick together). This is the "shared message-passing" layer.
• The Specialized Heads: Attached to this core are 16 different "heads" (specialists). One head is an expert on organic molecules, another on metals, another on catalysts.
• The Analogy: Think of a massive hospital. The "Shared Brain" is the general medical knowledge all doctors have (anatomy, physiology). The "Heads" are the specialists: a cardiologist, a neurologist, a dermatologist. They all share the same foundational knowledge but apply it to their specific area. This prevents the model from getting confused when switching between different types of data.

3. The Engine Room: Exascale Computing

To train this model, they used Frontier, one of the world's fastest supercomputers, utilizing 16,384 GPUs (graphics cards) working in perfect unison.

• The Analogy: If training a normal AI model is like a single person digging a hole with a spoon, this project was like 16,000 people digging with bulldozers simultaneously.
• The Logistics: They had to move massive amounts of data without clogging the system. They used a special pipeline (ADIOS2/DDStore) that acts like a high-speed conveyor belt, bringing the data right next to the workers so no one has to wait in line.

4. The Selection Process: The "Talent Show"

Before settling on the final model, they ran a massive Hyperparameter Optimization (HPO) campaign.

• The Analogy: Imagine holding a talent show with 6 different types of singers (different AI architectures). They auditioned hundreds of combinations to see which singer could perform the best song in the shortest time.
• The Winner: They found that a specific architecture called PaiNN was the "Goldilocks" model—it wasn't too heavy, wasn't too light, and was the fastest at learning while maintaining high accuracy.

5. The Result: Billion-Scale Screening

Once trained, this model can screen 1.1 billion atomic structures in 50 seconds.

• The Impact: Doing this with traditional methods would take 6.7 years of continuous supercomputer time.
• The Real-World Use: This allows scientists to instantly scan vast "chemical design spaces" to find rare, high-value materials (like a better battery or a new drug) that would have been impossible to find before. It turns a needle-in-a-haystack problem into finding the needle in a haystack in the blink of an eye.

6. Fine-Tuning: The "Apprentice" System

The paper also showed that this giant model can be easily adapted to specific tasks with very little new data.

• The Analogy: Normally, to learn a new skill, you have to start from zero. But because this model has already learned the "basics of the universe," it can become an expert in a new field (like predicting the strength of a specific new alloy) just by looking at a few examples. It's like a master chef who can instantly learn to bake a new type of cake after tasting it once, rather than needing to read a whole new cookbook.

Summary

This paper isn't just about making a bigger AI; it's about making AI practical for science. By combining a "team of specialists" approach, massive supercomputing power, and smart data management, they created a tool that turns the impossible task of exploring the entire universe of materials into a routine, 50-second job. This accelerates the discovery of everything from clean energy solutions to new medicines.

Technical Summary

1. Problem Statement

The discovery of new materials is bottlenecked by the high computational cost of first-principles methods like Density Functional Theory (DFT), which are too slow for screening vast chemical spaces (billions of candidates). While Machine Learning Interatomic Potentials (MLIPs) and Graph Foundation Models (GFMs) offer a faster alternative, current approaches face three critical challenges:

• Data Heterogeneity & Imbalance: No single dataset covers all chemical domains (organic, inorganic, hybrid) or fidelity levels. Jointly training on diverse, imbalanced datasets often leads to optimization instability, where the model overfits dominant datasets and ignores scientifically valuable minority data.
• Architecture Selection: No single neural network backbone (e.g., MACE, PaiNN, SchNet) is universally superior. Selecting the optimal architecture requires massive hyperparameter optimization (HPO) that balances accuracy, training time, and resource efficiency at the exascale.
• Scalability & Portability: Training on hundreds of millions of structures across heterogeneous supercomputers (Frontier, Aurora, Perlmutter) requires robust data pipelines and software that can scale without re-engineering.
• Precision vs. Performance: Scientific workflows often require high precision (FP64) for stability, but this conflicts with the throughput demands of billion-scale screening.

2. Methodology

The authors present a holistic exascale workflow built on HydraGNN, integrating advanced data management, multi-task learning, and large-scale hyperparameter optimization.

A. Data Pipeline & Infrastructure

• Datasets: The model was trained on 16 open-source first-principles datasets containing over 544 million atomistic structures covering 85+ elements. These include organic molecules, inorganic crystals, catalysts, and polymers, generated with varying DFT functionals and fidelities.
• Data Movement: The workflow utilizes ADIOS2 and DDStore (Distributed Data Store) to stage data on local NVMe storage. This reduces remote file system I/O latency and prevents bottlenecks during synchronized multi-node training epochs.
• Parallelism: Training employs sharded distributed execution and task-parallelism, where MPI ranks are allocated proportionally to dataset sizes to prevent idle resources on smaller data branches.

B. Multi-Task Learning (MTL) Architecture

• Shared Backbone: A shared message-passing neural network (MPNN) learns transferable atomistic interaction features common across all datasets.
• Dataset-Specific Heads: To handle imbalanced data and fidelity offsets, the architecture uses per-dataset output heads. This decouples dataset-specific supervision, reducing gradient interference and allowing the model to learn from minority datasets without being overwhelmed by dominant ones.
• Composition-Conditioned Routing: A lightweight MLP dynamically blends the outputs of the 16 decoding branches based on the chemical composition of the input structure, enabling the model to interpolate between training domains for novel chemical systems.

C. Exascale Hyperparameter Optimization (HPO)

• Scale: Six large-scale HPO campaigns were executed on the Frontier supercomputer using 16,384 GPUs (1,024 nodes per campaign).
• Search Space: The search covered six MPNN backbones (EGNN, SchNet, DimeNet, MACE, PaiNN, PNAEq) with varying depths, hidden dimensions, and interaction parameters.
• Selection Criteria: Models were selected based on a trade-off between validation loss, time-to-epoch, and computational cost, rather than accuracy alone.

D. Precision & Inference Optimization

• Precision Analysis: The study systematically evaluated BF16, FP32, and FP64. While lower precision increases throughput, the lead model was trained in FP64 to ensure scientific rigor and stability during fine-tuning.
• Inference Acceleration: Four key optimizations were applied to the inference pipeline:
  1. Encoder Reuse: Caching shared trunk computations.
  2. Branch Skipping: Skipping low-weight decoder branches based on composition.
  3. Fused Gradient: Combining 16 energy predictions into a single weighted sum before computing forces via a single autograd call.
  4. Kernel Fusion: Using torch.compile to fuse convolution blocks.

3. Key Contributions

1. First Exascale Multi-Task GFM: The first workflow to jointly train a foundation model on 544M+ structures from 16 heterogeneous datasets using 16,384 GPUs.
2. Novel MTL Strategy: A head-level specialization approach that stabilizes training on imbalanced, multi-fidelity data, outperforming standard mixture-of-experts (MoE) designs for cross-dataset heterogeneity.
3. Systematic HPO at Scale: Demonstrated that architecture selection must be an exascale optimization problem, identifying PaiNN as the optimal backbone for this specific workload due to its balance of accuracy and time-to-epoch.
4. Cross-System Portability: Proven seamless strong and weak scaling across three distinct leadership-class systems: Frontier (AMD MI250X), Aurora (Intel XPU), and Perlmutter (NVIDIA A100).
5. Billion-Scale Screening: Achieved the evaluation of 1.1 billion atomistic structures in 50 seconds on Frontier, a task that would take ~6.7 years of continuous first-principles computation.

4. Key Results

• Training Performance:
  • The selected PaiNN-based lead model (12.1M parameters) achieved convergence at 2,048 nodes.
  • Near-linear strong scaling was observed up to 2,048 GPUs on Perlmutter, 6,144 on Aurora, and 1,024 on Frontier.
• Downstream Fine-Tuning:
  • Data Efficiency: Fine-tuning the pre-trained model on downstream tasks (e.g., QM9, MD17, OQMD) reduced validation errors by up to one order of magnitude compared to training from scratch, even in low-data regimes (e.g., 150 structures).
  • Task Specificity: Unfrozen fine-tuning (updating all weights) significantly outperformed frozen backbones for Potential Energy Surface (PES) tasks but showed diminishing returns for non-PES classification tasks.
• Inference Throughput:
  • Optimized inference achieved 21.8 million structures/second system-wide (293 structures/s/GPU).
  • Precision Trade-offs: FP32 introduced a negligible energy deviation ($|\Delta E| \approx 0.021$ eV), while BF16 showed larger deviations ($\approx 0.080$ eV), confirming FP64 is necessary for high-fidelity scientific deployment.

5. Significance

This work transforms materials discovery from a combinatorial bottleneck into a tractable workflow. By demonstrating that exascale-trained foundation models can be:

1. Robustly trained on massive, imbalanced, multi-fidelity data.
2. Portably deployed across heterogeneous supercomputing architectures.
3. Efficiently screened at the billion-structure level in seconds.

The authors establish a new paradigm where "foundation models" are not just large AI artifacts but practical, high-throughput scientific tools. This enables the exploration of vast chemical design spaces that were previously inaccessible to first-principles methods, accelerating the discovery of new materials for energy, catalysis, and electronics. The work also highlights the critical importance of precision management and system-level co-design (data pipelines, memory staging) in achieving true exascale scientific AI.