[论文解读] Boosting Distributed Full-graph GNN Training with Asynchronous One-bit Communication

Meng Zhang, Qinghao Hu+3 more|arXiv (Cornell University)|Mar 2, 2023

Advanced Graph Neural Networks被引用 4

一句话总结

该论文提出Sylvie，一种用于分布式图神经网络（GNN）训练的框架，通过异步一比特通信显著降低GPU间通信开销。通过应用低比特量化与误差补偿的反量化，Sylvie在保持模型精度的同时，相较于基线方法实现最高达28.1倍的加速。

ABSTRACT

Training Graph Neural Networks (GNNs) on large graphs is challenging due to the conflict between the high memory demand and limited GPU memory. Recently, distributed full-graph GNN training has been widely adopted to tackle this problem. However, the substantial inter-GPU communication overhead can cause severe throughput degradation. Existing communication compression techniques mainly focus on traditional DNN training, whose bottleneck lies in synchronizing gradients and parameters. We find they do not work well in distributed GNN training as the barrier is the layer-wise communication of features during the forward pass & feature gradients during the backward pass. To this end, we propose an efficient distributed GNN training framework Sylvie, which employs one-bit quantization technique in GNNs and further pipelines the curtailed communication with computation to enormously shrink the overhead while maintaining the model quality. In detail, Sylvie provides a lightweight Low-bit Module to quantize the sent data and dequantize the received data back to full precision values in each layer. Additionally, we propose a Bounded Staleness Adaptor to control the introduced staleness to achieve further performance enhancement. We conduct theoretical convergence analysis and extensive experiments on various models & datasets to demonstrate Sylvie can considerably boost the training throughput by up to 28.1x.

研究动机与目标

• 解决分布式全图GNN训练中高GPU间通信开销的问题，该问题限制了可扩展性和吞吐量。
• 克服现有压缩技术的局限性，这些技术虽针对深度神经网络（DNNs）优化，但对GNNs效果不佳，原因在于GNNs中存在逐层的特征与梯度通信。
• 通过低比特量化与异步流水线执行，在大幅降低通信成本的同时保持高模型精度。
• 设计一个可泛化至多种GNN模型与数据集的系统，且不牺牲收敛性或性能。

提出的方法

• 引入低比特模块，将特征与梯度量化为1比特进行传输，并在接收端使用误差补偿重建方法恢复为全精度。
• 实现异步流水线（Sylvie-A），将通信与计算重叠，隐藏延迟并提升资源利用率。
• 提出有界延迟自适应模块，控制异步训练中过时特征与梯度的影响，提升收敛稳定性。
• 设计图引擎，将图在GPU间划分，并识别通信关键节点，以最小化冗余数据传输。
• 使用训练器模块，基于反量化后的数据执行前向与反向传播，确保在量化存在的情况下仍保持模型一致性。
• 仅对通信张量应用量化，保留模型核心部分的全精度计算。

实验结果

研究问题

• RQ1 一比特量化能否在不降低模型精度的前提下，有效减少分布式GNN训练中的通信开销？
• RQ2 结合低比特通信的异步流水线执行对GNN训练的吞吐量与收敛性有何影响？
• RQ3 特征与梯度的延迟在多大程度上影响模型收敛？是否可通过控制机制维持性能？
• RQ4 所提出的框架能否在不同GNN架构与真实世界图数据集上实现良好泛化？
• RQ5 在低比特GNN训练中，通信减少、计算开销与内存占用之间的权衡关系如何？

主要发现

• 与FP32基线相比，Sylvie在Reddit数据集上将通信开销降低89.8%，每轮训练时间减少84.2%。
• 在双服务器配置下，使用1比特量化，Sylvie在训练吞吐量上最高实现28.1倍加速，如图1所示。
• 将所有嵌入与梯度均量化为1比特会导致精度严重下降（如Ogbn-products上从78.86%降至70.6%），证明选择性量化至关重要。
• 在双服务器配置下，低比特模块仅占总训练时间的5.9%，表明计算开销可忽略不计。
• 有界延迟自适应模块在异步更新可能降低性能的场景中，显著改善了收敛性。
• 该框架在多种模型（GraphSAGE、GCN、GAT）与数据集（Reddit、Yelp、Ogbn-products）上均保持高精度，与全精度训练相比精度损失极小。