[论文解读] Legion: Automatically Pushing the Envelope of Multi-GPU System for Billion-Scale GNN Training

Jie Sun, Li Su+9 more|arXiv (Cornell University)|May 26, 2023

Caching and Content Delivery被引用 5

一句话总结

Legion 是一个用于大规模图神经网络训练的多GPU系统，可自动优化缓存使用。它采用感知NVLink的分层划分方法和基于热度的统一缓存机制，最大限度减少PCIe通信流量并提升吞吐量，在大规模图上相比最先进系统最高可实现4.32倍的加速。

ABSTRACT

Graph neural network(GNN) has been widely applied in real-world applications, such as product recommendation in e-commerce platforms and risk control in financial management systems. Several cache-based GNN systems have been built to accelerate GNN training in a single machine with multiple GPUs. However, these systems fail to train billion-scale graphs efficiently, which is a common challenge in the industry. In this work, we propose Legion, a system that automatically pushes the envelope of multi-GPU systems for accelerating billion-scale GNN training. First, we design a hierarchical graph partitioning mechanism that significantly improves the multi-GPU cache performance. Second, we build a unified multi-GPU cache that helps to minimize the PCIe traffic incurred by caching both graph topology and features with the highest hotness. Third, we develop an automatic caching management mechanism that adapts the multi-GPU cache plan according to the hardware specifications and various graphs to maximize the overall training throughput. Evaluations on various GNN models and multiple datasets show that Legion supports training billion-scale GNNs in a single machine and significantly outperforms the state-of-the-art cache-based systems on small graphs.

研究动机与目标

• 解决现有基于缓存的GNN系统在多GPU环境下的可扩展性问题，特别是针对百亿规模图的场景。
• 克服现有系统在缓存复制和粗粒度拓扑管理方面的局限性。
• 通过基于访问热度智能缓存图拓扑和特征，最大限度减少PCIe数据传输。
• 自动化缓存管理，适应多样化的硬件配置和图工作负载，无需用户调优。
• 在单机多GPU系统上实现高效、端到端的百亿规模图神经网络训练。

提出的方法

• 提出感知NVLink的分层图划分方法：首先以最小化边切割的方式划分图，将独占分区分配给NVLink集群，然后在每个集群内使用哈希划分将训练顶点映射到各个GPU。
• 设计基于热度的统一缓存机制，将特征和拓扑统一存储于以顶点为中心的数据结构中，优先缓存访问频率最高的数据。
• 在NVLink集群内的GPU之间实现增强型缓存共享，提升内存利用率并减少冗余数据存储。
• 实现自动缓存管理机制，根据硬件规格和图特征动态规划缓存分配。
• 在预采样阶段基于访问频率设计热度指标，指导缓存放置策略以最小化缓存未命中。
• 集成基于GPU的图采样与统一缓存机制，加速采样和特征提取阶段。

Figure 1: The workflow of 2-hop GraphSAGE training.

实验结果

研究问题

• RQ1 如何在不造成过度数据复制的前提下，提升百亿规模GNN训练在多GPU环境下的缓存可扩展性？
• RQ2 如何在多GPU环境下以统一且GPU高效的方式管理拓扑与特征缓存？
• RQ3 系统是否能自动生成最优缓存方案，而无需依赖基于硬件或图特性的手动调优？
• RQ4 基于NVLink拓扑的分层划分如何提升缓存性能并减轻PCIe带宽压力？
• RQ5 与现有系统相比，统一的、基于热度感知的缓存机制在大规模图上能将训练吞吐量提升多少？

主要发现

• Legion 支持在单台多GPU机器上实现百亿规模图的端到端训练，突破了先前系统在可扩展性方面的限制。
• 在中小型图上，该系统相比最先进基于缓存的GNN系统最高可实现4.32倍的加速，展现出卓越的吞吐性能。
• 在极低缓存比例下，Legion 仍保持高缓存命中率（如图9和图11所示），表明其缓存利用效率极高。
• 分层划分减少了缓存复制，同时保持了分区内数据局部性，显著提升了整体内存效率。
• 自动缓存管理机制消除了用户调优需求，能有效适应多样化的硬件和图工作负载。
• 统一缓存设计通过优先缓存高热度顶点，显著降低PCIe通信开销，有效减少数据传输量。

(a) 2 GPUs per NVLink clique