Co-GNN: A Co-optimization Framework for Memory and Computation in Sampling-Based GNN Training
摘要
Graph Neural Networks (GNNs) have achieved remarkable success in handling unstructured graph data, with sampling-based GNN training enhancing scalability for large graphs. However, due to the irregularity of graph data and the complexity of GNN training, existing GNN systems struggle to efficiently utilize multi-core accelerators, particularly GPUs. From a memory access perspective, graph topology caching improves the scalability of GPU sampling on large-scale graphs. However, cache misses necessitate fetching topology data from host memory, which introduces latency and degrades sampling performance. From a computational perspective, the same vertices are often sampled across multiple GNN layers. However, their scattered distribution leads to poor data locality, which reduces the computational efficiency. To address these issues, we propose Co-GNN, a framework designed to accelerate GNN training by collaboratively optimizing memory access and computation to enhance sampling efficiency. Co-GNN introduces two key innovations. First, it implements a graph topology cache prefetching mechanism that asynchronously overlaps the memory access for missing topologies with the sampling computations, effectively hiding memory access latency. Second, it develops a locality-aware sampling schedule that identifies frequently sampled vertices across sampling layers and caches them in the GPU’s read-only memory to enhance data locality. We evaluate Co-GNN on several real-world graph datasets using representative GNN models. Experimental results demonstrate that Co-GNN significantly outperforms DGL (UVA) by an average of 3.32 \(\times \) . Compared to the advanced GPU sampling method, Co-GNN achieves speedup of 1.62 \(\times \) (from 1.07 \(\times \) ).