[论文解读] Towards Efficient Large-Scale Graph Neural Network Computing

Lingxiao Ma, Zhi Yang+5 more|arXiv (Cornell University)|Oct 19, 2018

Advanced Graph Neural Networks参考文献 1被引用 11

一句话总结

NGra 是一种用于大规模图神经网络（GNNs）的新型并行框架，引入了 SAGA-NN 模型，将 GNN 表达为具有独立的 Scatter、ApplyEdge、Gather 和 ApplyVertex 阶段的顶点程序。通过自动图划分、基于数据块的流式处理以及基于环形的多 GPU 数据流机制，实现了高效、可扩展的 GPU 加速，相较于 TensorFlow 基线实现最高可提升 4 倍性能，并在多 GPU 上实现接近线性的可扩展性。

ABSTRACT

Recent deep learning models have moved beyond low-dimensional regular grids such as image, video, and speech, to high-dimensional graph-structured data, such as social networks, brain connections, and knowledge graphs. This evolution has led to large graph-based irregular and sparse models that go beyond what existing deep learning frameworks are designed for. Further, these models are not easily amenable to efficient, at scale, acceleration on parallel hardwares (e.g. GPUs). We introduce NGra, the first parallel processing framework for graph-based deep neural networks (GNNs). NGra presents a new SAGA-NN model for expressing deep neural networks as vertex programs with each layer in well-defined (Scatter, ApplyEdge, Gather, ApplyVertex) graph operation stages. This model not only allows GNNs to be expressed intuitively, but also facilitates the mapping to an efficient dataflow representation. NGra addresses the scalability challenge transparently through automatic graph partitioning and chunk-based stream processing out of GPU core or over multiple GPUs, which carefully considers data locality, data movement, and overlapping of parallel processing and data movement. NGra further achieves efficiency through highly optimized Scatter/Gather operators on GPUs despite its sparsity. Our evaluation shows that NGra scales to large real graphs that none of the existing frameworks can handle directly, while achieving up to about 4 times speedup even at small scales over the multiple-baseline design on TensorFlow.

研究动机与目标

• 解决现有深度学习与图处理框架中大规模图神经网络（GNNs）缺乏高效、可扩展的 GPU 支持的问题。
• 弥合基于数据流的深度学习框架与图结构神经网络计算之间的鸿沟。
• 在单个 GPU 内存无法容纳的超大规模图上实现端到端的 GNN 训练，利用主机内存和多 GPU 加速。
• 针对 GNN 中固有的不规则、稀疏数据访问模式，在 GPU 上优化稀疏图操作（Scatter/Gather）性能。

提出的方法

• 提出 SAGA-NN 模型，一种顶点程序抽象，将 GNN 表达为包含神经网络操作的 Scatter、ApplyEdge、Gather 和 ApplyVertex 阶段的序列。
• 将 SAGA-NN 模型映射为细粒度的数据流图，以块为单位，实现在单个或多个 GPU 上的细粒度、可扩展的流式处理。
• 采用自动图划分和基于数据块的处理方式，管理超出 GPU 内存容量的大规模图。
• 使用基于环形的流式传输机制，通过支持 GPU 间直接数据交换，最小化 GPU 间冗余数据移动。
• 通过利用每顶点数据访问的并行性并最小化内存延迟，优化 GPU 上的稀疏 Scatter 和 Gather 操作。
• 扩展 TensorFlow，添加自定义算子和顶点程序抽象，将 GNN 计算无缝集成到熟悉的深度学习生态系统中。

实验结果

研究问题

• RQ1 如何在大规模 GPU 加速系统上高效表达和执行图神经网络？
• RQ2 在 GPU 上处理 GNN 中固有的不规则、稀疏数据访问模式，需要哪些系统级优化？
• RQ3 如何利用多 GPU 和主机内存高效处理超出 GPU 内存容量的大规模图？
• RQ4 数据移动和调度策略对 GPU 上 GNN 训练性能的影响是什么？
• RQ5 结合数据流与顶点程序抽象的混合模型能否在 GNN 中同时实现表达力与高性能？

主要发现

• NGra 在小规模 GNN 工作负载上相较于多个 TensorFlow 基线实现最高提升 4 倍性能，展现出显著的性能优势。
• 基于环形的多 GPU 流式传输机制在 GPU 数量从 1 增加到 2 时实现 1.93 倍性能提升，而传统非环形机制仅实现 1.06 倍提升。
• NGra 支持训练包含多达数亿条边和数百万个顶点的图，其扩展能力超越现有框架的限制。
• 在多 GPU 环境中，系统在达到 NUMA 节点边界前表现出接近线性的可扩展性，通过手动 NUMA 感知张量分配可实现 7.09 倍性能提升。
• NGra 在 GPU 上优化的 Scatter 和 Gather 操作显著提升了内存访问效率，即使在 GNN 计算高度稀疏的情况下亦然。
• SAGA-NN 与 TensorFlow 的集成实现了 GNN 的无缝编程，同时通过系统级优化保持了高性能。