Academic Journal
NeutronJittor: Efficient CPU-GPU memory management for GNN training with Jittor.
| Τίτλος: | NeutronJittor: Efficient CPU-GPU memory management for GNN training with Jittor. |
|---|---|
| Συγγραφείς: | Lu, Shizhan1, Fu, Zhenbo1, Wang, Qiange1, Zhang, Yanfeng1 zhangyf@mail.neu.edu.cn, Xu, Yuehao1, Ai, Xin1, Wei, Zhewei2, Yu, Ge1 |
| Πηγή: | Fundamental Research. Jul2026, Vol. 6 Issue 4, p2602-2611. 10p. |
| Θεματικοί όροι: | *Graph neural networks, *Computer memory management, *Heterogeneous computing, *Machine learning, *Process optimization, *Graphics processing units |
| Περίληψη: | 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×–11.61× speedup compared to state-of-the-art systems (DGL and NeutronStar), while maintaining comparable model accuracy. [ABSTRACT FROM AUTHOR] |
| Βάση Δεδομένων: | Academic Search Index |
καταχωρήστε σχόλιο πρώτοι!