Malope: Memory-Aware and Locality-Preserved Graph Neural Network Training

Abstract

Training GNNs on large-scale graphs imposes significant memory constraints for storing substantial amounts of graph structures and node features. This often necessitates the use of memory extensions such as SSDs, leading to a memory hierarchy with disparities in capacity and access speed. Existing approaches focus on mitigating the read amplification of SSDs used as memory extensions to enhance overall performance. However, these methods fail to achieve optimal performance on heterogeneous memory architectures such as DRAM–NVM systems and overlook the efficient utilization of fast memory. In this paper, we propose Malope, an efficient memory-aware and locality-preserved GNN training framework designed for heterogeneous memory systems. First, Malope introduces a memory-aware graph partitioning strategy that preserves multi-hop connectivity and maximizes fast memory utilization. Second, Malope presents a novel locality-preserved GNN training mechanism that reorganizes mini-batches to enhance data locality, thereby improving fast memory hit rates and minimizing partition switching overhead. Additionally, Malope integrates pipelined GNN training and partition switching to minimize data transfer overhead under low bandwidth conditions. Lastly, Malope enables fine-grained model persistence, built on reorganized mini-batch training, for rapid failure recovery. Experimental results on large real-world datasets show that Malope significantly outperforms state-of-the-art GNN training frameworks, achieving an impressive average speedup of <inline-formula><tex-math notation="LaTeX">$1.51\times$</tex-math><alternatives><mml:math><mml:mrow><mml:mn>1</mml:mn><mml:mo>.</mml:mo><mml:mn>51</mml:mn><mml:mo>×</mml:mo></mml:mrow></mml:math><inline-graphic xlink:href="cheng-ieq1-3705364.gif"/></alternatives></inline-formula>.