[论文解读] A Distributed Multi-GPU System for Large-Scale Node Embedding at Tencent

Wanjing Wei, Yangzihao Wang+3 more|arXiv (Cornell University)|May 28, 2020

Caching and Content Delivery参考文献 27被引用 10

一句话总结

该论文提出了一种高性能、基于分布式多GPU的系统，用于大规模节点嵌入，通过结合模型并行和数据并行，在通用GPU集群上实现了高效扩展。通过采用分层数据分区、流水线训练以及解耦的随机游走与嵌入训练阶段，该系统在使用40块V100 GPU的情况下，仅用3分钟即可完成一个包含3000亿条边的网络的一轮训练，相比最先进方法实现了5.9倍至14.4倍的加速，同时保持了具有竞争力的准确率。

ABSTRACT

Real-world node embedding applications often contain hundreds of billions of edges with high-dimension node features. Scaling node embedding systems to efficiently support these applications remains a challenging problem. In this paper we present a high-performance multi-GPU node embedding system. It uses model parallelism to split node embeddings onto each GPU's local parameter server, and data parallelism to train these embeddings on different edge samples in parallel. We propose a hierarchical data partitioning strategy and an embedding training pipeline to optimize both communication and memory usage on a GPU cluster. With the decoupled design of CPU tasks (random walk) and GPU tasks (embedding training), our system is highly flexible and can fully utilize all computing resources on a GPU cluster. Comparing with the current state-of-the-art multi-GPU single-node embedding system, our system achieves 5.9x-14.4x speedup on average with competitive or better accuracy on open datasets. Using 40 NVIDIA V100 GPUs on a network with almost three hundred billion edges and more than one billion nodes, our implementation requires only 3 minutes to finish one training epoch.

研究动机与目标

• 为解决在通用GPU集群上扩展节点嵌入训练以处理百亿节点、万亿边网络的挑战。
• 克服现有系统在多GPU可扩展性方面不足或未能高效利用GPU内存与通信带宽的局限性。
• 设计一种将基于随机游走的网络增强与嵌入训练解耦的系统，以实现独立优化。
• 在大规模节点嵌入的分布式多GPU训练中实现高吞吐量与高效的通信重叠。
• 在真实世界与公开基准数据集上同时实现高性能与具有竞争力的准确率。

提出的方法

• 系统采用混合模型并行与数据并行：通过模型并行将节点嵌入在GPU间分区，通过数据并行将训练样本分布到各GPU上。
• 提出一种分层数据分区策略，以最小化通信开销并提升GPU集群中的数据局部性。
• 采用流水线嵌入训练流水线，重叠计算与通信，以最大化GPU利用率与带宽。
• 将基于CPU的随机游走生成与基于GPU的嵌入训练解耦，支持各阶段独立优化。
• 框架支持本地与远程参数服务器，并使用高效的嵌入查找与更新操作，以减少内核启动开销。
• 采用负采样策略以减少网络传输，提升训练效率。

实验结果

研究问题

• RQ1 如何设计一种分布式多GPU系统，以高效扩展至具有数千亿条边的大规模节点嵌入任务？
• RQ2 哪些设计模式能够实现多GPU节点嵌入系统中的高通信效率与内存利用率？
• RQ3 如何将随机游走生成与嵌入训练解耦，以实现独立优化与更优的资源利用率？
• RQ4 流水线训练与分层数据分区在多大规模图上能在多大程度上提升吞吐量并减少训练时间？
• RQ5 基于通用GPU集群构建的系统是否能在速度与准确率上超越现有最先进多GPU框架？

主要发现

• 在公开数据集上，该系统相比当前最先进的一体化多GPU单节点嵌入系统，平均实现了5.9倍至14.4倍的加速。
• 在一个拥有近3000亿条边、超过十亿个节点的网络上，系统仅用40块NVIDIA V100 GPU在3分钟内即完成一轮训练。
• 在标准公开基准数据集上，该系统保持了与基线模型相当或更优的准确率。
• 流水线训练策略成功实现了计算与通信的重叠，显著提升了GPU利用率与系统吞吐量。
• 分层数据分区有效降低了通信开销并提升了数据局部性，实现了GPU集群间的高效扩展。
• 解耦的架构使随机游走与嵌入训练可独立优化，提升了系统的灵活性与性能。