# Multi-task parallelism for robust pre-training of graph foundation models on multi-source, multi-fidelity atomistic modeling data

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 32 references that are in the map.
- First posted on arXiv 26 Jun 2025.

## Builds on (most cited first, 9 of 9)

- [Deep CORAL: Correlation Alignment for Deep Domain Adaptation](https://mlsanity.com/paper/1607.01719/) (2016)
- [AliGraph: A Comprehensive Graph Neural Network Platform](https://mlsanity.com/paper/1902.08730/) (2019)
- [DistGNN: Scalable Distributed Training for Large-Scale Graph Neural Networks](https://mlsanity.com/paper/2104.06700/) (2021)
- [Towards Foundational Models for Molecular Learning on Large-Scale Multi-Task Datasets](https://mlsanity.com/paper/2310.04292/) (2023)
- [From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction](https://mlsanity.com/paper/2310.16802/) (2023)
- [Transition1x -- a Dataset for Building Generalizable Reactive Machine Learning Potentials](https://mlsanity.com/paper/2207.12858/) (2022)
- [On the Scalability of GNNs for Molecular Graphs](https://mlsanity.com/paper/2404.11568/) (2024)
- [Distributed Hybrid CPU and GPU training for Graph Neural Networks on Billion-Scale Graphs](https://mlsanity.com/paper/2112.15345/) (2021)
- [Scalable Training of Trustworthy and Energy-Efficient Predictive Graph Foundation Models for Atomistic Materials Modeling: A Case Study with HydraGNN](https://mlsanity.com/paper/2406.12909/) (2024)

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

- [Out-of-Distribution Generalization in Graph Foundation Models](https://mlsanity.com/paper/2601.21067/) (2026)

## Abstract (arXiv)

Graph foundation models using graph neural networks promise sustainable, efficient atomistic modeling. To tackle challenges of processing multi-source, multi-fidelity data during pre-training, recent studies employ multi-task learning, in which shared message passing layers initially process input atomistic structures regardless of source, then route them to multiple decoding heads that predict data-specific outputs. This approach stabilizes pre-training and enhances a model's transferability to unexplored chemical regions. Preliminary results on approximately four million structures are encouraging, yet questions remain about generalizability to larger, more diverse datasets and scalability on supercomputers. We propose a multi-task parallelism method that distributes each head across computing resources with GPU acceleration. Implemented in the open-source HydraGNN architecture, our method was trained on over 24 million structures from five datasets and tested on the Perlmutter, Aurora, and Frontier supercomputers, demonstrating efficient scaling on all three highly heterogeneous super-computing architectures.