DiskGNN: Bridging I/O Efficiency and Model Accuracy for Out-of-Core GNN Training

TL;DR

DiskGNN tackles the challenge of training GNNs on graphs that do not fit in CPU memory by introducing offline sampling to decouple graph sampling from model computation and by optimizing data layout. It employs a four-level feature store, segmented disk cache with HashOrder-based node reordering, batched feature packing, and a pipelined training workflow to minimize disk I/O while preserving accuracy. Empirical results show DiskGNN achieves up to 8x end-to-end speedups over state-of-the-art out-of-core baselines while matching Ginex’s accuracy and outperforming MariusGNN in accuracy and efficiency. The approach meaningfully advances practical out-of-core GNN training on single machines and under cloud environments by leveraging memory hierarchy and proactive data layout optimization.

Abstract

Graph neural networks (GNNs) are machine learning models specialized for graph data and widely used in many applications. To train GNNs on large graphs that exceed CPU memory, several systems store data on disk and conduct out-of-core processing. However, these systems suffer from either read amplification when reading node features that are usually smaller than a disk page or degraded model accuracy by treating the graph as disconnected partitions. To close this gap, we build a system called DiskGNN, which achieves high I/O efficiency and thus fast training without hurting model accuracy. The key technique used by DiskGNN is offline sampling, which helps decouple graph sampling from model computation. In particular, by conducting graph sampling beforehand, DiskGNN acquires the node features that will be accessed by model computation, and such information is utilized to pack the target node features contiguously on disk to avoid read amplification. Besides, \name{} also adopts designs including four-level feature store to fully utilize the memory hierarchy to cache node features and reduce disk access, batched packing to accelerate the feature packing process, and pipelined training to overlap disk access with other operations. We compare DiskGNN with Ginex and MariusGNN, which are state-of-the-art systems for out-of-core GNN training. The results show that DiskGNN can speed up the baselines by over 8x while matching their best model accuracy.

Figures (19)

• Figure 1: An illustration for node-wise graph sampling. The seed node is $v_0$, and the sampled 1-hop and 2-hop neighbors are marked in yellow and green, respectively. We assume two node features take up a disk page.
• Figure 2: DiskGNN system architecture and workflow.
• Figure 3: An example of feature assembling in DiskGNN. $\mathsf{G\#}$, $\mathsf{C\#}$, $\mathsf{D\#}$, $\mathsf{P\#}$, and $\mathsf{PI\#}$ denote that a node feature locates in GPU cache, CPU cache, disk cache, packed chunk, and partial input, respectively. $\mathsf{PG}$ refers to disk pages in the disk cache.
• Figure 4: Feature access pattern before and after reordering. Minhash is used to reorder features in disk cache.
• Figure 5: Effect of node reordering for global and segmented disk cache on Friendster. (a) I/O amplification for an epoch, and (b) the ratio of I/O over No-reorder for each mini-batch (ranked by the I/O ratio). A segment has 50 min-batches.