LEO：通过分层缓存实现高效通信-计算流程，以支持图神经网络（GNN）的应用

《ACM Transactions on Architecture and Code Optimization》：LEO: Enabling Efficient Communication-Computation Pipeline for GNN via Hierarchical Caching

　　摘要人工智能总结：要查看此人工智能生成的摘要，您必须具有高级访问权限。了解更多信息，请登录。

摘要：在多GPU平台上使用GNN训练大规模图面临巨大的特征加载开销，导致资源利用率低和训练效率低下。克服这种通信瓶颈至关重要，当前的解决方案通过流水线化实现计算和通信的重叠来达到这一目标。然而，现有设计通常只关注有限的操作调度，忽略了流水线内的工作负载不平衡问题，从而限制了其在异构环境中的性能优化。本文介绍了一种名为LEO的新系统，用于加速多个GPU上的全图GNN训练。LEO的核心在于引入分层缓存来构建混合粒度的软件流水线，从而实现GPU内核内部和之间的高效计算-通信重叠。首先，LEO引入了一种分层缓存策略，有效平衡了本地和远程操作之间的工作负载。其次，LEO量身定制了基于缓存的流水线和混合工作负载映射，以促进操作重叠。广泛的评估表明，LEO在各种设置下的性能都优于现有的全图GNN系统：其速度分别比DGL快7.50倍、比LEO-UVM快11.90倍、比MGG快3.25倍。

人工智能生成摘要（实验性）：该摘要是使用自动化工具生成的，并非由文章作者撰写或审核。它旨在支持发现、帮助读者评估相关性，并协助来自相关研究领域的读者理解这项工作。它旨在补充作者提供的摘要，后者仍然是论文的官方版本。点击此处了解更多信息。点击此处对摘要的准确性、清晰度和实用性进行评论。这些建议将有助于改进未来的摘要版本。要查看此人工智能生成的简单语言摘要，您必须具有高级访问权限。

1 引言：图神经网络（GNN）已成为处理基于图的深度学习任务（例如节点分类[18, 19]和链接预测[24, 28]）的有效解决方案，这些任务能够捕捉非欧几里得图结构数据中的丰富关系信息。为了处理下游图相关任务，GNN利用节点特征和图结构来学习图内节点的低维向量表示（嵌入）。GNN通常由多个层组成，每一层接收前层的输出并执行聚合和更新操作。由于图拓扑的高稀疏性和该操作的复杂性，聚合操作占据了GNN计算时间的绝大部分[38, 48]。与传统图分析方法[20, 39]相比，GNN具有更高的表达能力和泛化能力[27, 21]。在包含数百万节点和数十亿边的大规模图上训练GNN需要利用多GPU平台来扩展性能。GNN实践通常采用数据并行策略，利用现代多GPU平台提供的高内存容量和带宽（例如HBM和NVLink）来扩展GNN训练。然而，当数据并行性划分输入图时，跨分区节点会产生复杂且交错的数据依赖关系。这些依赖关系由输入图的拓扑决定，导致不规则且昂贵的GPU间通信，这是多GPU平台上GNN训练的关键瓶颈。现有工作试图通过使用定制的通信操作[37, 45, 52]和内存缓存[30, 41, 53]来减轻昂贵的不规则GNN通信。然而，前者在没有利用图工作负载特性的情况下只能带来有限的性能提升，而后者虽然针对基于样本的GNN训练进行了优化，但会导致精度损失[25, 47, 49]。同时执行通信和计算操作被认为是重叠昂贵的通信开销的有希望的方法。最近的工作MGG[49]旨在通过区分每个GPU上的本地/远程数据并实现细粒度流水线来增强计算和通信的重叠。然而，它由于两个原因而效率低下：（1）图拓扑在多个GPU上的分布导致本地/远程数据比例过高，远程数据访问比本地数据访问更耗时；（2）在训练期间CPU资源被闲置，限制了其潜在的性能提升。这些限制揭示了优化GNN训练中通信-计算重叠的两个关键挑战。第一个挑战涉及全面和系统地探索通信-计算重叠的整个优化空间。我们的关键见解是，GNN执行可以分解为分布在一组设备上的工作负载，形成流水线的基本调度单元，以鼓励更广泛的通信-计算重叠。我们将内存访问操作分为三个维度：（1）本地数据访问；（2）GPU间通信；（3）CPU到GPU的通信。例如，对于GPU上的一个图节点，可以在其他GPU上并发执行远程邻居访问和CPU上的邻居访问。此外，这些操作可以与本地邻居的内存访问和计算重叠。第二个挑战是如何有效调度不同类型的通信和计算操作以构建高效的流水线。为此，我们引入了一种利用混合粒度软件流水线来加速多GPU GNN训练的范式。我们将整个GNN工作负载分配到GPU和CPU上，并进一步将其分解为混合粒度操作。通过仔细编排工作负载并设计考虑GPU的映射方案，我们可以构建高效且灵活的流水线，以促进操作重叠。

在本文中，我们介绍了LEO，这是一个整体系统，旨在通过构建多阶段混合粒度流水线来促进全图GNN训练中的高效计算-通信重叠。LEO包含两个核心组件：运行时管理器和运行时内核（图1），它们分别负责构建缓存和创建针对GNN的流水线。在运行时管理系统中，我们提出了分层缓存部分（第4节）。通过缓存最热的数据，我们创建了一个更广泛的优化空间。考虑到可用的平台存储、带宽层次结构和热度统计信息，LEO可以使用成本模型有效地分配内存用于嵌入。为了进一步促进操作重叠，LEO构建适当的数据布局，以充分利用硬件资源。此外，LEO采用异步缓存更新来最大化缓存带来的好处。在运行时内核中，LEO构建了针对多GPU GNN加速的混合粒度软件流水线部分（第5节）。通过特定的工作负载编排，LEO将不规则的GNN工作负载分解为不同粒度的操作。此外，LEO为不同类型的选择适当的通信原语，以提供高性能的通信支持。最后，LEO提出了一种基于块的映射策略，将流水线中的操作映射到GPU上，以最大化内核效率。

图1. LEO概述。LEO探索了一个新的维度，从而扩展了整个优化空间。具体来说，我们的关键见解是GPU间通信延迟并非硬件的不变属性；它可以通过算法进行重塑。通过引入分层缓存作为可调维度，LEO将固定的远程/本地工作负载比例转变为可控变量。这一扩展的空间解锁了新的流水线设计和映射策略。因此，LEO不仅仅是改进了以往工作的调度粒度；相反，它通过同时优化数据放置和访问时机，重新定义了GNN流水线构建的基本问题。我们在本文中做出了以下贡献：
- 我们对现有解决方案中的性能问题进行了深入和全面的分析（第3节），旨在实现多GPU平台上的高效训练。
- 我们引入了针对流水线定制的分层缓存（第4节），以减轻昂贵的GPU间通信并提高CPU利用率。
- 我们提出了针对GNN的流水线构建和映射（第5节），通过编排和映射不同类型的操作来实现高效的计算-通信重叠。
- 广泛的实验表明，LEO的性能优于现有的多GPU GNN系统，验证了其有效性和泛化能力（第6节）。

