NeutronOrch: Rethinking Sample-based GNN Training under CPU-GPU Heterogeneous Environments

TL;DR

NeutronOrch introduces a layer-based task orchestration system for CPU-GPU heterogeneous environments that improves GNN training efficiency by balancing resource utilization and reducing GPU memory load, achieving up to 11.51x speedup.

Contribution

NeutronOrch proposes a novel layer-based task orchestrating method that decouples training layers and balances CPU-GPU workloads for efficient GNN training.

Findings

Achieves up to 11.51x speedup over state-of-the-art systems.

Effectively balances CPU and GPU utilization during training.

Reduces GPU memory footprint significantly.

Abstract

Graph Neural Networks (GNNs) have demonstrated outstanding performance in various applications. Existing frameworks utilize CPU-GPU heterogeneous environments to train GNN models and integrate mini-batch and sampling techniques to overcome the GPU memory limitation. In CPU-GPU heterogeneous environments, we can divide sample-based GNN training into three steps: sample, gather, and train. Existing GNN systems use different task orchestrating methods to employ each step on CPU or GPU. After extensive experiments and analysis, we find that existing task orchestrating methods fail to fully utilize the heterogeneous resources, limited by inefficient CPU processing or GPU resource contention. In this paper, we propose NeutronOrch, a system for sample-based GNN training that incorporates a layer-based task orchestrating method and ensures balanced utilization of the CPU and GPU. NeutronOrch decouples the training process by layer and pushes down the training task of the bottom layer to the CPU. This significantly reduces the computational load and memory footprint of GPU training. To avoid inefficient CPU processing, NeutronOrch only offloads the training of frequently accessed vertices to the CPU and lets GPU reuse their embeddings with bounded staleness. Furthermore, NeutronOrch provides a fine-grained pipeline design for the layer-based task orchestrating method, fully overlapping different tasks on heterogeneous resources while strictly guaranteeing bounded staleness. The experimental results show that compared with the state-of-the-art GNN systems, NeutronOrch can achieve up to 11.51x performance speedup.