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

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 24 of its 37 references that are in the map.
- First posted on arXiv 2 Oct 2024.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Geometric deep learning: going beyond Euclidean data](https://mlsanity.com/paper/1611.08097/) (2016)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [Machine learning accelerated computational fluid dynamics](https://mlsanity.com/paper/2102.01010/) (2021)
- [Combining Differentiable PDE Solvers and Graph Neural Networks for Fluid Flow Prediction](https://mlsanity.com/paper/2007.04439/) (2020)
- [Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior](https://mlsanity.com/paper/2306.00258/) (2023)
- [Transformers for Modeling Physical Systems](https://mlsanity.com/paper/2010.03957/) (2020)
- [MultiScale MeshGraphNets](https://mlsanity.com/paper/2210.00612/) (2022)
- [Deep convolutional recurrent autoencoders for learning low-dimensional feature dynamics of fluid systems](https://mlsanity.com/paper/1808.01346/) (2018)
- [Simulating Continuum Mechanics with Multi-Scale Graph Neural Networks](https://mlsanity.com/paper/2106.04900/) (2021)
- [MAgNET: A Graph U-Net Architecture for Mesh-Based Simulations](https://mlsanity.com/paper/2211.00713/) (2022)
- [Multiscale Graph Neural Network Autoencoders for Interpretable Scientific Machine Learning](https://mlsanity.com/paper/2302.06186/) (2023)
- [Conditionally Parameterized, Discretization-Aware Neural Networks for Mesh-Based Modeling of Physical Systems](https://mlsanity.com/paper/2109.09510/) (2021)
- [A Finite Element-Inspired Hypergraph Neural Network: Application to Fluid Dynamics Simulations](https://mlsanity.com/paper/2212.14545/) (2022)
- [Multiscale graph neural networks with adaptive mesh refinement for accelerating mesh-based simulations](https://mlsanity.com/paper/2402.08863/) (2024)
- [Learning Closed-form Equations for Subgrid-scale Closures from High-fidelity Data: Promises and Challenges](https://mlsanity.com/paper/2306.05014/) (2023)
- [NekRS, a GPU-Accelerated Spectral Element Navier-Stokes Solver](https://mlsanity.com/paper/2104.05829/) (2021)
- [Importance of equivariant and invariant symmetries for fluid flow modeling](https://mlsanity.com/paper/2307.05486/) (2023)
- [Multi-GPU Approach for Training of Graph ML Models on large CFD Meshes](https://mlsanity.com/paper/2307.13592/) (2023)
- [In Situ Framework for Coupling Simulation and Machine Learning with Application to CFD](https://mlsanity.com/paper/2306.12900/) (2023)
- [Sampling-based Distributed Training with Message Passing Neural Network](https://mlsanity.com/paper/2402.15106/) (2024)
- [Combining Machine Learning with Computational Fluid Dynamics using OpenFOAM and SmartSim](https://mlsanity.com/paper/2402.16196/) (2024)

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

- [Learning Effective Dynamics across Spatio-Temporal Scales of Complex Flows](https://mlsanity.com/paper/2502.07990/) (2025)
- [A Mesh-Adaptive Hypergraph Neural Network for Unsteady Flow Around Oscillating and Rotating Structures](https://mlsanity.com/paper/2503.22252/) (2025)

## Abstract (arXiv)

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.