FreshGNN: Reducing Memory Access via Stable Historical Embeddings for Graph Neural Network Training

TL;DR

FreshGNN introduces a caching framework for GNN training that reduces memory access and accelerates training on large graphs while maintaining high accuracy by selectively reusing stable node embeddings.

Contribution

It proposes a novel historical cache policy for GNN training that balances re-computation and reuse to improve efficiency and accuracy.

Findings

Achieves up to 20.5x speedup on large datasets

Reduces memory access by 59%

Maintains less than 1% accuracy loss

Abstract

A key performance bottleneck when training graph neural network (GNN) models on large, real-world graphs is loading node features onto a GPU. Due to limited GPU memory, expensive data movement is necessary to facilitate the storage of these features on alternative devices with slower access (e.g. CPU memory). Moreover, the irregularity of graph structures contributes to poor data locality which further exacerbates the problem. Consequently, existing frameworks capable of efficiently training large GNN models usually incur a significant accuracy degradation because of the currently-available shortcuts involved. To address these limitations, we instead propose FreshGNN, a general-purpose GNN mini-batch training framework that leverages a historical cache for storing and reusing GNN node embeddings instead of re-computing them through fetching raw features at every iteration. Critical to its success, the corresponding cache policy is designed, using a combination of gradient-based and staleness criteria, to selectively screen those embeddings which are relatively stable and can be cached, from those that need to be re-computed to reduce estimation errors and subsequent downstream accuracy loss. When paired with complementary system enhancements to support this selective historical cache, FreshGNN is able to accelerate the training speed on large graph datasets such as ogbn-papers100M and MAG240M by 3.4x up to 20.5x and reduce the memory access by 59%, with less than 1% influence on test accuracy.