面向分布式网格数据驱动建模的可扩展且一致的图神经网络
分布式、并行与集群计算
2024-10-03 v1 机器学习
计算物理
摘要
本工作发展了一种用于网格建模应用的分布式图神经网络（GNN）方法，采用一致的神经消息传递层。所谓一致性，指的是：在一个秩（一个大图）上训练和评估的 GNN，其计算结果与在多个秩（分区图）上的评估在算术上等价。通过在子图边界处引入边界节点（halo nodes）来实现可扩展操作并满足物理一致性。本文通过将 GNN 与美国能源部 Argonne 国家实验室开发的 GPU 能力的超大规模 CFD 求解器 NekRS 接口，演示了这一概念。结果表明，NekRS 的网格分区可链接到分布式 GNN 的训练和推理例程，从而形成可扩展的网格数据驱动建模工作流程。我们研究了一致性对网格 GNN 可扩展性的影响，结果表明在 Frontier 超大规模超算上，一致 GNN 可高效扩展至 O(10\textsuperscript{9}) 个图节点。
引用
@article{arxiv.2410.01657,
title = {Scalable and Consistent Graph Neural Networks for Distributed Mesh-based Data-driven Modeling},
author = {Shivam Barwey and Riccardo Balin and Bethany Lusch and Saumil Patel and Ramesh Balakrishnan and Pinaki Pal and Romit Maulik and Venkatram Vishwanath},
journal= {arXiv preprint arXiv:2410.01657},
year = {2024}
}