[论文解读] MGG: Accelerating Graph Neural Networks with Fine-grained intra-kernel Communication-Computation Pipelining on Multi-GPU Platforms

Yuke Wang, Boyuan Feng+5 more|arXiv (Cornell University)|Sep 14, 2022

Advanced Graph Neural Networks被引用 7

一句话总结

MGG 提出了一种用于多 GPU 图神经网络（GNNs）的细粒度核内通信-计算流水线系统，通过 GNN 优化的流水线构建与 GPU 感知的映射策略，实现通信与计算的重叠。该系统在 ROC 等最先进系统的基础上实现了最高达 10.83× 的加速，通过动态优化流水线配置并减少 GPU 间通信瓶颈，显著提升了全图 GNN 推理速度。

ABSTRACT

The increasing size of input graphs for graph neural networks (GNNs) highlights the demand for using multi-GPU platforms. However, existing multi-GPU GNN systems optimize the computation and communication individually based on the conventional practice of scaling dense DNNs. For irregularly sparse and fine-grained GNN workloads, such solutions miss the opportunity to jointly schedule/optimize the computation and communication operations for high-performance delivery. To this end, we propose MGG, a novel system design to accelerate full-graph GNNs on multi-GPU platforms. The core of MGG is its novel dynamic software pipeline to facilitate fine-grained computation-communication overlapping within a GPU kernel. Specifically, MGG introduces GNN-tailored pipeline construction and GPU-aware pipeline mapping to facilitate workload balancing and operation overlapping. MGG also incorporates an intelligent runtime design with analytical modeling and optimization heuristics to dynamically improve the execution performance. Extensive evaluation reveals that MGG outperforms state-of-the-art full-graph GNN systems across various settings: on average 4.41X, 4.81X, and 10.83X faster than DGL, MGG-UVM, and ROC, respectively.

研究动机与目标

• 为解决多 GPU GNN 中由不规则、稀疏且细粒度的 GPU 间通信引起的性能瓶颈。
• 克服现有系统分别优化计算与通信的局限性，避免错失联合调度的机会。
• 通过在 GPU 内核内动态重叠通信与计算，实现高性能的 GNN 执行。
• 设计一种能够通过智能运行时优化适应多样化 GNN 工作负载与 GPU 硬件配置的系统。
• 通过利用高效、低开销的通信原语，将该方法推广至 GNN 之外的其他深度学习工作负载。

提出的方法

• MGG 引入一种动态软件流水线，实现单个 GPU 内核中通信与计算的细粒度重叠。
• 采用 GNN 优化的流水线构建方法，基于数据访问模式将粗粒度的 GNN 操作分解为细粒度、可重叠的阶段。
• 系统使用 GPU 感知的流水线映射策略，选择最优的线程、 warp 和 block 配置，以最大化 GPU 占用率并最小化延迟。
• 集成基于 NVSHMEM 的通信原语，支持跨 GPU 的低延迟、细粒度远程内存访问。
• 智能运行时引擎结合分析建模与优化启发式方法，动态调优流水线参数，如交错距离和 warp 到 block 的映射。
• 通过适配不同的通信与工作负载模式，系统支持多种图划分策略——包括基于局部性的、基于工作负载的以及基于学习的划分。

实验结果

研究问题

• RQ1 如何有效构建细粒度的核内流水线，以在不规则、稀疏的 GNN 工作负载中重叠通信与计算？
• RQ2 在多 GPU 平台上，针对 GNN 的最优流水线配置与 GPU 映射策略是什么？
• RQ3 系统如何动态适应不同的图结构与硬件配置以维持高性能？
• RQ4 通过智能运行时调优流水线参数，通信开销可降低到何种程度？
• RQ5 所提出的流水线方法能否推广至 GNN 之外的其他深度学习工作负载？

主要发现

• 在多种 GNN 工作负载下，MGG 相较于 DGL 实现平均 4.41× 加速，相较 MGG-UVM 实现 4.81× 加速，相较 ROC 实现高达 10.83× 加速。
• 在 Criteo Kaggle 数据集上，MGG 将 DLRM 中的嵌入查找时间从 315.27ms 降低至 119.66ms，提升 2.64×。
• 通过采用 GPU 主导的细粒度远程内存访问，而非依赖 NCCL，系统有效减少了冗余的 GPU 间通信流量。
• 通过分析建模实现的运行时优化，使系统能够动态适应工作负载与硬件特性，提升流水线效率。
• 该方法在其他深度学习工作负载中也表现出良好的可扩展性，如在 DLRM 中成功集成并取得显著性能提升。
• 使用 NVSHMEM 原语实现了低延迟、可扩展的通信，支持在多 GPU 间高效流水线处理。