论文 | GraNNDis：大型集群上深度 GNN 的高效统一分布式训练框架

本文约2500字，建议阅读5分钟
本文提出了 GraNNDis，这是一种高效的分布式 GNN 训练框架，用于在大型图和深层上训练 GNN。

图神经网络（GNN）是深度学习中发展最快的领域之一。根据数据集的增长和 GNN 使用的模型大小，一个重要的问题是几乎不可能将整个网络保留在 GPU 内存上。在众多尝试中，分布式训练是解决该问题的一种流行方法。然而，由于 GNN 的性质，现有的分布式方法的可扩展性较差，这主要是由于外部服务器通信缓慢。

在本文中，我们提出了 GraNNDis，这是一种高效的分布式 GNN 训练框架，用于在大型图和深层上训练 GNN。GraNNDis 引入了三种新技术。首先，共享预加载为多 GPU 服务器集群提供了训练结构。我们建议对基本顶点依赖项进行服务器方面的预加载，以减少低带宽外部服务器通信。其次，我们提出了扩展感知采样。

由于共享预加载本身因邻居爆炸而存在局限性，因此扩展感知采样可以减少跨越服务器边界的顶点依赖性。第三，我们提出协作批处理，为全图和小批量训练创建统一的框架。它显着减少了小批量训练中的冗余内存使用。由此看来，GraNNDis 通过统一在全图和小批量训练之间实现了合理的权衡，特别是当整个图不适合 GPU 内存时。通过在多服务器/多 GPU 集群上进行的实验，我们表明 GraNNDis 比最先进的分布式 GNN 训练框架提供了卓越的加速。

https://arxiv.org/abs/2311.06837

训练大图数据不可避免地面临GPU工程挑战。即使采用最好的多 GPU 设置，也不能保证获得“良好”的学习成果。在多 GPU 服务器上训练图神经网络（GNN）时，消息传递的具体过程以及如何跨服务器优化利用资源可以通过我们今天介绍的 GraNNDis 论文来了解。

许多人在为深度学习模型设置计算环境时都经历过这样的困境：每个参数需要多少 GB 的 VRAM。这包括使用 Colab 等基于云的解决方案。当您自信地将数据加载到 VRAM 时，您可能会遇到可怕的内存不足 (OOM) 错误。通过减少批量大小以及其他策略来缓解这种情况的尝试很常见，因此有一句格言说，有时实验的设置比实验本身花费更多的精力。

随着数据规模的不断增长，工程能力变得越来越重要。现实世界的图数据大多遵循幂律分布，因此更容易受到 OOM 的影响。例如，在将权重加载到 VRAM 时，一些边缘较少的节点计算的权重较少，而边缘较多的节点可能会超出 VRAM 容量，从而导致 OOM。

为了防止 GNN 中的 OOM，GraNNDis 提出了一种利用共享预加载、扩展感知采样和协作批处理的想法。这些策略通过增强 GPU 服务器通信效率、服务器间采样以及用于学习整个图的批处理技术来改进 OOM，并通过实验证明了卓越的性能。

在进行实验之前，有必要建立“良好”学习的标准。该论文根据每个服务器的 GPU 计算节点和边缘的数量以及 GPU 之间的通信时间来设置基准。简而言之，它是在服务器 GPU 上并行分配和计算过多的数据，并测量集成这些计算所需的时间。

共享预加载

同一服务器上的 GPU 网络和不同服务器上的 GPU 网络之间的通信时间必然有所不同。显然，如果所有图形都加载到单个服务器的 GPU 上，速度会更快，因为跨多个服务器的 GPU 加载它们需要额外的服务器访问权限来收集信息，从而导致处理时间更长。

为了解决这个问题，我们使用了一种称为共享预加载的方法。这种方法不同于传统的通过在服务器阶段预先分配和预加载节点以供各层使用的跨服务器均匀分布图的方法。当发生外部-内部操作时，它利用预加载节点的特性进行操作。

扩展感知采样