2 背景：
2.1 图神经网络（GNN）：GNN是一种以图结构数据为输入的神经网络。图中的节点代表现实世界数据中的实体（例如社交网络中的用户），边表示节点之间的关系。每个节点与一个特征向量（嵌入）相关联，结合图结构，使GNN能够学习和理解数据中的关系。在GNN中，节点v在k层邻居聚合后的嵌入可以从第k-1层的节点嵌入（CNs）中得出，其中：
avk = Aggregate(k(huk?1|u∈N(v)∪hvk?1)
hvk = Update(k(avk)

（1）
这里，hvk表示第k层节点v的嵌入；N(v)表示节点v的邻居集，avk通过聚合邻居的嵌入获得中间结果。聚合函数主要收集邻居的特征。一些方法通过结合边权重和端点信息（例如图注意力网络（GAT）[44]）来利用边属性。更新函数使用一些经典的神经网络操作（例如均值、最大值和MLP）来更新节点的嵌入。
2.2 GNN训练策略：在处理包含数十亿节点和边的大规模图时，数据往往超出单个GPU的可用内存，因此需要将整个工作负载分布在多个GPU上。此外，现有实践采用两种训练策略来处理大规模图，包括基于样本的小批量训练和全图训练。
基于样本的小批量GNN训练：基于样本的训练包括三个主要步骤：采样、提取和训练。在采样阶段，对于每个训练节点，根据特定的采样算法[20, 21, 23, 39]对k跳邻居进行采样。接下来，GNN模型根据采样的子图提取相应的特征嵌入，这些子图由高维特征向量组成。与图拓扑相比，嵌入的体积更大，从而导致昂贵的存储和传输开销。最后，在训练阶段，将采样的子图及其相关特征作为输入来训练k层GNN模型。基于样本的策略可以有效减轻GNN训练中的计算和内存压力。然而，它不可避免地会带来一定的精度损失[12, 25]。
全图GNN训练：与基于样本的GNN训练不同，全图训练保留了整个图拓扑，使得可以在输入图的完整嵌入上进行学习。全图训练将整个图拓扑分布在多个GPU上，并在每一层交替进行特征提取和训练过程。在逐层训练过程中，每个GPU需要接收分区子图中节点的1跳邻居的消息，并从远程GPU获取相关数据。全图训练要求GPU之间频繁的数据交换，这带来了昂贵的通信开销。此外，由于每个图分区需要不同量的节点信息，不规则的P2P通信会进一步加剧通信瓶颈。
2.3 多GPU GNN训练加速：在处理大规模数据时，现有的深度学习应用通常将模型扩展到多GPU平台，并探索新的优化和设计以提高性能。针对结构化数据的深度学习应用（如DNN）已经探讨了各种方面的可扩展性，包括并行策略（例如数据和模型并行性）、高性能通信库设计（例如NCCL [7]）和操作符设计（例如cuDNN [2]）。然而，针对非结构化数据的GNN由于数据的稀疏性和不可预测性，无法简单地迁移这些方法。为此，现有应用开发了一系列针对GNN的系统优化，主要集中在三个方向：
操作符定制用于数据传输：这种解决方案主要关注为稀疏GNN通信定制操作符以加速通信（图2(a)），旨在加速异构环境中的GNN训练通信。DGL [45]是当前的GNN框架，它利用统一虚拟寻址（UVA）[10]来简化大规模GNN训练中的内存管理。UVA将图数据固定在中处理器内存中，并使GPU内核能够基于零拷贝内存直接访问特征[10]。WhloeGraph [52]利用GPUDirect Peer-to-Peer（P2P）[6]和CUDA IPC技术，使不同的GPU在一个CUDA内核中访问内存。来自[37]的工作利用BaM [40]使GPU线程可以直接从存储中获取特征嵌入。然而，这些解决方案仅关注加速通信，忽略了通过改善内存资源利用来减少通信的潜力，而这可以有效地加速特征提取。

图2. 不同的多GPU GNN训练加速方法：在每个子图中，y轴左侧的框代表存储在不同设备上的节点，而右侧的框代表各种类型工作负载的计算/通信操作。
缓存优化用于消除通信：第二种解决方案将远程数据缓存到本地GPU上，以避免潜在的通信（图2(b)）。它们为基于样本的小批量GNN训练定制了缓存策略，消除了采样和特征提取期间的数据移动。PaGraph [30]和GNNLab [53]分别开发了基于度和预采样的策略，在每个GPU中缓存最热的嵌入。Quiver [42] 利用 NVLink 高带宽在多个 GPU 之间进行数据传输，并在多 GPU 平台上尽可能多地缓存特征。UGache [41] 提出在缓存更多数据以提高全局命中率和复制更多热点数据以提高局部命中率之间取得平衡。它将特征提取建模为一个混合整数线性规划问题，并构建了特定于平台的缓存方案。然而，基于缓存的方法是为基于样本的 GNN 训练设计的，这会降低模型的准确性 [12, 25, 49]。此外，单独加速 GNN 训练的某个阶段（例如，样本或提取）将错失联合优化计算和通信操作的机会，导致硬件资源利用不足。为隐藏计算延迟而设计的流水线。第三种方法通过构建流水线使计算和通信操作重叠。图 2(c) 示例说明了当 GPU-0 中的节点 1 从邻居节点收集信息时，它可以在获取远程邻居节点（节点 5）的同时执行局部计算（节点 {2,3}）。MGG [49] 打破了特征提取和训练阶段之间的壁垒，以基于节点的粒度使这两个阶段内的计算和通信操作重叠。通过构建细粒度的软件流水线，MGG 成功地将训练时间隐藏在高水平的 GPU 间通信开销中。尽管基于流水线的解决方案看起来很有前景，但现有的设计存在工作负载划分不平衡的问题，导致不同阶段之间的操作重叠效果不佳（第 3.1 节）。总之，现有的设计主要关注在一个（通信）或两个（计算-通信重叠）视角内探索高效解决方案。这些努力错过了从全局视角构建完整优化空间的机会，包括计算、通信和存储。因此，它们未能充分利用现有硬件资源的潜力。

在本节中，我们首先通过在平台 A 上进行的初步实验揭示现有解决方案中的瓶颈（细节见第 6.1 节）。然后，我们分析了数据缓存和操作并行性的潜力，以解决上述问题。

3.1 性能瓶颈分析
定制通信操作的开销。我们首先对 UVA 进行了实验分析，这是一种广泛用于将 GNN 训练扩展到超出 GPU 内存容量的数据集的主流技术。我们使用 DGL 与 UVA（称为 DGL-UVA）来分析一个两层 GCN 模型的整体性能。图 3(a) 将 GNN 训练时间分为 CPU-GPU 数据传输和计算。如图所示，CPU-GPU 数据传输占平均执行时间的 85%。昂贵的传输开销大大削弱了 GPU 高计算能力的优势，导致训练效率低下。
图 3. (a) DGL-UVA 中 GNN 训练的延迟分解。(b) MGG 中分解后的内核延迟比较。在这个实验中，我们将 MGG 内核分解为两个独立的内核。
当前流水线设计的低效率。我们对 MGG 进行了实验分析，这是一种利用流水线构建来促进计算-通信重叠的先进方法。图 2(c) 展示了 MGG 的核心思想，即在单个 GPU 内核中同时执行本地和远程邻居的聚合。为了进一步分析流水线设计的性能，我们将 MGG 的内核分解为本地聚合和远程聚合，用于两层 GCN 模型。如图 3(b) 所示，由于 GPU 间通信的高成本，远程邻居聚合占据了总体执行时间的绝大部分（超过 95%）（更多硬件利用率的结果见图 16(b)）。它还显示，现有的流水线设计存在严重的停滞和空洞，导致 GPU 资源利用率不足。为了进一步调查远程邻居聚合高开销的原因，我们在一个配备了八个 A100 GPU 的单个节点上，按照每个图的顶点 ID 顺序划分了五个真实世界的图拓扑结构。表 1 显示了每个图分区中本地边缘和远程边缘的比例（源节点和目标节点位于不同的分区）。可以观察到，对于大多数图拓扑结构，将其划分为多个子图会显著增加远程边缘的比例。加之远程数据访问的成本远高于同等大小的本地数据访问成本，这导致了本地和远程邻居聚合之间的开销严重不平衡。此外，图中顶点度分布的高度不对称性进一步加剧了图分区之间的工作负载不平衡，从而阻碍了 GNN 训练。先前的研究尚未解决这个问题。

表 1. 数据集类型 | 本地边数 | 远程边数
--- | --- | ---
Part. 1-4 | 1.95M | 13.0M | 1.83M | 12.6M | 1.86M | 12.7M | 1.77M | 12.5M | 12.8M |
Part. 5-8 | 1.83M | 12.7M | 1.77M | 12.5M | 1.69M | 12.2M | 1.64M | 12.1M | 20.6M | 14.0M |
EN | 1.31M | 13.1M | 11.9M | 12.4M | 10.5M | 15.4M | 12.9M | 8.8M | 9.3M | 17.2M | 10.2M |
IT | 116M | 4.86M | 101M | 2.58M | 126M | 2.56M | 146M | 169M | 17.4M | 4.5M | 118M |
PD | 7.5M | 0M | 0.59M | 3.57M | 0.6M | 4.8M | 0.6M | 0.6M | 0.6M | 10.7M | 10.9M |
COP | 10.9M | 0(0%) | 4M | 7.87M | 7.57M | 9.27M | 6.49M | 12.0M | 3.72M | 15.0M | 2.45M | 11.7M | 11.8M | 1.52M |

3.2 缓存和并行性的机会
缓存。受到之前基于缓存设计的探索和第 3.1 节观察结果的启发，我们认识到将缓存技术集成到全图训练的流水线设计中的潜力。通过将频繁访问的连接特征（CNs）预加载到 GPU 中，我们可以提高数据重用率并避免潜在的通信开销。具体来说，缓存可以最小化 GPU 之间的数据传输，从而有效减少远程邻居聚合所花费的时间。图 2(d) 展示了一个简单的例子。对于 GPU-0，我们将节点 5 的嵌入预加载到本地内存中，避免了从 GPU-1 获取数据以更新节点 {1,3} 的信息，从而减少了 GPU 间的通信开销。此外，缓存解决方案还可以改善本地/远程数据的比例，减少现有流水线中的空洞。
并行性。首先，我们观察到当前设计中主机端资源未被充分利用。对于图 2(b) 所示的缓存优化设计，CPU 在特征提取阶段向 GPU 发送数据，在训练阶段保持空闲。对于图 2(c) 所示的基于流水线的设计，CPU 端资源始终处于空闲状态。其次，我们观察到图中每个节点从其不同邻居检索嵌入不会引入任何数据依赖性。基于上述两个观察结果，现有设计的低效率突显了在异构环境中并行执行来自不同邻居的信息聚合的机会。通过利用异步的 CPU-GPU 数据传输，我们可以进一步提高设备利用率，从而加速整体流水线执行。图 2(d) 提供了一个示例，我们在 CPU 端缓存节点嵌入（CNEs）（节点 {3,4}），同时执行 CPU 到 GPU 的通信（传输节点 {3,4}）、GPU 间通信（传输节点 2）和计算。这种解决方案一举实现了两个目标：首先，它可以将部分通信负载转移到 CPU 上，减少由随机访问模式引起的 GPU 间的链接拥塞；其次，利用基于 CPU 的异步传输可以进一步隐藏通信开销。

4. 用于 GNN 流水线的定制分层缓存
在构建 GNN 流水线定制缓存时，我们总结了三个主要挑战：(1) 如何确定缓存策略以最大化流水线效率；(2) 如何在异构环境中布局模型输入以确保流水线高效执行；(3) 如何减少缓存更新的开销，以便最大化缓存的好处。LEO 通过分层缓存设计、混合数据放置和异步缓存更新来解决这些挑战。
4.1 分层缓存设计
与基于样本的 GNN 不同。如第 2.2 节所述，基于样本的训练和全图训练具有不同的数据处理模式，这使得在全图 GNN 定制的流水线中管理缓存数据变得具有挑战性。首先，现有的主流基于样本的 GNN [43, 45, 47] 在单次上传中将小批次的 k 跳计算图及其所有相关特征从 CPU 传输到 GPU。因此，在逐层训练过程中不需要进一步的特征传输。相比之下，训练全图 GNN 具有更大的挑战性，因为它引入了不规则的四层间通信模式，需要精心设计的解决方案。其次，基于样本的 GNN 遵循均匀采样固定数量邻居的范式（即扇出），而全图 GNN 的工作负载是不规则的，因为它由图划分决定。每个分区的子图具有不同数量的节点和边。第三，基于样本的 GNN 通常涉及密集、规则的工作负载的单一传输，可以利用高度优化的通信库（例如 NCCL）。相比之下，全图 GNN 的计算-通信重叠涉及更细粒度的工作负载（例如单个 NE），使得这些小工作的调度更加困难。
适当的缓存策略对于提高 GNN 流水线的效率至关重要。首先，高效的缓存空间分配可以充分利用 CPU 和 GPU 上的可用内存。其次，适当选择嵌入以构建缓存可以显著减少 GPU 间的通信开销并提高性能。基于上述考虑，我们提出了分层缓存设计，其中“分层”指的是同时在 CPU 和 GPU 两侧缓存嵌入。
全图 GNN 训练在每层聚合图中所有节点的邻居，这避免了小批量训练中采样引入的随机性。在这种情况下，训练阶段的数据访问模式是可预测的。因此，我们选择了基于度数的缓存策略 [30]。这个选择基于一个基本事实：节点的出度越高，被访问的概率就越高。具体来说，对于每个图分区，我们按照节点度的降序对远程节点进行排序，并根据缓存比例 αD 和 αH 分别将它们缓存在设备和主机上。给定一个图 (V, E)，其中 V 是节点集，E 是边集。缓存后，V 进一步分为 VL、VR、VD 和 VH，分别代表本地 GPU、远程 GPU、本地 GPU 缓存和 CPU 缓存上的节点。对于 N-GPU 平台中的每个图分区，节点集 VL、VR、VD 和 VH 的大小分别为 ?|V|/N?、|V|?|VL|?|VD|?|VH|、|V|?αD 和 |V|?αH。给定可用的主机容量 CapH 和 GPU 容量 CapD，LEO 使用 Se 来表示存储单个条目 e 的内存容量，进一步限制在 CPU 和 GPU 侧存储的嵌入条目数量：(|VL|+|VD|)?Se?Ke≤CapD|VH|?Se?Ke≤CapH，(2) 其中 Ke 是与模型相关的参数。根据所分析平台的详细信息，LEO 可以构建加载本地和远程数据的时间，分别表示为 TL 和 TR：TL=∑i∈VL TD?ai←j?j∈VL∪VD，TR=∑i∈VL TD2D?ai←j?j∈VR (3)，其中 TD 和 TD2D 分别表示从本地 GPU 和远程 GPU 读取特征嵌入的时间。ai←j 表示节点 i 和 j 在图拓扑中是否相连。同样，LEO 可以构建从 CPU 侧传输邻居的时间 TH：TH=|VH|?TH2D (4)。在上述方程中，TH2D 表示从 CPU 读入一个特征嵌入到 GPU 的时间。我们的目标是调整 αD 和 αH 以最小化 max{TL,TR}，同时确保 TH 满足 TH≤TR 的约束，以避免使主机到设备的数据传输成为新的瓶颈。缓存策略的选择可以表述为一个混合整数线性规划问题，可以使用成熟的优化工具（如 Gurobi [3]）来解决。我们的分区策略可以促进工作负载在GPU之间的均匀分布，从而避免额外的内存浪费。其次，我们利用第(4.1)节中的缓存策略将远程邻居的嵌入预加载到本地GPU的DRAM中作为CNE。我们不将CNE存储在共享内存中，因为缓存的嵌入是在本地处理的，并且不同GPU上的缓存条目数量是不同的。最后，我们根据第(4.1)节中的策略使用UVM [9]来存储数据作为主机节点嵌入（HNE）。图4. LEO的数据存储布局和索引查询模块。“LNE-i”和“CNE-i”分别指的是存储在第i个GPU上的本地/CNE。“HNE”指的是存储在主机上的CNE。“GP”指的是图的分区。“lb”和“ub”分别指的是NEs的下限/上限。

对于规模更小的分区图结构（GP）数据，包括对应于四种类型NE的图拓扑，我们将它们存储在GPU的全局内存中。对于这种放置方法，我们有两个关键见解：(1) 每个GPU中的GP数据是不同的，且只会被本地GPU访问；(2) GP数据相对于NEs来说足够小，可以完全容纳在本地GPU中。LEO提供了一个索引查询单元，用于在主机/设备端处理GP数据以获取NE索引。如图4所示，我们对不同类型的节点采用不同的查询策略。对于本地/远程节点，我们可以通过处理全局节点ID来确定目标GPU ID和偏移量，因为NEs在多个GPU上均匀分布。对于缓存的节点（在GPU/CPU端），我们提供了一个数组来记录CNE/HNE中的索引（ID映射），这有助于处理相应的GP数据以获取正确的NE。

4.3 异步缓存更新

对于基于样本的小批量GNN训练，样本（包括采样的节点及其k-hop邻居）及其对应的特征被输入到k层模型中。这种训练范式在训练阶段不依赖于GPU之间的通信，从而避免了信息滞后。图5(a)展示了一个简单示例，其中CNEs（Node-{1,3}）减少了第一层聚合过程中的远程访问次数。然而，由于缺乏邻居聚合，缓存的嵌入在后续层中会变得过时，导致信息丢失。图5展示了包含缓存的全图训练的2层GCN示例。在图中，白色/黄色圆圈代表本地/缓存节点；带有虚线边框和条纹填充的圆圈代表信息过时或损坏的节点；黑色箭头表示聚合过程，而红色箭头表示数据在GPU之间的复制。

为了解决这个问题，一个简单的解决方案是从每一层的远程GPU中获取最新的NEs（图5(b)）。然而，由于缓存更新只能在聚合操作之后进行，这引入了额外的开销。如图6(a)所示，对于CNEs，模型需要执行额外的GPU间通信（更新缓存）和计算（执行更新操作）以确保模型信息的完整性。为了解决这个问题，我们基于两个关键观察开发了异步缓存更新：(1) 不同类型的数据操作可以并行执行，因为它们之间没有数据依赖；(2) 通过在GPU硬件处理单元上同时调度多个GPU内核，可以协同执行不同的操作。具体来说，我们在不同的CUDA流中处理不同类型的数据，以最大化不同操作之间的重叠，如图6(b)所示。异步缓存更新还可以进一步填充空闲的GPU周期，提高GPU硬件的利用率。在反向传播过程中，可以采用类似的方法异步获取与缓存嵌入对应的远程梯度，从而实现通信和计算的重叠。图6. 两种缓存更新方法的比较。在图中，“UC”代表更新缓存，“U”代表更新函数（第2.1节）。下标表示正在处理的数据类型，包括LNE、CNE和HNE（第4.2节）。

工作流程和竞争避免。具体来说，我们根据算法1逐层进行前向推理。首先，我们基于NVSHMEM执行聚合操作以更新共享嵌入（LNE），并使用nvshmem_barrier()函数确保所有GPU达到同步点。然后，在所有分区完成聚合后，我们将远程嵌入从共享内存复制到本地GPU上以更新缓存数据（CNE和HNE）。请注意，在这一步中，我们只在共享内存中对数据执行读取操作，在本地私有缓存中执行写入操作，这不会引起数据冲突。最后，在更新缓存数据后，我们对本地数据（LNE、CNE和HNE）执行更新操作，此操作不涉及读取外部数据。

5 管道构建和映射

在异构环境中构建高效的GNN训练管道需要解决两个主要挑战：(1) 如何协调GNN内的混合粒度工作负载以最大化管道效率；(2) 如何将协调好的工作负载映射到底层硬件处理单元以确保管道的高效执行。为了解决这些挑战，LEO提出了针对管道的工作负载协调和基于块的 workload 映射。

5.1 针对管道的工作负载协调

与当前管道设计的不同之处在于，我们提出的分层缓存引入了新的工作负载，使得管道的构建比现有设计更具挑战性。首先，我们需要仔细协调各种操作以实现高效的重叠，包括跨GPU的数据加载、CPU到GPU的数据传输以及从共享/虚拟内存中加载数据。简单的方法会削弱管道带来的优势（见第(3.1)节）。其次，由于每个节点的不同类型邻居数量各异，这种变化突显了定制映射方案的需求。这种工作负载不平衡会在现有设计中由于落后者而放大内核效率低下问题。基于第(3.2)节中的观察，我们提出了一个多级动态管道设计。首先，我们从全局角度定义管道的基本工作负载，包括加载本地数据（LL）、加载远程数据（LR）、加载主机数据（LH）和聚合计算（AC）。具体来说，我们将每个GPU上的邻居集分为四类，包括本地/远程/缓存1/主机邻居，如图7(b)所示。特定类型邻居的聚合包括两个步骤：数据加载和计算。例如，远程邻居的聚合需要依次执行LR和AC。请注意，由于缓存邻居和本地邻居都存储在本地（第4.2节），它们的聚合都需要LL和AC。其次，我们根据三个观察来进一步调度这些基本工作负载以实现管道化：(1) 不同类型的邻居的聚合可以并行执行，因为它们之间没有数据竞争或依赖；(2) 通过同时调度多个GPU内核，GPU可以共享硬件处理单元（例如SM），允许计算和通信操作同时运行。具体来说，我们在不同的CUDA流中处理不同类型的数据，以最大化不同操作之间的重叠，如图6(b)所示。异步缓存更新可以进一步填充闲置的GPU周期，提高GPU硬件的利用率。在反向传播过程中，也可以采用类似的方法，通过异步获取与缓存嵌入对应的远程梯度，从而实现通信和计算的重叠。

图6. 两种缓存更新方法的比较。在图中，“UC”代表更新缓存，“U”代表更新函数（第2.1节）。下标表示正在处理的数据类型，包括LNE、CNE和HNE（第4.2节）。

工作流程和竞争避免。具体来说，我们根据算法1逐层执行前向推理。首先，我们基于NVSHMEM执行聚合操作以更新共享嵌入（LNE），并使用nvshmem_barrier()函数确保所有GPU达到同步点。然后，在所有分区完成聚合后，我们将远程嵌入从共享内存复制到本地GPU上以更新缓存数据（CNE和HNE）。请注意，在此阶段，我们只对共享内存中的数据执行读取操作，对本地私有缓存执行写入操作，这不会引起数据冲突。最后，在更新缓存数据后，我们对本地数据（LNE、CNE和HNE）执行更新操作，此操作不涉及读取外部数据。

5 管道构建和映射

在异构环境中构建高效的GNN训练管道需要解决两个主要挑战：(1) 如何协调GNN内的混合粒度工作负载以最大化管道效率；(2) 如何将协调好的工作负载映射到底层硬件处理单元以确保管道的高效执行。为了解决这些挑战，LEO提出了针对管道的工作负载协调和基于块的 workload 映射。

5.1 针对管道的工作负载协调

与当前管道设计的不同之处在于，我们提出的分层缓存引入了新的工作负载，使得管道的构建比现有设计更加具有挑战性。首先，我们需要仔细协调各种操作以实现高效的重叠，包括跨GPU的数据加载、CPU到GPU的数据传输以及从共享/虚拟内存中加载数据。简单的方法会削弱管道带来的好处（见第(3.1)节）。其次，由于每个节点的不同类型邻居数量各异，这种变化突显了定制映射方案的需求。这种工作负载不平衡会在现有设计中由于落后者而加剧内核效率低下问题。

基于第(3.2)节中的观察，我们提出了一个多级动态管道设计。首先，我们从全局角度定义管道的基本工作负载，包括加载本地数据（LL）、加载远程数据（LR）、加载主机数据（LH）和聚合计算（AC）。具体来说，我们将每个GPU上的邻居集分为四类，包括本地/远程/缓存1/主机邻居，如图7(b)所示。特定类型邻居的聚合包括两个步骤：数据加载和计算。例如，远程邻居的聚合需要依次执行LR和AC。请注意，由于缓存邻居和本地邻居都存储在本地（第4.2节），它们的聚合都需要LL和AC。其次，我们根据三个观察来进一步调度这些基本工作负载以实现管道化：(1) 不同类型的邻居的聚合可以并行执行，因为它们之间没有数据竞争或依赖；(2) GPU可以通过同时调度多个逻辑抽象层次（线程/线程块）来共享硬件处理单元，允许计算和通信操作同时运行；(3) 先进的通信技术可以促进管道的构建，例如使用NVSHMEM进行细粒度的GPU间通信和UVM进行粗粒度的CPU到GPU数据传输。通过将不同类型邻居的聚合分配到不同的逻辑处理单元并调度操作，我们的设计实现了管道中计算和通信的高效重叠。

图7.(a)(b) LEO和现有工作[49]中的工作负载分区比较。“LNP”/“RNP”表示本地/远程邻居分区，作为本地/远程邻居中的细粒度组。每个框代表本地/远程/主机邻居的嵌入加载，其长度表示相对的内存访问时间。在图(b)中，带有对角线条纹的框表示邻居已被缓存。(c)(d)(e) 不同的管道设计。“LL”/“LR”/“LH”表示加载本地/远程/主机数据。“AC”表示聚合计算。由于空间限制，我们将(c)中的管道沿着时间轴分为了两部分。

在本节中，我们将详细阐述如何逐步实现最优管道，从最简单的优化策略开始。

基于粒度意识的管道构建。现有解决方案[49]试图构建初始管道，主要包括两个步骤：1) 通过区分本地和远程邻居（图7(a)）并并行执行相应的聚合，如图7(c)所示，可以重叠LL、LR和AC；2) 通过对邻居进行固定数量的分组，可以将粗粒度和不规则的工作负载转换为细粒度和规则的工作负载。图7(a)提供了一个简单的示例。首先，对于GPU-0中的节点，将其所有邻居分为两部分：本地邻居和远程邻居。其次，通过根据固定数量（图7(a)中的3）进一步划分每个节点的邻居集，可以获得更细粒度的组LNP/RNP。

