提出可扩展的分布式图神经网络，实现网格建模中物理一致性与高效并行计算。
Scalable and Consistent Graph Neural Networks for Distributed Mesh-based Data-driven Modeling
- 通过边界环节点设计一致的消息传递机制，保证分布式训练等价性。
- 在前沿超算上实现百亿级节点图神经网络的高效并行，性能稳定提升。
- 适用于大规模流体仿真等科学计算场景，尤其适合高性能计算环境。
本文针对基于网格的建模应用，提出一种分布式图神经网络（GNN）方法，核心是采用一致的神经消息传递层。通过在子图边界引入环形节点（halo nodes），确保在单个大型图上训练和评估的GNN，其数学运算结果等同于在多个分片图上进行分布式计算的结果，从而实现物理一致性。该方法与阿贡国家实验室开发的GPU支持的百亿亿次级流体动力学求解器NekRS无缝集成，将NekRS的网格划分方式与分布式GNN训练推理流程结合，构建了可扩展的网格数据驱动建模工作流。研究验证了该一致性机制对可扩展性的关键作用，在前沿超算系统上实现了高达O(1B)节点规模的高效扩展。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。