Optimizing Graph Neural Network Training on Large Graphs

Abstract

Graphs can be used to represent many important classes of structured real-world data. For this reason, there has been an increase of research interest in various machine learning approaches to solve tasks such as link prediction and node property prediction. Graph Neural Network models demonstrate good performance on such tasks. However, the depth of the models and the size of the graphs they can be trained on is constrained either by the low processing throughput of CPUs or by the limited memory capacity of GPUs. Techniques such as neighborhood sampling are often used to create smaller mini-batch training examples that fit in GPU memory. In this thesis, I provide a systematic performance analysis of GNN training codes written using PyTorch Geometric, the most popular machine learning framework for GNNs. Through this performance analysis, I uncover significant performance bottlenecks related to neighborhood sampling and GPU data transfers. To address these issues, I create FastPyG: a performance-engineered fork of PyTorch Geometric, which achieves a 3-6× speedup over comparable PyTorch Geometric codes without impacting model accuracy. The core contribution included in FastPyG is fast_sampler, an efficient and parallel neighborhood sampling implementation in C++.