PyTorch-Direct：以 GPU 为中心的数据访问实现不规则访问下超大规模图神经网络训练
机器学习
2021-01-21 v1 性能
摘要
随着图神经网络（GNNs）在机器学习社区中的采用日益增加，GPU 已成为加速 GNN 训练的重要工具。然而，在无法放入 GPU 内存的超大图上训练 GNN 仍是一项具有挑战性的任务。与传统神经网络不同，GNN 中的小批量输入样本需要遍历相邻节点并收集其特征值等复杂任务。虽然此过程占训练时间的很大一部分，我们发现使用 PyTorch 等流行深度神经网络（DNN）库的现有 GNN 实现，在整个数据准备步骤上局限于以 CPU 为中心的方法。这种“全在 CPU”的方法对整体 GNN 训练性能有负面影响，因其过度利用 CPU 资源并阻碍了 GNN 训练的 GPU 加速。为克服此类限制，我们引入了 PyTorch-Direct，其实现了一种以 GPU 为中心的 GNN 训练数据访问范式。在 PyTorch-Direct 中，GPU 能够直接高效访问主机内存中的复杂数据结构而无需 CPU 干预。我们的微基准和端到端 GNN 训练结果表明，PyTorch-Direct 平均减少 47.1% 的数据传输时间，并将 GNN 训练加速至多 1.6 倍。此外，通过降低 CPU 利用率，PyTorch-Direct 在训练期间还节省了 12.4% 至 17.5% 的系统功耗。为最小化程序员工作量，我们引入了新的“统一张量”类型以及对 PyTorch 内存分配器、调度逻辑和放置规则的必要修改。因此，用户对于每个张量对象最多只需修改两行其 PyTorch GNN 训练代码即可利用 PyTorch-Direct。
引用
@article{arxiv.2101.07956,
title = {PyTorch-Direct: Enabling GPU Centric Data Access for Very Large Graph Neural Network Training with Irregular Accesses},
author = {Seung Won Min and Kun Wu and Sitao Huang and Mert Hidayetoğlu and Jinjun Xiong and Eiman Ebrahimi and Deming Chen and Wen-mei Hwu},
journal= {arXiv preprint arXiv:2101.07956},
year = {2021}
}