Accelerating Storage-Based Training for Graph Neural Networks

AI summary

Overview

Research area: Graph neural network (GNN) systems — specifically storage-based (out-of-core) GNN training that keeps graph topology and node features on NVMe SSDs and streams only what is needed into main memory and GPU.

Technical level: Advanced. The paper assumes familiarity with GNN training loops (minibatching, k-hop sampling, forward/backward propagation), storage hierarchies, I/O bandwidth, and cache/buffer management.

Scope (one sentence): The paper identifies a large number of small storage I/Os as the dominant bottleneck in storage-based GNN training and proposes AGNES, a framework combining block-wise storage I/O processing with hyperbatch-based processing to eliminate that bottleneck.

What This Paper Is About

Training GNNs on web-scale graphs no longer fits in GPU memory (e.g., 80GB for an NVIDIA H100) or even in a single machine's main memory (e.g., 256GB) — the paper notes that training a 3-layer GAT on the yahoo-web graph (1.4B nodes, 6.6B edges) requires about 1.5TB. Storage-based training addresses this by keeping the full graph on external storage, but the authors observe that the data preparation stage (finding neighbors, gathering features, moving them to the GPU) dominates execution — up to 96% of total time — because real-world graphs have power-law degree distributions that cause a flood of tiny storage reads.

The goal is to make each storage I/O large and useful enough to saturate NVMe SSD bandwidth, rather than merely improving cache hit ratios as prior work does.

Key Contributions

1. Observation: The authors identify and characterize an under-explored challenge in storage-based GNN training — how to handle a large number of small storage I/Os — and show via preliminary experiments that data preparation consumes up to 96% of total execution time.
2. Framework: They propose AGNES (Accelerating storage-based training for Graph NEural networkS), a 3-layer architecture with storage, in-memory, and operation layers, built around block-wise storage I/O processing and a new data layout that places data accessed together into the same or adjacent blocks.
3. Strategy: They introduce hyperbatch-based processing, a simple yet effective technique that extends the processing scope from the target nodes of a single minibatch to those of multiple minibatches, so that a loaded block is reused across many target nodes instead of being evicted and reloaded.
4. Evaluation: Across five real-world graphs, AGNES finishes training up to 4.1× faster than the best competitor and achieves up to 4.5× greater I/O bandwidth utilization. Code is publicly released at https://doi.org/10.5281/zenodo.18112452 and https://github.com/Bigdasgit/agnes-kdd26.

Main Findings

• Data preparation dominates: In preliminary experiments with Ginex and GNNDrive on twitter-2010, ogbn-papers100M, and com-friendster using GCN and GraphSAGE, the data preparation stage takes up to 96% of total execution time.
• Storage I/Os are overwhelmingly small: The distribution of storage I/O sizes shows that a large number of I/Os are small while only a few are very large, which degrades computing resource utilization (e.g., GPU utilization).
• Root cause identified as graph structure: Real-world graphs follow a power-law degree distribution, so most nodes have only a few edges; existing methods read a few nodes at a time on demand, generating many small I/Os instead of addressing the problem fundamentally.
• Simply enlarging the I/O unit backfires: On the PA dataset, increasing Ginex's storage I/O unit size from 4KB to 1024KB caused total storage I/Os to surpass 15TB while the cache hit ratio dropped below 0.06%.
• Hyperbatch processing cuts I/O counts dramatically: In the paper's illustrative example, hyperbatch-based processing reduces storage I/Os from 20 to 5.
• Consistent wins over four state-of-the-art methods: AGNES outperforms Ginex, GNNDrive, MariusGNN, and OUTRE across all datasets. In Setting 1 (32GB), it achieves up to a 3.1× speedup over the best competitor, Ginex; in Setting 2 (8GB), the gap widens to up to 4.1×.
• Larger gains under memory pressure: The improved margin in Setting 2 indicates AGNES handles storage I/O effectively even when memory is constrained relative to graph size.
• Comparison against a distributed system: AGNES is reported to achieve comparable training performance to the distributed method DistDGL while using only limited computing resources (the detailed numbers for this comparison are in Figure 7 and are truncated in the provided content).
• Object index overhead is negligible: The object index table occupies less than 0.01% of the size of the original graph and is pinned in main memory.
• Not all baselines can run all models: MariusGNN and OUTRE support only the GraphSAGE model (marked N.A), and several method/dataset combinations resulted in out-of-memory (O.O.M) or out-of-time (O.O.T, more than 48 hours) outcomes.

Methodology in Plain English

The design rests on two ideas that work together.

Block-wise storage I/O processing. Instead of reading individual nodes on demand, AGNES divides the graph topology and node features into fixed-size blocks (the paper uses 1MB). It keeps an object index table — storing only the first and last object index for each block, sorted by node ID — so it can map a node to its block quickly. To make this efficient, the authors adopt a data layout that assigns consecutive node IDs to nodes likely to be accessed together at the same or adjacent iterations, so a single block-wise read fetches data that will actually be used, and accesses become more sequential.

Hyperbatch-based processing. Because a block loaded for one target node usually contains neighbors of many other target nodes, processing minibatches one at a time wastes the loaded data and forces reloads. AGNES instead groups multiple minibatches into a "hyperbatch" (size 1024) and processes all target nodes in a hyperbatch that fall inside the same loaded block within a single iteration. It tracks this with a bucket matrix whose rows correspond to blocks and whose columns correspond to minibatches, so it can scan a row to identify exactly which nodes in a block belong to which minibatch. Combined with LRU-based dynamic caching that pins graph blocks until the current iteration finishes, and with asynchronous I/O so threads issue reads and move on to other work rather than idling, this hides I/O latency.

