PDGNN: Efficient Micro-batch GNN Training via Degree-Pruned Partitioning and Redundancy Elimination

作者：Lizhi Zhang, Menghan Jia, Ping Gong, Zhiquan Lai, Dongsheng Li, Yongquan Fu, Ao Shen, Kai Lü · 发表于：ACM Transactions on Architecture and Code Optimization · 年份：2025 · DOI：10.1145/3767325 · 研究领域：Advanced Neural Network Applications、Brain Tumor Detection and Classification、Advanced Image and Video Retrieval Techniques

Graph Neural Networks (GNNs) have become essential for graph-based applications. Existing GNN training frameworks primarily adopt sampling-based mini-batch approach, which samples vertex neighborhoods to construct subgraphs for training iteratively, but faces a memory-accuracy tradeoff. To address this, naive micro-batch training with gradients accumulation can be applied, yet introduces severe redundancy from repeated neighbor sampling. Recent micro-batch training methods overcome this by adopting batch-level graph partitioning strategies that partition full-batch subgraph sampled from all target vertices to generate micro-batches. Nevertheless, these approaches struggle to balance partitioning overhead with redundancy. This article proposes PDGNN, a system designed for efficient micro-batch GNN training on single machine, reducing both training time and memory consumption. First, we propose a degree-pruning-based graph partitioning method that efficiently processes on simplified and sparsified multi-level bipartite graphs by removing high-degree vertices, significantly reducing partitioning overhead. Second, we propose a dependency-aware redundancy matching and elimination technique that maximizes inter-micro-batch redundancy matching, and eliminates redundant computation by embedding caching and reuse while maintaining model accuracy. Experiments on real-world datasets demonstrate that PDGNN substantially improves the efficiency of large-scale GNN training, achieving up to...