降低基于磁盘的图神经网络训练中的内存争用与 I/O 瓶颈
分布式、并行与集群计算
2024-06-21 v1 机器学习
摘要
图神经网络（GNN）日益受到欢迎。大规模图结构及其高维特征数据变得常见，对普通机器上的 GNN 训练变得不为人知。给定巨大的图结构，即使采用样本化方法的 GNN 训练也难以高效运行，因为在训练过程中难以将图结构的全部数据保留在内存中。利用固态驱动器（SSD）或其他存储设备扩展内存空间的训练方法已被广泛研究。因此，内存和 I/O 成为磁盘化训练的关键因素。我们发现，当前最先进的（SoTA）磁盘化 GNN 训练系统严重受制于诸多问题，如图结构与特征数据之间的内存争用，以及从 SSD 加载数据进行训练时的严重 I/O 瓶颈。为此，我们开发了 GNNDrive。GNNDrive 1）通过在抽样和特征提取之间进行整体性缓冲区管理，以最小化内存占用；2）通过异步特征提取策略，避免 I/O 瓶颈。它还规避了关键路径上的昂贵数据准备工作，充分利用软件和硬件资源。实验表明，GNNDrive 性能卓越。例如，在使用 Papers100M 数据集和 GraphSAGE 模型进行训练时，GNNDrive 相对于 SoTA 的 PyG+、Ginex 和 MariusGNN 分别快 16.9 倍、2.6 倍和 2.7 倍。
引用
@article{arxiv.2406.13984,
title = {Reducing Memory Contention and I/O Congestion for Disk-based GNN Training},
author = {Qisheng Jiang and Lei Jia and Chundong Wang},
journal= {arXiv preprint arXiv:2406.13984},
year = {2024}
}
备注
This is a full version for the paper with almost the same title accepted by the 53rd International Conference on Parallel Processing (ICPP 2024)