BPB-GNN: Resource-Efficient Distributed GNN Training via Boundary-Aware Prefetching and Cost-Balanced Batch Construction

Abstract

Distributed graph neural network (GNN) training is essential for large-scale graph learning, but existing mini-batch frameworks still suffer from substantial inter-partition communication overhead and indirect, coarse-grained GPU memory control. In distributed mini-batch training, remote sampled neighbors cause frequent feature fetching, while fixed seed-node batch sizes cannot accurately reflect the computation graph expansion induced by multi-layer sampling. This paper presents BPB-GNN, a resource-efficient distributed GNN training system that jointly optimizes inter-partition feature access, hierarchical storage management, and mini-batch construction. First, BPB-GNN introduces a training-centric boundary-aware prefetching strategy to localize high-value remote neighbor features associated with boundary nodes in training-node receptive fields. Second, it uses a DRAM/NVMe SSD hierarchy with priority-aware placement and dynamic migration to manage prefetched features under limited DRAM capacity. Finally, it constructs seed-node batches according to estimated sampling-induced load costs rather than equal seed-node counts, thereby balancing batch-level workloads and reducing peak GPU memory usage and its inter-batch fluctuation. These resource-efficient mechanisms may also improve the efficiency of large-scale GNN computation in computing power networks and JointCloud environments. Experiments on real-world graph datasets show that BPB-GNN achieves up to 5× speedup over the baselines while reducing peak GPU memory usage by up to 36%.