基于缓存的管道工作负载平衡。基于先前优化的管道由于远程/本地操作比率高而导致严重的工作负载不平衡。如图7(c)所示，远程数据加载占据了管道的大部分时间，原因有两个：(1) 节点在GPU上分布均匀，导致远程邻居的数量远多于本地邻居；(2) 获取远程数据比访问本地数据更耗时。此优化步骤通过缓存频繁访问的远程邻居的嵌入来平衡远程/本地操作比率，从而减少管道的总体执行时间。这种优化的关键在于缓存远程邻居的嵌入。具体来说，我们利用第(4.1)节中的缓存策略将嵌入缓存在每个GPU的本地内存中。图7(b)示例中，我们将Node-{4,7}的嵌入缓存到GPU-0的本地全局内存中。缓存邻居和本地邻居都只涉及本地内存访问，其执行时间比访问远程内存要短。我们进一步压缩管道，分别处理这三种类型的邻居以实现并行聚合。如图7(d)所示，这种缓存优化的内存管理平衡了LL和LR之间的总体执行时间，使得不同工作负载之间的重叠更加紧密。

虽然第二种优化通过平衡工作负载提高了管道效率，但仍有改进的空间。首先，由于特征嵌入的数量庞大，缓存太多特征可能会超出GPU的内存容量。其次，我们观察到由于本地内存访问和计算之间的不匹配，管道中仍然存在小间隙（图7(d)）。这种优化步骤旨在进一步平衡工作负载，同时填补这些小间隙。这种优化的关键在于引入新的工作负载来分摊高昂的远程成本，从而进一步压缩管道。我们的关键观察是，在训练期间CPU资源是空闲的。通过在CPU端缓存数据并执行异步的CPU到GPU数据传输，我们可以进一步减少管道中的执行时间。具体来说，我们利用第(4.1)节中的缓存方案将嵌入卸载到CPU端（图7(b)），这自然产生了一个新的工作负载——主机邻居。我们将多GPU GNN的数据加载分为三种类型，每种类型分别涉及本地数据访问、远程数据访问和主机数据访问，相应的执行时间依次增加。我们调度这三种访问模式和计算，以实现管道设计中的最大计算-通信重叠（图7(e)）。这种设计的好处是有双重的：首先，它通过缓存进一步减少了远程访问；其次，它可以利用主机邻居的计算（AC）来填补管道中的小间隙，从而提高GPU利用率。

