CDFGNN: a Systematic Design of Cache-based Distributed Full-Batch Graph Neural Network Training with Communication Reduction

TL;DR

This work tackles the high communication overhead in distributed full-batch GNN training by introducing CDFGNN, a cache-based framework that adaptively caches neighbor features and gradients, quantizes inter-node messages, and employs a hierarchical vertex-cut graph partitioning strategy. The combined approach reduces remote data accesses and cross-node traffic while preserving convergence, with theoretical convergence guarantees under bounded caching error. Empirically, CDFGNN achieves a 63.14% average reduction in remote vertex accesses and a 30.39% improvement over state-of-the-art distributed full-batch frameworks across multiple large graphs. The results demonstrate practical potential for scalable, accurate full-batch GNN training on heterogeneous multi-GPU clusters, and future work plans to leverage NVLink for further speedups.

Abstract

Graph neural network training is mainly categorized into mini-batch and full-batch training methods. The mini-batch training method samples subgraphs from the original graph in each iteration. This sampling operation introduces extra computation overhead and reduces the training accuracy. Meanwhile, the full-batch training method calculates the features and corresponding gradients of all vertices in each iteration, and therefore has higher convergence accuracy. However, in the distributed cluster, frequent remote accesses of vertex features and gradients lead to huge communication overhead, thus restricting the overall training efficiency. In this paper, we introduce the cached-based distributed full-batch graph neural network training framework (CDFGNN). We propose the adaptive cache mechanism to reduce the remote vertex access by caching the historical features and gradients of neighbor vertices. Besides, we further optimize the communication overhead by quantifying the messages and designing the graph partition algorithm for the hierarchical communication architecture. Experiments show that the adaptive cache mechanism reduces remote vertex accesses by 63.14% on average. Combined with communication quantization and hierarchical GP algorithm, CDFGNN outperforms the state-of-the-art distributed full-batch training frameworks by 30.39% in our experiments. Our results indicate that CDFGNN has great potential in accelerating distributed full-batch GNN training tasks.

Figures (8)

• Figure 1: Distributed full-batch GNN Training.
• Figure 2: The sample process of mini-batch training.
• Figure 3: The workflow of CDFGNN.
• Figure 4: The communication pattern of CDFGNN.
• Figure 5: Comparison of average training time per epoch.

Theorems & Definitions (4)

• Lemma 1
• Lemma 2
• Lemma 3
• Theorem 1