PyTorch-Direct: Enabling GPU Centric Data Access for Very Large Graph Neural Network Training with Irregular Accesses

With the increasing adoption of graph neural networks (GNNs) in the machine\nlearning community, GPUs have become an essential tool to accelerate GNN\ntraining. However, training GNNs on very large graphs that do not fit in GPU\nmemory is still a challenging task. Unlike conventional neural networks,\nmini-batching input samples in GNNs requires complicated tasks such as\ntraversing neighboring nodes and gathering their feature values. While this\nprocess accounts for a significant portion of the training time, we find\nexisting GNN implementations using popular deep neural network (DNN) libraries\nsuch as PyTorch are limited to a CPU-centric approach for the entire data\npreparation step. This "all-in-CPU" approach has negative impact on the overall\nGNN training performance as it over-utilizes CPU resources and hinders GPU\nacceleration of GNN training. To overcome such limitations, we introduce\nPyTorch-Direct, which enables a GPU-centric data accessing paradigm for GNN\ntraining. In PyTorch-Direct, GPUs are capable of efficiently accessing\ncomplicated data structures in host memory directly without CPU intervention.\nOur microbenchmark and end-to-end GNN training results show that PyTorch-Direct\nreduces data transfer time by 47.1% on average and speeds up GNN training by up\nto 1.6x. Furthermore, by reducing CPU utilization, PyTorch-Direct also saves\nsystem power by 12.4% to 17.5% during training. To minimize programmer effort,\nwe introduce a new "unified tensor" type along with necessary changes to the\nPyTorch memory allocator, dispatch logic, and placement rules. As a result,\nusers need to change at most two lines of their PyTorch GNN training code for\neach tensor object to take advantage of PyTorch-Direct.\n