Scalable and Consistent Graph Neural Networks for Distributed Mesh-based Data-driven Modeling

TL;DR

This paper introduces a scalable, consistent distributed graph neural network framework for mesh-based modeling, enabling efficient data-driven simulations on exascale supercomputers while maintaining arithmetic equivalence across partitions.

Contribution

It develops a novel distributed GNN methodology with a consistent message passing layer that ensures arithmetic equivalence across partitions, linked with mesh partitioning for scalable CFD modeling.

Findings

Achieved scalable GNN training on up to 1 billion nodes on Frontier supercomputer.

Demonstrated that consistency improves scalability and accuracy in mesh-based GNNs.

Integrated GNNs with NekRS CFD solver for a unified data-driven modeling workflow.

Abstract

This work develops a distributed graph neural network (GNN) methodology for mesh-based modeling applications using a consistent neural message passing layer. As the name implies, the focus is on enabling scalable operations that satisfy physical consistency via halo nodes at sub-graph boundaries. Here, consistency refers to the fact that a GNN trained and evaluated on one rank (one large graph) is arithmetically equivalent to evaluations on multiple ranks (a partitioned graph). This concept is demonstrated by interfacing GNNs with NekRS, a GPU-capable exascale CFD solver developed at Argonne National Laboratory. It is shown how the NekRS mesh partitioning can be linked to the distributed GNN training and inference routines, resulting in a scalable mesh-based data-driven modeling workflow. We study the impact of consistency on the scalability of mesh-based GNNs, demonstrating efficient scaling in consistent GNNs for up to O(1B) graph nodes on the Frontier exascale supercomputer.