LEO设计可以扩展到常见的GNN模型，包括GCN[15, 27]、GAT[44]和其他变体[21, 51]，只需进行少量调整。具体来说，由于不同模型的内部设计不同，它们的本地/远程操作比率也不同。我们可以通过调整缓存比率来平衡本地/远程工作负载，从而实现高效的管道执行。同样，LEO设计可以通过预先调整缓存策略和配置扩展到分布式环境。

5.2 基于块的 workload 映射

高效的管道需要一个有效的映射方案，将高级管道工作负载映射到GPU的底层逻辑处理单元。LEO基于两个主要考虑因素引入了一种定制的映射方案。(1) 在逻辑层面上应选择哪种类型的工作负载作为基本调度单元？我们将每个节点的邻居集分解为不同类型的子集（图8(a)），包括本地邻居（LN）、远程邻居（RN）、缓存邻居（CN）和主机邻居（HN），作为基本的工作负载。例如，图8(a)中的LN-0代表节点0的所有本地邻居。这种工作负载的选择有助于调度并最大化GPU内部的线程资源利用率。一方面，选择更粗粒度的工作负载（例如，合并多个节点的邻居）会降低并行性，并导致逻辑处理单元内的工作负载分散。另一方面，选择更细粒度的工作负载（例如，图8(a)中的邻居分组）会产生更多的同步开销。由于每个节点的不同类型邻居的比例各不相同，简单的划分可能会导致负载不平衡。（2）应该选择哪个GPU逻辑处理单元来处理流水线工作负载？首先，我们选择Block作为基本的工作单元来处理流水线工作负载。因为一个Block可以提供多个具有并行处理能力的warp，它在并行性和同步开销之间取得了良好的平衡。当一个Block由于昂贵的远程访问而被阻塞时，可以切换到其他Block来执行本地计算，从而提高GPU的利用率。在同一Block内，共享内存和寄存器等硬件资源可以被共享，从而促进数据重用。图8中，“LNP”/“RNP”指的是本地/远程邻居分区。“LN”/“RN”/“CN”/“HN”分别对应于本地/远程/缓存/主机邻居。在确定了流水线工作负载和基本执行单元之后，选择合适的映射方案仍然不清楚。可以通过将不同的工作负载映射到Block上来构建不同的映射方案，如图8(b)所示。LEO采用了“1对1”的映射方案，其中每种类型的工作负载都映射到一个Block上，以增加并行性并最大化计算-通信重叠。我们提出这种设计基于以下两个关键见解：（1）在同一Block中混合多种工作负载（图8(b)中的“4对1”和“2对1”）可能会导致严重的线程阻塞和资源竞争；（2）不同类型的工作负载之间没有数据依赖关系，可以独立执行以实现并行化，从而加速整体执行速度。

