由 Meta 的 PyTorch 团队：杨春志、朱富、胡赞、徐亮贝、张欣、杨继岩、范德·斯泰伊·丹尼斯、周王、方金、聂佳、胡宇熙

在 Meta，推荐系统是全球向数十亿用户提供相关和个性化广告的基石。通过像 PyTorch 的 TorchRec 这样的技术，我们已经成功开发了能够跨数百个 GPU 进行模型训练的解决方案。虽然这些系统为我们提供了良好的服务，但最近关于扩展定律的研究揭示了一个令人信服的机会：通过训练显著更大的神经网络，我们可以实现显著更好的模型性能。

然而，这一洞察给我们带来了新的挑战。我们当前的训练基础设施，虽然针对数百个 GPU 进行了高度优化，但无法有效地扩展到训练这些更大模型所需的数千个 GPU。从数百个 GPU 到数千个 GPU 的飞跃引入了复杂的技术挑战，特别是在处理推荐模型中的稀疏操作方面。这些挑战需要全新的分布式训练方法，我们通过一种新颖的并行化策略来应对这些挑战。

为了解决这些问题，我们引入了 2D 嵌入并行，这是一种新颖的并行策略，可以克服在数千个 GPU 上训练大型推荐模型时固有的稀疏扩展挑战。今天，您可以通过 TorchRec 的 DMPCollection API 使用它。这种方法结合了两种互补的并行化技术：数据并行用于模型的稀疏组件，模型并行用于嵌入表，利用 TorchRec 强大的分片能力。通过战略性地整合这些技术，我们创建了一个可以扩展到数千个 GPU 的解决方案，现在它为 Meta 最大的推荐模型训练运行提供动力。

什么是稀疏扩展挑战？

我们确定了三个关键挑战，阻止我们将模型天真地扩展到数千个 GPU：

  • 不平衡和落后者问题：随着 GPU 数量的增加，实现平衡分片变得更加困难，某些排名可能需要更重的嵌入计算工作负载，这可能会减慢整个训练过程。
  • 节点间通信：随着训练作业使用越来越多的 GPU，在特定的网络拓扑下，全互连通信带宽可能会下降，这会显著增加通信延迟。
  • 内存开销：输入特征使用的内存通常可以忽略不计，然而，当我们使用数千个 GPU 时，我们可以引入更大的输入特征，内存需求可能会变得显著。

在 2D 嵌入并行的情况下，我们可以这样描述我们的新并行方案，在这个例子中，我们有 2 个模型副本（副本 1：GPU1/GPU3，副本 2：GPU2/GPU4）

Flow diagram

图 1：2D 稀疏并行布局示意图

通过 2D 稀疏并行处理，我们应对这些挑战，而不是将表分片到所有等级，我们首先将所有等级平均分成几个并行组：

  1. 在每个组内，我们使用模型并行处理嵌入表，例如列/行分片。在规模较大的情况下，对于我们最大的表，我们还开发了网格分片，将嵌入表分片在行和列维度上。
  2. 在组之间，我们进行数据并行，使得每个组中的每个等级都有其在其他组中的对应副本等级（副本等级意味着存储相同的嵌入表分片）。
    1. 每个组完成自己的反向传播后，我们将所有副本的嵌入表权重进行汇总，以保持它们同步。

我们的生产解决方案

TorchRec 是我们用于在原生 PyTorch 中构建推荐模型稀疏部分的库。传统的 API 是 DistributedModelParallel，它将模型并行应用于嵌入表。我们在此 API 旁边引入了一个新的 API，称为 DMPCollection，它是启用 TorchRec 模型 2D 并行的主要入口点。我们设计它，使其更改尽可能简单，就像应用 FSDP/DDP 一样简单。

要理解 DMPCollection 做什么，我们首先需要了解 DistributedModelParallel (DMP) 做什么：

  1. 创建嵌入表，称为 EmbeddingBagCollection 和 EmbeddingCollections。
  2. 根据 GPU 拓扑、嵌入表、可用内存、输入数据等因素生成分片计划。
  3. 将模型包装在 DMP 中，并传递相关的分片计划。
  4. DMP 根据分片计划初始化并分片嵌入表。
  5. 在训练步骤中，DMP 接收一个输入批次，将其传递给包含所需嵌入表分片的相关 GPU，查找值，并将其返回给请求 GPU，所有这些操作都在全局进程组中进行，对于特殊分片（如表行分片）有一些例外。

分布式模型并行是为模型并行设计的，许多部分在假设分片和绕过全局世界大小的情况下工作。我们需要对这些部分进行修改，以便在不丢失 TorchRec 的优化和功能集的情况下引入额外的并行维度。

DMPCollection 通过改变几个关键部分，以可扩展的方式实现 2D 并行。

  • 一次性生成较小分片组的分片计划，一旦传入，我们就向全局组中的适当排名进行通信，并将排名重新映射以适应新的分片组排名。
  • 创建两个新的 NCCL 进程组，称为分片和副本进程组。分片进程组被传递到 TorchRec 的分片和训练步骤组件中。副本进程组用于权重和优化器状态同步，所有 reduce 调用都通过此进程组进行。
    • 子 NCCL 进程组使我们能够高效地仅在与特定通信相关的排名之间进行通信。每个排名将有两个相关的进程组。

