Accelerating Graph Sampling in GNN Systems Through Accuracy-Aware Data Reuse
摘要
Graph Neural Networks (GNNs) have become fundamental tools for processing and analyzing graph-structured data. However, when handling large-scale graph data, GNNs face exponential computational and memory challenges. While sampling-based training frameworks alleviate this issue, the sampling process itself becomes a performance bottleneck. Although existing optimizations like feature caching strategies and Unified Memory Access (UMA) kernels enhance memory efficiency, they overlook the redundant computation problem in sampling: the repeated sampling of identical nodes across different layers leads to resource wastage (consuming an average overhead of 15%). To address redundant computation, prior work proposes a complete cross-layer reuse strategy, eliminating redundancy by restricting each node to be sampled only once per mini-batch and indiscriminately reusing results. However, our work shows this indiscriminate reuse will destroy vertex neighbor diversity and reduce training stability to a certain extent, thus degrading the accuracy of GNN. (Training stability refers to the smooth convergence of the loss function during optimization, avoiding oscillations or divergence that impede model convergence.) In this paper, we propose an accuracy-aware sampling framework that accelerates sampling while preserving model accuracy. Specifically, by dissecting two critical execution phases in GNN training—feature aggregation and backpropagation—we introduce dual-dimensional vertex selective metrics for accuracy-aware cross-layer sampling-result reuse: (1) Coordinately considering in-degree and neighbor label entropy to identify nodes with high neighbor diversity, whose reuse is blocked to preserve structural information, and (2) Considering gradient norm metric to identify nodes that significantly influence parameter update directions, whose reuse is blocked to ensure training stability. Merging them all, our approach achieves up to 2.74 \(\times \) speedup on DGL while keeping accuracy loss within 1.5%.