A Distributed Multi-GPU System for Large-Scale Node Embedding at Tencent
arXiv:2005.13789
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.
arXiv commentAccepted by Scalable Deep Learning over Parallel And Distributed Infrastructures, An IPDPS 2021 Workshop
References in corpus (8)
- Efficient Estimation of Word Representations in Vector Space
- Distributed Representations of Words and Phrases and their Compositionality
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
- PyTorch-BigGraph: A Large-scale Graph Embedding System
- GraphVite: A High-Performance CPU-GPU Hybrid System for Node Embedding
- DGL-KE: Training Knowledge Graph Embeddings at Scale
- Graph Embeddings at Scale