6 评估
6.1 实验设置
模型和数据集。我们评估了两种不同且具有代表性的GNN模型：GCN [27] 和 GIN [51]。对于GCN和GIN的模型配置，我们采用了2层具有16个隐藏维度和5层具有64个隐藏维度，这些配置来自原始论文[27, 51]。我们使用了7个大规模图数据集进行多GPU GNN训练。表2显示了每个数据集的详细信息。
表2.
#数据集#顶点数#边数#图拓扑体积#特征体积
Reddit (RD)[45]0.2M114M60241438 MB535 MB
Enwiki-2013 (EN)[4]4.2M202M30012789 MB4.7 GB
IT-2004 (IT)[14]41M1.1B256644.4 GB39.4 GB
ogbn-Papers100M (PA)[17]111M1.6B128646.4 GB53.0 GB
ogbn-Products (PD)[22]2.4M61M10047245 MB934 MB
Com-Orkut (CO)[4]3M117M12832459 MB1.5 GB
MAG240M (MA)[22]244M1.75B7681535.3 GB
用于评估的数据集。VolG和VolF表示图拓扑和特征的数据体积。

平台和工具。我们在两个不同的平台上部署了实验，这些平台具有不同的配置，其中平台B是进行广泛测试的主要平台。
- 平台A：通过PCIe 4.0连接的八个A100（80 GB，SXM4）GPU，两个Intel(R) Xeon(R) Gold 6330 CPU @ 2.00 GHz（总共56个核心），以及2.9 TB的宿主内存。
- 平台B：通过NVSwitch连接的八个A100（80 GB，SXM4）GPU，提供400 GB/s的GPU至GPU双向带宽，两个AMD EPYC 7713处理器（总共128个核心），以及1.2 TB的宿主内存。
LEO的实现由13K行的C++和CUDA代码组成。我们使用CUDA（v12.1）、OpenMPI（v4.1.1）、NVSHMEM（v2.11.0）、cuDNN（v9.2）和OpenBLAS（v0.3.28）库编译和链接LEO。所有实验都在Docker中使用NVIDIA的镜像进行，同时使用Python v3.9、PyTorch v1.10.0和DGL v2.0。

