NeutronJittor: Efficient CPU-GPU memory management for GNN training with Jittor

Area of Science:

• Machine Learning
• Computer Science
• Graph Neural Networks

Background:

• Full-graph training in Graph Neural Networks (GNNs) offers accuracy but requires substantial GPU memory.
• Irregular graph data access complicates manual memory management, necessitating user-transparent solutions.
• Existing memory management frameworks struggle with GNN workloads due to redundant computations and inefficient data movement.

Purpose of the Study:

• To introduce NeutronJittor, an efficient GNN training system with user-transparent memory management.
• To address the inefficiencies of current memory management systems for GNNs.
• To improve the performance and reduce the memory footprint of GNN training.

Main Methods:

• Implemented an embedding-reuse execution strategy to reuse intermediate embeddings of shared vertices, reducing redundant computations.
• Developed a spatial-temporal-aware data scheduler (ST-Aware scheduler) to prioritize data movement between CPU and GPU.
• Evaluated NeutronJittor using GCN and GAT models on real-world datasets.

Main Results:

• NeutronJittor significantly reduced communication volume by 65.7%-94.3%.
• Achieved substantial speedups ranging from 3.71x to 11.61x compared to state-of-the-art systems.
• Maintained comparable model accuracy to existing methods.

Conclusions:

• NeutronJittor provides an efficient solution for GNN training with user-transparent memory management.
• The embedding-reuse strategy and ST-Aware scheduler effectively mitigate redundant computations and data movement.
• NeutronJittor offers a promising approach for scalable and performant GNN training.