CNS Core: Small: Transparently Scaling Graph Neural Network Training to Large-Scale Models and Graphs

Large-scale graphs with billions of edges are ubiquitous in many industry, science, and engineering fields such as recommendation systems, social graph analysis, knowledge bases, materials science, and biology. In particular, Graph Neural Networks (GNN), an emerging class of machine learning (ML) models, are increasingly adopted due to their superior performance in many tasks.

Unfortunately, the progress towards training GNNs on large-scale real-world graphs is undermined by the lack of adequate systems support for ML practitioners. This project will develop fundamental research on algorithms, systems, and infrastructures to meet the pressing and growing need for GNN training systems that can scale to both large graph datasets and large expressive GNN models transparently to users.

First, this project will develop split parallelism, a novel parallel training paradigm designed to support arbitrarily large-scale graphs and GNN models by scaling out to distributed and multi-GPU (graphics processing unit) systems. Split parallelism is tailored to the specific bottlenecks of GNNs and introduces a set of techniques to transparently split the training computation across GPUs.

Second, this project will develop systems for scalable graph sampling, which can be a major performance bottleneck in GNN training. It will develop a novel fragment-based in-GPU sampling approach that transparently splits samples into multiple fragments to maximize data access locality and scalability.

Supporting large-scale graphs and GNN models will unleash innovation in a wide range of domains by making it easier for ML practitioners to develop large and expressive models without having to work around the scalability limitations of current GNN training systems. The project will develop novel approaches for parallel training and sampling and will introduce innovations in algorithms, infrastructure, and system design for the areas of general machine learning and graph analytics.

This project will stress technology transfer to integrate the findings into popular open-source GNN training tools such as the Deep Graph Library (DGL). The PIs will also support colleagues at their department working on question answering using knowledge graphs. The project will improve the training of both graduate and undergraduate students, emphasizing demographic diversity.

This award reflects NSF's statutory mission and has been deemed worthy of support through evaluation using the Foundation's intellectual merit and broader impacts review criteria.