Armada：大规模图神经网络分布式训练的内存高效方法
机器学习
2025-02-26 v1 分布式、并行与集群计算
摘要
我们研究了在跨机器划分的十亿级图上进行图神经网络 (GNN) 的分布式训练。高效训练依赖于 min-edge-cut 划分算法，这些算法由于 GNN 邻域采样而最小化跨机器通信。然而，min-edge-cut 划分在大型图上仍具有挑战：最先进 (SoTA) 的离线方法（如 METIS）有效，但需要比 GNN 训练本身高出多个数量级的内存和运行时间，而计算效率高的算法（如流式贪心方法）则会导致 edge cuts 增加。因此，在本工作中我们引入了 Armada，一个用于分布式 GNN 训练的全新端到端系统，其关键贡献是 GREM，即一种能够高效扩展到大型图的 novel min-edge-cut 划分算法。GREM 基于流式贪心方法，关键添加了一项：在流式过程中持续细化 prior vertex 分配，而不是在初始贪心选择后冻结。我们的理论分析和实验结果表明，这种细化对于最小化 edge cuts 至关重要，使 GREM 能够以 8-65 倍更少的内存和 8-46 倍更快的速度达到与 METIS 相当的划分质量。给定一个划分好的图，Armada 利用一种新的 disintegrated 架构进行分布式 GNN 训练，以进一步提高效率；我们发现即使在 common 云机器上，即使没有通信，GNN 邻域采样和特征加载也会成为训练的瓶颈。 disintegration 允许 Armada 独立分配这些操作的资源，确保昂贵的 GPU 在计算上保持饱和。我们对 Armada 与 SoTA 系统进行评估，用于分布式 GNN 训练，我们发现 disintegrated 架构导致运行时间提高最高可达 4.5 倍，成本降低最高可达 3.1 倍。
引用
@article{arxiv.2502.17846,
title = {Armada: Memory-Efficient Distributed Training of Large-Scale Graph Neural Networks},
author = {Roger Waleffe and Devesh Sarda and Jason Mohoney and Emmanouil-Vasileios Vlatakis-Gkaragkounis and Theodoros Rekatsinas and Shivaram Venkataraman},
journal= {arXiv preprint arXiv:2502.17846},
year = {2025}
}