CDFGNN:一种基于缓存的分布式全批次图神经网络训练系统设计与通信优化
分布式、并行与集群计算
2024-08-02 v1 机器学习
摘要
图神经网络训练主要分为小批量和全批量训练方法。小批量训练方法在每个迭代中从原始图中对子图进行采样。这种采样操作引入额外的计算开销并降低训练准确性。而全批量训练方法在每个迭代中计算所有顶点的特征及相应的梯度,因此具有更高的收敛准确性。然而,在分布式集群中,频繁的远程访问顶点特征和梯度会导致巨大的通信开销,从而限制了整体训练效率。本文引入基于缓存的分布式全批量图神经网络训练框架(CDFGNN)。我们提出了自适应缓存机制,通过缓存邻居顶点的历史特征和梯度来减少远程顶点访问。此外,我们进一步通过量化消息并设计用于分层通信架构的图分区算法来优化通信开销。实验表明,自适应缓存机制平均减少了 63.14% 的远程顶点访问。结合通信量化和分层 GP 算法,CDFGNN 在实验中比当前最先进的分布式全批量训练框架性能提升 30.39%。我们的结果表明,CDFGNN 在加速分布式全批量 GNN 训练任务方面具有巨大的潜力。
引用
@article{arxiv.2408.00232,
title = {CDFGNN: a Systematic Design of Cache-based Distributed Full-Batch Graph Neural Network Training with Communication Reduction},
author = {Shuai Zhang and Zite Jiang and Haihang You},
journal= {arXiv preprint arXiv:2408.00232},
year = {2024}
}