基线。在实验评估中，我们将LEO与三个支持多GPU平台上全图GNN训练的代表性系统进行了比较。（1）Deep Graph Library（DGL）[45]是一个先进的GNN框架，广泛用于多GPU场景的大规模训练。它利用UVA[10]通过启用从GPU端进行零拷贝访问来高效提取特征。（2）LEO-UVM是LEO的一个变体，它利用统一虚拟内存（UVM）[9]作为通信后端。UVM在处理基于图结构的不规则计算和通信工作负载方面表现出效率[26, 43, 49]。我们将缓存设计和流水线优化扩展到一个新的基线，该基线利用UVM进行所有数据传输过程。请注意，在LEO-UVM中，UVM同时处理CPU到GPU和GPU之间的通信。此外，LEO-UVM不支持零拷贝数据访问，这使其与基于UVA的解决方案不同。LEO-UVM的目的是展示，为了最大化流水线效率，需要根据通信模式选择不同的通信技术以有效协作。（3）MGG [49]是一个利用流水线技术在全图GNN训练中重叠计算和通信的代表性工作。MGG强调其动态软件流水线以实现细粒度的计算-通信重叠。MGG将数据分布在多个GPU上，并使用NVSHMEM[8]进行GPU间通信。

对于LEO-UVM，我们遵循提供的官方函数接口来实现它。具体来说，我们使用cudaMallocManaged来分配统一内存，这允许CPU和GPU之间隐式的数据处理和迁移。我们还遵循LEO的流水线设计来提高LEO-UVM的性能，并使用cudaMemPrefetchAsync来主动重叠CPU到GPU的数据传输、GPU间传输和GPU计算。我们保留了统一内存的默认页面大小，并遵守所有其他默认配置，以确保我们的实现符合官方发布的性能优势。实际上，这种实现的性能与先前工作中报告的典型基于UVM的设计相当[13]，并且所有运行时参数（例如块大小和线程配置）都与LEO中使用的相同，从而确保了公平的比较。

6.2 端到端性能
与DGL相比。首先，我们比较了利用UVA进行CPU到GPU通信的先进框架DGL。我们在两个平台上评估了不同的数据集。如图9所示，LEO在GCN和GIN模型上分别实现了平均10倍和5倍的加速。我们观察到，即使在配备了PCIe 4.0的平台A上，LEO仍然实现了平均5倍的性能提升，主要有两个原因。（1）LEO利用GPU互连有效地 overlap了来自远程GPU和宿主的数据加载，加速了特征提取。（2）因为多个GPU同时向宿主发送大量请求来获取特征嵌入，DGL遭受了严重的通信拥塞。另一方面，仅依赖UVA在处理大规模数据时存在潜在风险。当数据无法适应CPU内存时，基于UVA的解决方案可能会失败[40]。在这种情况下，LEO可以将数据分布到多个GPU和宿主上，从而减轻内存压力。与仅依赖UVA进行CPU到GPU通信的DGL相比，LEO还利用NVSHMEM进行GPU间通信，从而提高了整体通信效率。实验结果表明，LEO的混合通信模式优于现有主流框架的通信后端，更好地服务于异构环境中的GNN训练。

与LEO-UVM相比。在这个实验中，我们将LEO与其基于UVM的变体LEO-UVM进行了比较，后者在GNN训练中使用了UVM进行通信。注意，LEO-UVM使用UVM进行CPU到GPU的通信，并且还用NVSHMEM替换了GPU间的通信，这与DGL采用的UVA不同。图10显示，LEO在GCN和GIN模型上分别实现了平均11.91倍和11.89倍的加速。UVM利用页面故障机制进行隐式的数据迁移，提供了更粗粒度的数据（单个页面大小为4 KB）。然而，在流水线执行期间，基本调度工作负载通常由单个节点的邻居嵌入组成（小于256B），导致带宽浪费和高的数据迁移开销。在多GPU场景中，流水线由于大多数warp等待耗时的页面错误处理而更多次停滞。

与MGG相比。在这项实验中，我们将LEO与MGG进行了比较，后者利用NVSHMEM实现细粒度的计算-通信重叠。图10显示，LEO在GCN和GIN模型上分别实现了2.54倍和3.95倍的加速。MGG将完整的图嵌入分布在多个GPU上，并依赖高速互连（如NVLink）来加速远程数据访问。然而，如第（3.1）节所分析的，MGG仍然受到高远程通信的影响，导致性能不佳。我们观察到，在处理更高维度的嵌入（例如GIN）时，相对于MGG的加速效果更高。这是因为更大的数据移动量可能导致通信饱和，增加链接拥塞的可能性。我们观察到，在某些数据集（例如EN）上，LEO的性能不如MGG。观察到的性能下降是由于图划分结果造成的。具体来说，更平衡的划分（其中本地和远程访问延迟变得趋于一致）倾向于削弱我们提出的方法在减轻分区间工作负载不平衡方面的有效性。从表1可以看出，EN数据集在划分后显示出更平衡的工作负载分布，从而减少了LEO和MGG之间的差距。总之，我们设计的好处主要来自于解决图拓扑划分引起的工作负载不平衡问题，并且与现有的图划分技术（如Metis [5]）的优化方向正交。

在缺乏高速互连的情况下（例如平台A），MGG的性能显著下降，因为它无法及时处理大量的远程数据访问请求。LEO利用缓存和来自宿主端的异步数据传输来减少GPU间的通信，从而确保在通信带宽有限的场景中具有竞争力。实验结果表明，LEO的设计既有效又通用，可以扩展到不同的平台配置。

6.3 优化分析
分层缓存。为了研究分层缓存对原始流水线的性能改进，我们在平台B上累积启用了不同类型的缓存，并进行了实验验证。对于每个数据集，我们将GPU上的缓存比率设置为20%，CPU上的缓存比率设置为15%。图11(a)显示，在GPU和CPU端引入缓存分别可以带来平均3.12倍和1.22倍的性能提升。原始流水线设计由于其本地/远程邻居分离的设计，面临大量的远程访问请求，导致严重的通信瓶颈。

异步缓存更新。我们将LEO与在平台A上未应用异步缓存更新的基线进行了比较。缓存配置参数与之前的实验相同。图11(b)显示，LEO在一致的情况下优于具有同步缓存更新的基线，实现了平均1.38倍的性能提升。如果没有异步缓存更新，基于缓存的流水线的性能优势可能会被更新缓存的开销所削弱，导致更高的延迟。

