提出高效分布式图神经网络训练系统Armada,解决大规模图训练内存与通信瓶颈。
Armada: Memory-Efficient Distributed Training of Large-Scale Graph Neural Networks
- 引入GREM算法,在流式处理中持续优化顶点分配,减少跨机通信。
- 相比METIS,内存降低8-65倍、速度提升8-46倍,且分区质量相当。
- 采用解耦架构,分离计算与数据加载,使GPU利用率显著提升。
我们研究在跨多台机器划分的十亿级图上进行图神经网络(GNN)的分布式训练。高效训练依赖于最小边割划分算法,以减少因邻居采样导致的跨机通信。然而,大规模图上的最小边割划分仍具挑战:现有先进离线方法(如METIS)虽有效,但所需内存和运行时间比GNN训练本身高数个数量级;而计算高效的流式贪心算法则面临边割增加的问题。为此,本文提出Armada,一个端到端分布式GNN训练系统,其核心是新型最小边割划分算法GREM。GREM基于流式贪心方法,关键改进在于:在流式处理过程中持续优化先前的顶点分配,而非冻结初始贪心选择。理论分析与实验结果表明,这种优化对减少边割至关重要,使GREM达到与METIS相当的分区质量,同时内存降低8-65倍、速度提升8-46倍。在获得划分图后,Armada采用新型解耦架构进一步提升效率;实验发现,在常见云服务器上,即使零通信,邻居采样与特征加载仍会成为训练瓶颈。解耦设计允许独立分配资源,确保昂贵的GPU持续饱和计算。与现有最优系统对比,该架构带来最高4.5倍的运行时加速和最高3.1倍的成本降低。
原文摘要 · Abstract (English)
We study distributed training of Graph Neural Networks (GNNs) on billion-scale graphs that are partitioned across machines. Efficient training in this setting relies on min-edge-cut partitioning algorithms, which minimize cross-machine communication due to GNN neighborhood sampling. Yet, min-edge-cut partitioning over large graphs remains a challenge: State-of-the-art (SoTA) offline methods (e.g., METIS) are effective, but they require orders of magnitude more memory and runtime than GNN training itself, while computationally efficient algorithms (e.g., streaming greedy approaches) suffer from increased edge cuts. Thus, in this work we introduce Armada, a new end-to-end system for distributed GNN training whose key contribution is GREM, a novel min-edge-cut partitioning algorithm that can efficiently scale to large graphs. GREM builds on streaming greedy approaches with one key addition: prior vertex assignments are continuously refined during streaming, rather than frozen after an initial greedy selection. Our theoretical analysis and experimental results show that this refinement is critical to minimizing edge cuts and enables GREM to reach partition quality comparable to METIS but with 8-65x less memory and 8-46x faster. Given a partitioned graph, Armada leverages a new disaggregated architecture for distributed GNN training to further improve efficiency; we find that on common cloud machines, even with zero communication, GNN neighborhood sampling and feature loading bottleneck training. Disaggregation allows Armada to independently allocate resources for these operations and ensure that expensive GPUs remain saturated with computation. We evaluate Armada against SoTA systems for distributed GNN training and find that the disaggregated architecture leads to runtime improvements up to 4.5x and cost reductions up to 3.1x.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。