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

Graph Neural Networks (GNNs) have been widely adopted across domains, and full-graph training is often preferred for its accuracy and stability. However, this approach demands large memory footprints, requiring dynamic memory management during GPU execution. The highly irregular data access of graphs makes manual management impractical, motivating the need for user-transparent memory management. We observe that existing user-transparent memory management frameworks are inefficient when applied to GNN workloads. Their inefficiency stems from two fundamental sources: redundant computation caused by repeatedly processing shared vertices across subgraphs, and frequent CPU-GPU communication caused by suboptimal data movement strategies. To overcome these limitations, we introduce NeutronJittor, an efficient GNN training system with user-transparent memory management. Firstly, we design an embedding-reuse execution strategy that reuses the intermediate embeddings of shared vertices within an epoch, allowing later subgraphs to skip redundant computations and release temporary data sooner, reducing redundant computation and lowering GPU memory pressure. Secondly, we develop a spatial-temporal-aware data scheduler (ST-Aware scheduler) that ranks vertices by structural and temporal priority, enabling the system to decide which data should stay on GPU and reducing unnecessary CPU-GPU data movement. Extensive experiments on real-world datasets with GCN and GAT models show that NeutronJittorreduces communication volume by 65.7%-94.3% and achieves $3.71\times$ - $11.61\times$ speedup compared to state-of-the-art systems (DGL and NeutronStar), while maintaining comparable model accuracy.