Part of Proceedings of Machine Learning and Systems 2 (MLSys 2020)
Zhihao Jia, Sina Lin, Mingyu Gao, Matei Zaharia, Alex Aiken
Graph neural networks (GNNs) have been demonstrated to be an effective model for learning tasks related to graph structured data. Different from classical deep neural networks which handle relatively small individual samples, GNNs process very large graphs, which must be partitioned and processed in a distributed manner. We present Roc, a distributed multi-GPU framework for fast GNN training and inference on graphs. Roc is up to 4.6x faster than existing GNN frameworks on a single machine, and can scale to multiple GPUs on multiple machines. This performance gain is mainly enabled by Roc's graph partitioning and memory management optimizations. Besides performance acceleration, the better scalability of Roc also enables the exploration of more sophisticated GNN architectures on large, real-world graphs. We demonstrate that a class of GNN architectures significantly deeper and larger than the typical two-layer models can achieve new state-of-the-art classification accuracy on the widely used Reddit dataset.