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

TL;DR

This paper introduces a multi-task parallelism approach for pre-training graph neural network models on large, multi-source atomistic datasets, improving scalability and transferability across diverse chemical data.

Contribution

It presents a novel multi-task parallelism method implemented in HydraGNN, enabling efficient training on multi-million structure datasets across supercomputers.

Findings

Efficient scaling on three major supercomputers.

Successful training on over 24 million structures.

Enhanced transferability of graph models.

Abstract

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.