[论文解读] Accelerating Training and Inference of Graph Neural Networks with Fast Sampling and Pipelining

Tim Kaler, Nickolas Stathas+5 more|arXiv (Cornell University)|Oct 16, 2021

Advanced Graph Neural Networks参考文献 35被引用 21

一句话总结

该论文提出SALIENT系统，通过优化邻域采样、实现共享内存并行批量准备以及流水线化数据传输与GPU计算，加速图神经网络（GNN）训练与推理。在单张GPU上，相较于标准的PyTorch-Geometric，SALIENT实现3倍加速；在16张GPU上实现8倍并行加速，于ogbn-papers100M数据集上，训练一个3层GraphSAGE模型仅需每轮2.0秒，测试准确率达到64.58%。

ABSTRACT

Improving the training and inference performance of graph neural networks (GNNs) is faced with a challenge uncommon in general neural networks: creating mini-batches requires a lot of computation and data movement due to the exponential growth of multi-hop graph neighborhoods along network layers. Such a unique challenge gives rise to a diverse set of system design choices. We argue in favor of performing mini-batch training with neighborhood sampling in a distributed multi-GPU environment, under which we identify major performance bottlenecks hitherto under-explored by developers: mini-batch preparation and transfer. We present a sequence of improvements to mitigate these bottlenecks, including a performance-engineered neighborhood sampler, a shared-memory parallelization strategy, and the pipelining of batch transfer with GPU computation. We also conduct an empirical analysis that supports the use of sampling for inference, showing that test accuracies are not materially compromised. Such an observation unifies training and inference, simplifying model implementation. We report comprehensive experimental results with several benchmark data sets and GNN architectures, including a demonstration that, for the ogbn-papers100M data set, our system SALIENT achieves a speedup of 3x over a standard PyTorch-Geometric implementation with a single GPU and a further 8x parallel speedup with 16 GPUs. Therein, training a 3-layer GraphSAGE model with sampling fanout (15, 10, 5) takes 2.0 seconds per epoch and inference with fanout (20, 20, 20) takes 2.4 seconds, attaining test accuracy 64.58%.

研究动机与目标

• 为解决在分布式多GPU环境下，由于昂贵的迷你批量准备和数据传输导致的GNN训练与推理性能瓶颈。
• 识别并解决现有GNN框架（如PyTorch-Geometric和DGL）中尚未充分探索的瓶颈，特别是在批量处理期间的CPU利用率和内存带宽方面。
• 通过实证验证采样对精度影响极小，统一训练与推理的邻域采样方法，从而简化模型实现。
• 通过流水线化数据传输与计算，实现高GPU利用率，确保端到端训练时间仅受最慢组件限制。
• 构建一个与PyTorch和PyG兼容的通用系统，可轻松扩展至其他框架和大规模分布式环境。

提出的方法

• 设计一种面向CPU架构优化的性能增强型邻域采样器，降低多跳邻域扩展中的开销。
• 实现批量准备的共享内存并行化，以克服现有框架中CPU和内存带宽瓶颈。
• 引入CPU与GPU之间的流水线化数据传输，将数据移动与GNN计算重叠，提升GPU利用率。
• 使用PyTorch的DDP模块实现跨多台机器和GPU的分布式训练，支持通过极少代码修改实现水平扩展。
• 复用现有GNN架构与框架（如GraphSAGE），无需修改架构设计，确保向后兼容性与易用性。
• 通过实证评估采样对推理精度的影响，验证在训练与推理中统一使用采样的可行性。

实验结果

研究问题

• RQ1 在多GPU系统上使用邻域采样进行GNN训练与推理时，主要性能瓶颈是什么？
• RQ2 如何优化批量准备与数据传输，以减少其对整体训练吞吐量的影响？
• RQ3 邻域采样在多大程度上影响推理精度？是否可以在训练与推理中统一使用采样而不会造成显著精度损失？
• RQ4 将数据传输与GPU计算流水线化，能否在GNN工作负载中实现接近理想的GPU利用率？
• RQ5 所提出的系统在多张GPU和多台机器上的可扩展性如何？与标准实现相比，可达到的加速比是多少？

主要发现

• 在单张GPU上，SALIENT在ogbn-papers100M数据集上训练3层GraphSAGE模型时，相较标准PyTorch-Geometric实现3倍加速。
• 在16张GPU上，SALIENT实现8倍并行加速，相同模型与数据集的训练时间缩短至每轮2.0秒。
• 使用fanout (20, 20, 20)进行推理时，每批次耗时2.4秒，测试准确率达到64.58%，表明采样带来的精度损失极小。
• 批量准备、数据传输与GPU计算之间实现了近乎完美的重叠，使端到端训练时间几乎等于最慢组件的处理时间。
• 通过优化采样与共享内存并行化，显著降低了批量准备与传输环节的性能瓶颈，尤其在CPU与内存带宽方面。
• 所提出的优化具有通用性，可应用于其他GNN框架与分布式环境，未来还可与GPU级切片或数据缓存等技术进一步集成。