在讨论我们的方法之前，有必要简要介绍一下图神经网络（GNN）中如何进行采样，因为这将有助于向大家更好地概述本文。在 GNN 中，传统的采样方法可以分为两大类：逐层采样和子图采样。逐层方法涉及对每层最多 n 个邻居进行采样，而子图方法涉及对连接到特定节点的子图进行采样。

逐层采样方法的优点在于，它允许对每一层的消息传递操作的邻居进行任意控制。另一方面，子图采样方法通过对反映良好的子图进行采样，有效地捕获原始图的结构特征。举一个很多人都熟悉的例子，在 torch_geometric 模块的上下文中，分层采样可以等同于 NeighborLoader，子图采样可以等同于 GraphSAINTRandomWalkSampler。

回到我们的主题，即使使用共享预加载，随着层数的加深和采样节点数量的增加，预加载也不可避免地会受到限制，因为大多数现实世界的图数据都表现出强大的能力——规律分布。为了缓解这一限制，我们采用扩展感知采样。

这个概念很简单。与传统采样方法不同，扩展感知采样通过引入最大跳数和扇出等附加参数来关注“扩展”方面进行采样。它允许对距特定节点的跳数以及每个邻居要加载的最大连接数进行可调规范。此方法有助于设置标准来决定在共享预加载期间包含或排除哪些节点。

协作批处理

到目前为止，我们已经讨论了如何使用共享预加载和扩展感知采样来进行全图训练，即训练整个图。另一方面，协作批处理解决了如何为每个 GPU 获取节点的问题。

传统方法首先指定目标节点并将其划分，然后根据划分的节点进行读取，这种方法称为 split-then-fetch。协作批处理采用先获取后分割的方法，在决定哪个 GPU 将训练它们之前，会为每个批次预先获取服务器要使用的节点。

传统方法将每个目标节点分配给不同的GPU进行训练，需要为每个节点获取所有相关的邻居。然而，协作批处理通过跨批次分配公共节点，允许 GPU 分担计算负载。这使得它在内存利用率方面更加高效，因为它减少了每个 GPU 独立获取和存储所有邻居的需要。

实验

我们使用 PyG 和 DGL 等框架作为基线进行实验。值得注意的是，虽然DGL的pipeGCN和DistDGL有局限性，例如分别不支持小批量和全图训练，但GraNNDis适用于这两种场景。

在所有有趣的实验结果中，我发现图 11 是最有趣的。它侧重于测量内存效率以及内存不足 (OOM) 发生的时间点。考虑到 56 层和 256 维特征的巨大性，参数的绝对数量通常会令人难以承受。通常情况下，我们根本不会想到在这些条件下进行训练，但这篇论文让我们能够间接体验和观察这些结果，我觉得这很有趣。

此外，人们可能会认为，由于不同采样方法更新参数的方式发生变化，准确性等性能可能会下降。然而实验表明，在不牺牲性能的情况下，训练速度可以提高1.23倍到1.95倍。

结论

在当今大型语言模型（LLM）成为热门话题的时代，有很多关于如何在给定可用资源的情况下有效地训练大型数据集的讨论。与其他类型的数据不同，图数据由于其对拓扑的强烈依赖而具有独特的特征。节点学习的信息量根据指定的拓扑而变化，这与性能直接相关并且被认为是关键元素。

然而，在某些情况下，这种拓扑结构没有得到有效利用。由于拓扑的富者愈富的性质，一些 GPU 内存可能在加载少量信息时未得到充分利用，而另一些 GPU 内存可能因加载大量信息而过载，从而导致训练期间出现内存不足 (OOM) 问题。

为了防止这种情况，策略通常涉及减少批次大小或样本大小，以最大限度地减少每批次的信息量。然而，这可能会导致错过有意义的信息。

本文通过引入 GraNNDis 架构来解决这个问题，该架构具有共享预加载、扩展感知采样和协作批处理等功能。对于那些只关注消息传递的人来说，这是一篇很有价值的论文，因为它通过研究 GPU 上参数的计算过程拓宽了视野。

编辑：于腾凯

校对：林亦霖