Systems for Graph Neural Networks

My group was involved in many projects on efficient systems for graph representation learning.

We developed NextDoor, the first system specifically designed to perform graph sampling on GPUs. NextDoor introduces a high-level API based on a novel paradigm for parallel graph sampling called transit-parallelism. We implement several graph sampling applications, and show that NextDoor runs them orders of magnitude faster than existing systems

Our follow-up work developed split-parallelism, a novel hybrid parallelism strategies for mini-batch training. The traditional data parallel approach poses inherent limitations to scalable GNN training because it introduces redundant data loading and computation. This project proposes a fundamental shift towards a new paradigm for GNN training called split parallelism. The main idea is to split the computation of mini-batches across multiple GPUs, which cooperatively perform each iteration. We also proposed probabilistic splitting algorithms to balance the load across splits and minimize communication cost. We implemented split parallelism in a system called GSplit.

We also performed and extensive experimental comparison of two common methods for training GNNs: mini-batch training and full-graph training. These two methods require different training pipelines and systems optimizations, resulting in two separate classes of GNN training systems. We provided a comprehensive empirical comparison of representative full-graph and mini-batch GNN training systems. We found that mini-batch training systems consistently converge faster than the full-graph training ones across multiple datasets, GNN models, and system configurations. Mini-batch training techniques converge to similar to or often higher accuracy values than full-graph training ones, showing that mini-batch sampling is not necessarily detrimental to accuracy.

Our latest work is on efficient training of memory-based Temporal Graph Neural Networks (M-TGNNs). We developed PRISM, an M-TGNN system that tackles memory staleness in TGNN training using multi-version memory and iterative memory refinement, improving model accuracy and training/inference efficiency without changing the model itself. Using PRISM, we could make existing models reach new state-of-the-art accuracy with high training efficiency.