Ginex：通过可证明最优的内存缓存于单机上借助 SSD 实现十亿规模图神经网络训练
机器学习
2022-08-22 v1
摘要
近来，图神经网络（GNNs）作为能够有效服务图结构数据上各种推断任务的强大工具而备受关注。随着真实世界图的规模持续扩大，GNN 训练系统面临可扩展性挑战。分布式训练是通过扩展 CPU 节点来应对该挑战的流行方法。然而，基于磁盘的 GNN 训练却未受太多关注，其可通过利用 NVMe SSD 等高性能存储设备以更具成本效益的方式扩展单节点系统。我们观察到，主存与磁盘间的数据移动是基于 SSD 的训练系统的主要瓶颈，且传统 GNN 训练流水线在未考虑此开销时是次优的。因此，我们提出 Ginex，首个基于 SSD 的 GNN 训练系统，可在单台机器上处理十亿规模图数据集。受编译器优化中检查者-执行者执行模型的启发，Ginex 通过分离采样与聚集阶段来重构 GNN 训练流水线。该分离使 Ginex 能实现一种可证明最优的替换算法（即 Belady 算法）用于内存中特征向量的缓存，而特征向量占据了 I/O 访问的主要部分。根据我们在四个十亿规模图数据集上的评估，Ginex 的平均训练吞吐量比 SSD 扩展的 PyTorch Geometric 高 2.11 倍（最高达 2.67 倍）。
引用
@article{arxiv.2208.09151,
title = {Ginex: SSD-enabled Billion-scale Graph Neural Network Training on a Single Machine via Provably Optimal In-memory Caching},
author = {Yeonhong Park and Sunhong Min and Jae W. Lee},
journal= {arXiv preprint arXiv:2208.09151},
year = {2022}
}
备注
Published in 2022 International Conference on Very Large Databases (VLDB)