Feature vectors are handled separately: because they occupy far more space than topology, AGNES counts accesses per feature vector and caches only those above a threshold, writing the rest back to storage each minibatch. The pipeline repeats for k iterations (one per GNN layer), then gathers features into a contiguous memory region and transfers nodes and features to the GPU for forward and backward propagation.

Experiments use three 3-layer models (GCN, SAGE, GAT) with sampling size per layer (10, 10, 10), block size 1MB, minibatch size 1000, and hyperbatch size 1024, trained on five graphs: IGB-medium (10M nodes, 120M edges), twitter-2010 (41.65M nodes, 1.47B edges), ogbn-papers100M (111.06M nodes, 1.62B edges), com-friendster (68.35M nodes, 2.29B edges), and yahoo-web (1.4B nodes, 6.6B edges). Graph sizes range from 6GB (IG, |F|=128) to 1.4TB (YH, |F|=256). Each configuration is trained five times and averaged. Hardware is a Dell R750 with an NVIDIA A40 (48GB), two Intel Xeon Silver 4309Y CPUs (8 physical cores each), 128GB main memory, and PCIe Gen 4.0 NVMe SSDs each with a maximum I/O bandwidth of about 6.7GB/s, using 16 CPU threads.

Why This Matters

Impact on research. The paper reframes the storage-based GNN training problem: prior work optimized cache hit ratios on the assumption that reading fewer, smaller pieces is better, whereas this work shows that the number and size of I/Os is the binding constraint on SSDs. This motivates a shift in how out-of-core graph systems lay out and schedule data. Because the I/O unit is a block rather than a node, the approach also decouples performance from degree skew, which is the property that breaks naive on-demand loading.

Real-world applications:

• Social network analysis on graphs with billions of edges (e.g., twitter-2010, com-friendster) on a single commodity server rather than a cluster.
• Web-scale link prediction and page/content recommendation on graphs such as yahoo-web.
• Academic citation network tasks like node classification on ogbn-papers100M.
• Any production environment where the graph exceeds GPU memory but buying a multi-node distributed cluster is impractical or too costly.

Industry relevance. The results target a concrete cost trade-off: the paper positions storage-based single-machine training as an alternative to distributed training, which requires inter-machine communication overhead plus the cost and effort of maintaining high-performance clusters. A 4.1× speedup or 4.5× better I/O bandwidth utilization directly reduces the wall-clock time and hardware footprint of large graph model training, and the fact that AGNES works across three GNN architectures and five datasets suggests the gains are not tied to one model family.

Future Directions

• Extending beyond the tested scope: DistDGL supports multi-machine scaling while AGNES is single-machine; whether block-wise and hyperbatch processing compose with graph partitioning across machines is an open question the paper's presentation leaves to future work.
• Completing the evaluation questions: The paper states four evaluation questions (training performance, ablation study on hyperbatch-based processing, sensitivity to hyperparameter settings, and accuracy). The provided content reports EQ1 in detail; the ablation, sensitivity, and accuracy results are referenced (Section 4.3) but not included in the available text, so their outcomes are not reported here.
• Tuning the new knobs: Hyperbatch size, block size, and the feature-cache access threshold are all parameters introduced by this design. The sensitivity analysis is flagged but not shown, leaving the question of how to set them for an arbitrary graph unanswered in this excerpt.
• Handling graphs with different structural properties: The power-law degree distribution is the stated motivation; whether block-wise I/O and hyperbatch processing remain equally effective on more uniformly connected graphs is not addressed.

Target Audience

This paper is most valuable to systems researchers and engineers building or operating large-scale GNN training pipelines, particularly those working on out-of-core or single-machine training over datasets that exceed GPU and main memory. It also suits GNN practitioners who need to train on billion-edge graphs without a distributed cluster, and database/storage researchers interested in data layouts that align on-disk organization with graph access patterns. Readers without background in GNN sampling and storage I/O scheduling will likely need to consult the referenced prior work (Ginex, GNNDrive, MariusGNN, OUTRE) before the design trade-offs are fully legible.

Authors’ abstract

Graph neural networks (GNNs) have achieved breakthroughs in various real-world downstream tasks due to their powerful expressiveness. As the scale of real-world graphs has been continuously growing, a storage-based approach to GNN training has been studied, which leverages external storage (e.g., NVMe SSDs) to handle such web-scale graphs on a single machine. Although such storage-based GNN training methods have shown promising potential in large-scale GNN training, we observed that they suffer from a severe bottleneck in data preparation since they overlook a critical challenge: how to handle a large number of small storage I/Os. To address the challenge, in this paper, we propose a novel storage-based GNN training framework, named AGNES, that employs a method of block-wise storage I/O processing to fully utilize the I/O bandwidth of high-performance storage devices. Moreover, to further enhance the efficiency of each storage I/O, AGNES employs a simple yet effective strategy, hyperbatch-based processing based on the characteristics of real-world graphs. Comprehensive experiments on five real-world graphs reveal that AGNES consistently outperforms four state-of-the-art methods, by up to 4.1X faster than the best competitor. Our code is available at https://github.com/Bigdasgit/agnes-kdd26.