通过预取优化提升大规模图神经网络分布式训练效率
MassiveGNN: Efficient Training via Prefetching for Massively Connected Distributed Graphs
- 在DistDGL框架上设计参数化连续预取与淘汰策略
- 在NERSC Perlmutter超算上实现15%-40%训练加速
- 适合大规模图表示学习的分布式训练场景
图神经网络（GNN）在图结构数据学习中不可或缺，但其在高度连接图上的计算成本持续上升，严重制约执行性能。分布式内存方案通过分块图并行训练多个GNN副本成为主流，但此类方法常因邻域子批次采样不规则导致通信开销和负载不均，即便采用最优分块与通信策略也难以避免。本文针对基于GraphSAGE架构的分布式图表示学习，提出一种在先进Amazon DistDGL框架基础上的参数化连续预取与淘汰机制，有效缓解采样与通信开销，在国家能源研究科学计算中心（NERSC）的Perlmutter超算上对多个OGB数据集实现了约15%-40%的端到端训练性能提升。
原文摘要 · Abstract (English)
Graph Neural Networks (GNN) are indispensable in learning from graph-structured data, yet their rising computational costs, especially on massively connected graphs, pose significant challenges in terms of execution performance. To tackle this, distributed-memory solutions such as partitioning the graph to concurrently train multiple replicas of GNNs are in practice. However, approaches requiring a partitioned graph usually suffer from communication overhead and load imbalance, even under optimal partitioning and communication strategies due to irregularities in the neighborhood minibatch sampling. This paper proposes practical trade-offs for improving the sampling and communication overheads for representation learning on distributed graphs (using popular GraphSAGE architecture) by developing a parameterized continuous prefetch and eviction scheme on top of the state-of-the-art Amazon DistDGL distributed GNN framework, demonstrating about 15-40% improvement in end-to-end training performance on the National Energy Research Scientific Computing Center's (NERSC) Perlmutter supercomputer for various OGB datasets.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。