Work

Sampling-based Distributed Training with Message Passing Neural Network

Year: 2024
Type: preprint
Abstract: In this study, we introduce a domain-decomposition-based distributed training and inference approach for message-passing neural networks (MPNN). Our objective is to address the challenge of scaling edge-based graph neural networks as the number of nodes increases. Through our distributed training approach, coupled with Nyström-approximation sampling techniques, we present a scalable graph neural network, referred to as DS-MPNN (D and S standing for distributed and sampled, respectively), capable of scaling up to $O(10^5)$ nodes. We validate our sampling and distributed training approach on two cases: (a) a Darcy flow dataset and (b) steady RANS simulations of 2-D airfoils, providing comparisons with both single-GPU implementation and node-based graph convolution networks (GCNs). The DS-MPNN model demonstrates comparable accuracy to single-GPU implementation, can accommodate a significantly larger number of nodes compared to the single-GPU variant (S-MPNN), and significantly outperforms the node-based GCN.
Authors: Kakka, Priyesh, Nidhan, Sheel, Ranade, Rishikesh, Pathak, Jay, MacArt, Jonathan F.
Language: English
Open Access status: green
OpenAlex ID: W4392182047
Cites:
Cited by: 1
Related to: 10