映射方案。为了分析映射方案的性能影响，我们在平台B上进行了两项消融研究。首先，我们分别在线程、warp和块级别采用了LEO，训练了一个具有16个隐藏维度的两层GCN模型。对于不同的级别，我们使用了不同的NVSHMEM通信原语来提供高性能支持。图11(c)显示，基于块的映射方案相比其他方案平均性能提升了2.86倍，在并行性和同步开销之间取得了良好的平衡。其次，我们在平台B上评估了各种映射方案，将不同类型的工作负载映射到单个或多个块上。MGG设计实现了本地/远程邻居的分割，其映射策略被归类为“2对1”（图8(b)）。图11(d)表明，采用“1对1”的映射方案平均性能提升了21%，始终优于其他映射方案。本研究证明，我们选择的映射方案有效避免了由于不同类型工作负载不对齐而导致的流水线停滞，使得各种数据访问和本地计算操作之间的重叠更加紧密。

为了量化各个因素的贡献，我们进行了消融实验，提供了LEO相对于MGG的性能提升的详细分解（图12(a)）。为了分离缓存的好处，我们在原始的MGG流水线上增加了一个层次化缓存构建作为基线（+Cache），其中数据通过异步CUDA流批量传输，而不是使用本文提出的流水线设计。为了分离我们的流水线和映射设计的性能提升，我们构建了一个没有层次化缓存的LEO基线（+Pipeline）。图12(a)显示，在六个数据集上测试的GIN模型中，+Cache基线相对于MGG平均提升了1.83倍，而+Pipeline基线平均提升了1.40倍。当所有优化结合在一起时，LEO相对于MGG平均性能提升了2.56倍，这与各个优化提升的乘积非常接近。这验证了我们的优化技术是互补的，层次化缓存是主要的驱动力，而流水线和映射增强提供了额外的提升。

图12。(a) 不同配置下的消融研究和性能分解；(b) 不同分割策略下的LEO性能。

6.4 对图分割质量的敏感性分析。本研究分析了LEO在不同图分割策略下的性能。LEO采用了MGG的默认分割方案，即块分割（BP）。给定图中的节点ID和预设的分割数量，每个节点的分割ID计算为 partition_id = node_id // num_partitions。这种方法产生的预处理开销较轻，并且可以很容易地扩展到处理大型图；然而，它的分割质量相对较低。为了研究LEO在不同图分割策略下的性能，我们增加了一个使用Metis [5]的额外版本，这是一个主流的图分割工具，其多层分割方案显著减少了边切割，同时确保了跨分割的负载平衡。我们使用基于BP和基于Metis的LEO变体作为分别代表低和高新分割质量的基线。图12(b)展示了在六个数据集上相对于MGG的性能提升。我们观察到，在较低的分割质量下，LEO获得了更高的性能提升，平均加速比为2.56倍。在较高的分割质量下，LEO仍然实现了显著的性能提升，平均加速比为2.42倍。

实验结果证实，LEO在不平衡的工作负载下表现出优异的性能，并且在各种真实世界的图中保持了持续的改进，证明了其鲁棒性和通用性。

6.5 缓存的好处。本研究将分析在8个NVIDIA A100 GPU上引入缓存的好处和开销之间的权衡。我们通过调整不同的缓存比例，在Enwiki-2013数据集上训练了一个2层GCN模型。图13表明，引入层次化缓存可以有效地加速推理和训练过程，而更新缓存的开销微不足道。此外，实验结果表明，随着缓存比例的增加，前向推理的时间不断减少，而反向传播的时间增加。这种现象是由于在GPU上本地缓存数据，减少了推理期间昂贵的远程数据访问，但引入了额外的计算和更新开销。在实际应用中，我们使用成本模型（4.1）来确定缓存策略，这通常比一次迭代的时间要短。考虑到端到端GNN训练通常涉及超过100次迭代，因此获得的性能提升是显著的。我们提供了详细的内存分解分析（表3），包括LNE、CNE和HNE，在不同的数据集和平台上进行了测试。在确保分配了NVSHMEM共享内存（LNE）之后，我们再分配CNE，它通常占最大部分。对于GPU间通信带宽较低的场景，我们增加CNE的比例以掩盖昂贵的远程邻居获取成本。对于像MAG240M这样的大规模图，我们采用成熟的量化技术[33]来压缩特征维度，以避免内存不足错误。在某些存储受限的场景下，例如在固定内存预算下，我们调整成本模型（4.1）的约束，从完全隐藏昂贵的远程数据访问改为尽可能多地分配缓存空间，从而使模型训练也能从我们的设计中受益。

图13. 缓存带来的好处评估。“D”/“H”代表设备/主机侧的缓存比例。

表3. 数据集 LNE(A) CNE(A) HNE(A) LNE(B) CNE(B) HNE(B) RD 68 MB 348 MB 85.6 MB 68 MB 269 MB 107 MB EN 624.4 MB 3.05 GB 721.9 MB 624.4 MB 2.35 GB 963.6 MB IT 4.92 GB 9.85 GB 2.76 GB 4.92 GB 7.88 GB 3.9 GB PA 6.67 GB 11.66 GB 4.77 GB 6.67 GB 10.6 GB 5.3 GB PD 120.1 MB 653.8 MB 130.76 MB 120.1 MB 467 MB 186.8 GB CO 194 MB 1.05 GB 184.32 MB 194 MB 768 MB 307.2 MB MA 10.87 GB 7.35 GB 2.6 GB 10.87 GB 6.09 GB 2.65 GB

不同数据集和平台上的内存分解。A和B代表测试平台。

6.6 准确性-延迟权衡。在本研究中，我们比较了LEO与基线模型在不同数据集和架构上的准确度-时间。我们评估了表1中列出的数据集，除了IT-2004和Com-Orkut，因为它们缺乏真实标签。对于MAG240M，我们对特征嵌入进行了量化，使其可以在我们的平台上进行训练。对于LEO，我们实现了一个使用同步缓存更新的基线，以调查我们提出的异步缓存更新是否会导致最终准确度的任何下降。图14中的实验结果表明，LEO与基于采样的DGL基线相比，实现了更高的收敛准确度，提高了多达6.8个百分点。与全批量基线MGG相比，LEO收敛速度更快，同时没有牺牲准确度，训练时间最多减少了4倍。实验表明，异步缓存更新在减少训练时间的同时，不会降低最终准确度，平均减少了30%，验证了我们设计的有效性。为了验证LEO的通用性，我们在Ogbn-Papers100M数据集上使用了四种代表性的GNN模型进行了实验。如图14的下方面板所示，LEO在不同架构下始终实现了更高的准确度，并显著减少了训练时间。

6.7 远程访问减少。我们进一步分析了层次化缓存对远程访问减少的效果在平台B上的表现。我们统计分析了在不同GPU配置下使用和不使用缓存时的远程访问请求数量。对于每个GPU，我们保持了相同的数据缓存比例，25%的节点缓存在设备侧，15%缓存在主机侧。图15(a)报告了缓存后的标准化结果，显示层次化缓存可以在三个数据集中分别将远程GPU数据访问请求减少了85.9%、88.2%和94.7%。结果表明，我们的缓存方案可以显著减少远程内存访问需求。对于基于样本的小批量训练，LEO可以通过修改缓存方案（例如，预采样[53]）扩展到需要洗牌的场景。我们将这个扩展留待未来的工作。

图15. (a) 利用层次化缓存的远程访问。(b) 在4×V100平台上的性能分析。