对于用户来说，这种变化非常简单，同时去除了将并行策略应用于模型的所有复杂性。

我们如何创建这些分片和复制组？

这些进程组是 DMPCollection 高效实现的关键之一。从我们之前的图中，我们展示了简单的 2x2 GPU 设置，然而，在规模上，我们如何分配哪些排名属于给定的分片组以及它们在分片组中的副本排名？

考虑以下配置，有 2 个节点，每个节点有 4 个 GPU。在 2D 并行下的分片和复制组将是，

分片组 分片秩
0 0, 2, 4, 6
1 1, 3, 5, 7
复制组 复制等级
0 0, 1
1 2, 3
2 4, 5
3 6, 7

我们使用以下公式，

  1. 将所有训练师分为 G 个分片组，每个组有 L 个训练师
    1. 组数 G 由 G = T / L 确定，其中 T 是训练师总数
  2. 对于每个组 G，我们根据其所属的组分配了非连续的训练器排名，具体为，
    1. [i, G+i, 2G+i, …, (L - 1) G+i]，其中* i = 0 到 G-1*
  3. 从组 G 中，我们可以创建复制组，即每个 G 个连续的排名
    1. （0 到 G-1，G 到 2* G - 1）每个连续集合存储了副本嵌入表碎片。

这意味着我们的分片组 G 的大小为 L，这可以称为应用模型并行时的排名数量。这反过来又给我们提供了副本组，每个大小为 G，这是我们数据并行的排名。

在 DMPCollection 中，我们可以通过使用 DeviceMesh 高效地创建这些进程组，我们创建整个 GPU 拓扑为一个 2x2 矩阵，其中每一行代表分片排名的组，每一列代表相应的副本排名。

create peer matrix
num_groups = global_world_size // sharding_group_size
for each group_rank in num_groups:
	peers = [num_groups * rank + group_rank for rank in range(sharding_group_size)]
	add peer to peer matrix

initalize DeviceMesh with two dimensions (shard, replicate)
slice DeviceMesh on shard for sharding process group
slide DeviceMesh on replicate for replica process group

使用我们的 DeviceMesh 方法，如果我们想要在未来更改拓扑或提供更多的灵活性，我们可以轻松地将我们的创建逻辑扩展到任何形式的拓扑，甚至可以根据需要扩展到更多并行维度的进一步扩展。

2D 并行的性能

我们的排名分区策略通过在同一个计算节点内战略性地放置每个分片的数据副本，优化了通信模式。这种架构为权重同步操作提供了显著的性能优势。在反向传播之后，我们执行全局归约操作以同步模型权重——考虑到我们必须通信和同步的大量参数数量，这是一个昂贵的流程——在我们的设置中，我们将副本放置在同一节点上，利用节点内的高带宽，而不是依赖于较慢的节点间带宽。

这种设计选择对其他通信集体的影响通常可以降低延迟。这种改进源于两个因素。

  1. 通过将嵌入表分片到较少的排名上，并在较小的组内进行模型通信，我们实现了更低的点到点延迟。
  2. 在 2D 并行复制中，我们的嵌入查找延迟在排名上减少，我们可以将本地批次大小减少到等效全局批次大小的 1/N，其中 N 是模型副本的数量。

生产模型追踪示例说明了这两个因素，在这里我们在 1024 个 GPU 上运行 2D 并行作业，分片组大小为 256 个 GPU。

State diagram

图 2：比较非 2D 并行和 2D 并行工作负载的延迟

用户要调整以最大化其工作负载性能的两个关键杠杆是：

  1. 模型分片组大小相对于全局世界大小。全局世界大小除以分片组大小表示我们将拥有的模型副本数量。
    1. 为了最大化性能，用户可以将模型扩展至 8 倍，这种扩展因子可以保持主机内所有 reduce 操作。
      1. 对于进一步的扩展，所有 reduce 操作需要在主机间进行。根据我们的实验，我们没有看到明显的性能下降，实际上还发现了一些主机间 all reduce 的优势。我们可以更改我们的分片和副本拓扑结构，以实现主机间 all reduce，这有助于我们在特定主机出现故障时引入容错策略。
  2. 所有 reduce 同步的频率，DMPCollection 提供了 sync()调用，可以调整为每 N 个训练步骤调用一次，执行一种局部 SGD 训练。随着规模的扩大，减少同步频率可以显著提高性能。

未来工作

读者请注意，2D 稀疏并行训练与非并行训练不同，因为我们同步的是嵌入表权重而不是梯度。这种做法得益于 TorchRec 对 FBGEMM 的使用，它底层提供了优化的内核。FBGEMM 的关键优化之一是在反向传播中融合优化器。我们不是完全实现嵌入表梯度——这会消耗大量内存——而是直接将它们传递给优化器更新。尝试实现和同步这些梯度将产生大量开销，使得这种方法不切实际。

我们的探索发现，为了达到与基线相当的训练结果，我们按延迟计划同步优化器状态，时间取决于分片/副本组的数量（例如：对于 Adagrad，我们通过一个同步步骤更新其后的动量）。这种方法还允许用户实现本地 SGD 或半同步训练策略，这些策略可以实现收敛，并可能产生比基线更好的损失曲线。

感谢您阅读我们的帖子！这是我们遇到的一个令人兴奋的方向，我们希望进一步发展以最大化推荐系统的性能并推动技术前沿。