[论文解读] PyTorch-Direct: Enabling GPU Centric Data Access for Very Large Graph Neural Network Training with Irregular Accesses

Seung Won Min, Wu Kun+6 more|arXiv (Cornell University)|Jan 19, 2021

Advanced Graph Neural Networks被引用 11

一句话总结

PyTorch-Direct 为大规模图神经网络（GNN）训练引入了一种以 GPU 为中心的数据访问范式，使 GPU 能够直接访问主机内存中不规则、分散的节点特征，而无需 CPU 参与。通过引入统一张量类型并优化内存访问模式，其平均数据传输时间减少 47.1%，GNN 训练速度最高提升 1.6 倍，同时系统功耗降低 12.4%–17.5%。

ABSTRACT

With the increasing adoption of graph neural networks (GNNs) in the machine learning community, GPUs have become an essential tool to accelerate GNN training. However, training GNNs on very large graphs that do not fit in GPU memory is still a challenging task. Unlike conventional neural networks, mini-batching input samples in GNNs requires complicated tasks such as traversing neighboring nodes and gathering their feature values. While this process accounts for a significant portion of the training time, we find existing GNN implementations using popular deep neural network (DNN) libraries such as PyTorch are limited to a CPU-centric approach for the entire data preparation step. This "all-in-CPU" approach has negative impact on the overall GNN training performance as it over-utilizes CPU resources and hinders GPU acceleration of GNN training. To overcome such limitations, we introduce PyTorch-Direct, which enables a GPU-centric data accessing paradigm for GNN training. In PyTorch-Direct, GPUs are capable of efficiently accessing complicated data structures in host memory directly without CPU intervention. Our microbenchmark and end-to-end GNN training results show that PyTorch-Direct reduces data transfer time by 47.1% on average and speeds up GNN training by up to 1.6x. Furthermore, by reducing CPU utilization, PyTorch-Direct also saves system power by 12.4% to 17.5% during training. To minimize programmer effort, we introduce a new "unified tensor" type along with necessary changes to the PyTorch memory allocator, dispatch logic, and placement rules. As a result, users need to change at most two lines of their PyTorch GNN training code for each tensor object to take advantage of PyTorch-Direct.

研究动机与目标

• 解决由 CPU 为中心的数据准备导致的大规模 GNN 训练性能瓶颈，该问题使 CPU 过载并限制 GPU 利用率。
• 克服现有深度学习框架（如 PyTorch）在处理不规则图数据访问时效率低下的问题，这些框架需要 CPU 协助收集数据后才能将数据传送到 GPU。
• 使 GPU 能够直接访问主机内存中分散存储的节点特征，以减少延迟和数据移动开销。
• 通过引入统一张量抽象和对 PyTorch 内存分配器及分发逻辑的向后兼容性修改，最小化程序员工作量。
• 在保持与现有 PyTorch GNN 代码库兼容的前提下，实现显著的端到端训练加速和系统功耗降低。

提出的方法

• 引入一种新型的“统一张量”类型，使 GPU 内核能够直接访问主机内存中的数据，而无需事先在 CPU 端收集数据。
• 修改 PyTorch 的内存分配器和内存放置规则，以支持对 GPU 可访问的主机内存，并实现适当的对齐与合并优化。
• 实现一个具备内存对齐感知能力的 GPU 内核，以减少访问主机内存中对齐不规则的数据结构时的延迟。
• 利用 GPU 的高并发能力，直接向主机内存中分散的数据发起内存访问，绕过 CPU 在数据收集中的参与。
• 通过最多每张量修改两行代码，将新范式无缝集成到现有 PyTorch 训练工作流中，实现最小代码改动。
• 通过基于对齐感知的内核分发和合并内存访问策略，优化内存访问模式，以提升带宽利用率。

实验结果

研究问题

• RQ1 GPU 是否能够无需 CPU 参与，高效访问主机内存中的不规则、分散数据，从而减少 GNN 训练中的数据传输开销？
• RQ2 与 CPU 为中心的方法相比，GPU 为中心的数据访问模型在多大程度上能够减少数据准备时间并提升端到端 GNN 训练性能？
• RQ3 PyTorch-Direct 在 GNN 训练过程中对系统功耗有何影响，尤其是在降低 CPU 利用率的情况下？
• RQ4 在现有 PyTorch GNN 代码库中采用 GPU 为中心范式所需的最小程序员工作量是多少？
• RQ5 内存对齐如何影响 GNN 工作负载中 GPU 对主机内存中不规则数据的直接访问性能？

主要发现

• 与基线 PyTorch 方法相比，PyTorch-Direct 在端到端 GNN 训练中平均将数据传输时间减少 47.1%。
• 该框架在 GNN 训练性能上最高实现 1.6 倍加速，不同数据集和架构下的实际加速比范围为 1.01× 至 1.45×。
• 由于 CPU 利用率降低，训练期间系统功耗降低了 12.4% 至 17.5%，该结果通过直接电能表测量得到验证。
• 微基准测试结果表明，PyTorch-Direct 平均性能比基线 PyTorch 方法高出 2.39 倍，且在多种系统配置下均实现了接近理想的带宽利用率。
• PyTorch-Direct 中的内存对齐优化使 2052 字节特征的性能相比朴素 GPU 访问方法提升 1.95 倍，证明了合并访问模式的有效性。
• 用户仅需每张量修改最多两行代码即可启用 PyTorch-Direct 的全部优势，确保向后兼容性并降低采用门槛。