6.8 在其他平台上的性能。我们测试了LEO和DGL在4×PCIe-V100平台上训练2层GCN模型的性能。该平台有4个40GB的PCIe GPU通过PCIe 3.0连接，两个Intel Xeon Gold 6252N CPU（总共48核心），以及1.2 TB的主机内存，代表了低速GPU互连平台。注意，缓慢的GPU间互连加剧了远程访问的尾延迟，从而削弱了我们精心设计的流水线的好处。更严重的是，GPU的内存容量有限，严重限制了我们缓存嵌入的能力。图15(b)表明，在GPU内存和GPU间通信受限的条件下，LEO相对于DGL实现了平均2.64倍的性能提升，验证了我们设计的有效性和通用性。结合第(6.2)节的评估结果，我们有理由相信LEO可以在未来的更快互连平台（如CXL [1]）上实现良好的性能。我们将这个探索留待未来的工作。

6.9 对其他模型的通用性。为了进一步验证LEO的通用性，我们在两个额外的GNN模型上评估了其性能：GraphSAGE [21]和GAT [44]。我们在平台A上测试了这两个模型，它们都配置了两层和16的隐藏维度。对于GraphSAGE，我们采用了“10-25”的邻居采样方法。与GCN和GIN不同，GAT结合了基于边注意力的更复杂的消息传递方案。我们进一步将最初为节点级聚合阶段设计的LEO优化扩展到边注意力计算。为了公平比较，我们还扩展了MGG的核心流水线设计以处理GAT模型。图16(a)表明，LEO相对于DGL在GraphSAGE模型上实现了平均5.75倍的性能提升，在GAT模型上实现了7.37倍的性能提升。与MGG相比，LEO在GraphSAGE上实现了平均3.64倍的加速，在GAT上实现了4.16倍的加速。如第(5.1)节所讨论的，通过微调本地/远程操作比例，我们的设计在其他GNN模型上也保持了稳健的性能。因此，即使其他GNN模型（例如GAT）表现出不同的计算和内存访问模式，我们仍然可以通过调整缓存方案来实现持续的加速。

图16. (a) 在其他模型上的性能分析。(b) 硬件利用率比较。

6.10 资源利用率。为了评估硬件资源的利用率，我们在EN数据集上训练了一个GCN模型。我们使用NVIDIA Nsight Systems [36]来记录GNN训练期间CPU和GPU的活动。该工具提供了系统级别的跟踪和采样，使我们能够捕获CPU线程、CUDA API调用和GPU内核执行的统一时间线。图16(b)展示了GPU和CPU的利用率对比。LEO通过基于多阶段的缓存流水线实现了更高的硬件利用率，CPU和GPU的利用率分别为78.4%和94.4%。实验结果表明，充分利用硬件资源可以显著提升性能。

7 相关工作。全图GNN系统。关于全图GNN训练的扩展工作有很多。例如NeuGraph [32]、ROC [25]和DistGNN [34]等努力设计了智能的图分割策略，以减少全图训练期间的内存消耗和通信开销。GNNAdvisor [48]和MaxK-GNN [38]利用输入图的结构特征和中间特征表示来加速全图训练，同时提供用户友好的编程接口。TCGNN [50]和DTC-SpMM [16]设计了高性能的稀疏矩阵乘法内核，专为全图GNN量身定制。此外，HongTu [46]引入了一个去重通信框架，消除了多余的CPU-主机传输，而MGG [49]提出了一种流水线方案，旨在重叠并从而掩盖GPU间通信开销。

基于采样的GNNs的缓存。基于采样的GNN系统通过缓存频繁访问的数据来加速特征提取。PaGraph [30]提出了一种基于度数的缓存策略，保留了高度节点的特征。GNNLab [53]引入了一种基于预采样的缓存策略，实现了更高的命中率。UGache [41]利用高带宽GPU互连来加速远程GPU内存访问，并提供了针对不同硬件配置的数据放置方案。BGL [31]将缓存设计进一步推进了一步，在静态和动态缓存之间实现了开销和效率的平衡。它利用了小批量内节点的时间局部性来提高缓存命中率。此外，它还设计了一种针对图采样的图分割算法，缓解了节点之间的负载不平衡。尽管上述工作做出了重要贡献，但它们的优化都是为基于采样的GNN设计的（第4.1节）。LEO是第一个解决将缓存技术移植到全图GNN训练中的根本挑战的。

流水线技术。流水线技术已在深度神经网络训练中得到了广泛采用。PipeDream [35] 和 PipeSGD [29] 将 DNN 训练组织成流水线，并故意使用过时的参数来减少流水线中的等待时间。这些方案专为密集且结构良好的 DNN 工作负载设计，并未扩展到图神经网络（GNNs）所展示的稀疏、不规则的通信模式。NeutronOrch [11] 将 GNN 训练分解为 CPU 侧计算、CPU 到 GPU 的通信以及 GPU 侧计算，并通过软件流水线将它们重叠起来，但它未能充分利用 GPU 之间的高通信带宽。MGG [49] 实现了 GPU 之间的通信与计算重叠，但仍未能消除由于远程/本地工作负载比例极度不平衡而导致的严重通信瓶颈。

**讨论：**
**可移植性和泛化能力。**
虽然 LEO 目前已经在利用 CUDA、NVSHMEM 和 UVM 的 NVIDIA 平台上实现和评估，但其核心算法贡献——基于度的分层缓存策略、混合粒度流水线编排以及基于块的 workload 映射——都是与架构无关的设计原则，不依赖于特定供应商。为了阐明 LEO 的可移植性，我们将其组件分为三层：
(1) **完全可移植的组件（算法层）**：缓存替换策略（第 4.1 节）、用于平衡本地/远程/主机 workload 的流水线构建逻辑（第 5.1 节）以及“1 对 1”块映射策略（第 5.2 节）都是高级逻辑设计。它们仅依赖于图拓扑（度）和性能分析时间（TL, TR, TH），这些在任何平台上都是通用的。这些组件可以直接应用于其他加速器而无需修改。
(2) **特定于 NVIDIA 的后端组件（实现层）**：当前的运行时严重依赖于 NVIDIA 特定的库来实现流水线：
(a) NVSHMEM 用于细粒度的 GPU 间读取（LNE），
(b) UVM 用于粗粒度的主机到 GPU 的数据迁移（HNE），
(c) CUDA Streams/Events 用于异步内核调度。
(3) **适用于其他平台的直接替代方案**：
对于 AMD (ROCm) 或 Intel (oneAPI) 生态系统的可移植性，这些后端有直接的功能等效物：
(a) ROC_SHMEM/OpenSHMEM 或 Intel SHMEM 可以完美替代 NVSHMEM，用于对称的全局内存访问；
(b) ROCm Unified Memory 和 Intel Shared Virtual Memory 提供与 UVM 相同的页面迁移语义；
(c) HIP Streams/SYCL Queues 提供相同的异步执行能力。
在迁移到不同的硬件/软件生态系统时，超参数的选择可以由我们提出的求解器自动处理，而无需对核心系统逻辑进行任何修改。

**未来工作：**
**扩展到基于样本的训练。**
在扩展到基于样本的 GNN 训练时，我们的动态流水线和基于块的映射策略保持不变，因为它们与训练范式无关。缓存策略的调整需要解决两个方面的问题：
(1) 基于 MILP 的缓存比率选择必须从全图度分布调整为预期采样分布（由批量大小和扇出量决定）；
(2) 静态的基于度的缓存更新必须增加动态驱逐策略（例如 LFU 或 LRU），以处理随机打乱引入的时间局部性。
这些挑战是我们模型的算法扩展，我们将它们作为未来的工作。

**结论**
本文介绍了 LEO，这是一种新型系统，它利用分层缓存技术在多 GPU 平台上加速 GNN 训练。通过对现有系统的全面分析，我们发现了硬件资源利用不足和严重通信瓶颈等问题。LEO 通过引入分层缓存和构建混合粒度软件流水线来解决这些问题。具体而言，LEO 在本地 GPU 和 CPU 两侧缓存最热的特征嵌入，以减少昂贵的远程数据访问。此外，LEO 构建流水线以实现不同 workload 和操作的最大重叠。在各种平台上的广泛实验验证了 LEO 的有效性和泛化能力。