通过分层 SGD 在 PyTorch DDP 中缓解掉队者问题

作者：Yi Wang (Cruise AI), Rohan Varma (Meta AI)

PyTorch DDP 在业界分布式训练中已被广泛采用，它默认运行同步 SGD，在每个步骤同步模型副本间的梯度。这项技术的性能对于模型探索期间的快速迭代以及节省资源和成本至关重要。性能对于模型开发和探索的快速迭代和成本节约至关重要。为了解决大规模训练中慢速节点引入的普遍存在的性能瓶颈，Cruise 和 Meta 共同开发了一种基于分层 SGD 算法的解决方案，以在存在这些掉队者的情况下显著加速训练。

为何需要缓解掉队者问题

在 DDP 设置中，当一个或多个进程（“掉队者”）运行速度远慢于其他进程时，就可能发生掉队者问题。此时，所有进程在同步梯度和完成通信之前必须等待掉队者，这实质上将分布式性能瓶颈限制在最慢的工作节点上。因此，即使在训练相对较小的模型时，通信成本仍然可能是主要的性能瓶颈。

掉队者的潜在原因

严重的掉队者问题通常是由同步前工作负载不均衡引起的，许多因素都可能导致这种不均衡。例如，分布式环境中的一些数据加载工作节点可能成为掉队者，因为一些输入示例在数据大小方面可能是异常值，或者由于不稳定的网络 I/O 导致某些示例的数据传输速度大幅减慢，或者动态数据转换成本可能具有较高的方差。

除了数据加载之外，梯度同步之前的其他阶段也可能导致掉队者问题，例如在推荐系统的正向传播过程中，嵌入表查找的工作负载不均衡。

掉队者的表现

如果我们对存在掉队者的 DDP 训练任务进行性能分析，可能会发现在某个步骤，一些进程的梯度同步成本（即 allreduce 梯度）远高于其他进程。因此，即使模型非常小，分布式性能也可能被通信成本主导。在这种情况下，一些进程在某个步骤比掉队者运行得更快，因此它们必须等待掉队者，并在 allreduce 上花费更长时间。

下图显示了 PyTorch 性能分析器在一个用例中输出的两个追踪文件的截图。每个截图分析了 3 个步骤。

• 第一个截图显示，一个进程在第一步和第三步的 allreduce 成本都非常高，因为该进程比掉队者更早达到同步阶段，并花费了更多时间等待。另一方面，第二步的 allreduce 成本相对较小，这表明 1) 此步骤没有掉队者；或者 2) 此进程是所有进程中的掉队者，因此它不需要等待其他任何进程。

第一步和第三步都因掉队者而变慢

• 第二个截图显示了没有掉队者的正常情况。在这种情况下，所有的梯度同步都相对较短。

没有掉队者的正常情况

PyTorch 中的分层 SGD

最近，分层 SGD 被提出用于优化通信成本，主要通过减少大规模分布式训练中的总数据传输量，并提供了多项收敛性分析（示例）。本文的一个主要创新点是，我们在 Cruise 利用了分层 SGD 来缓解掉队者问题，这个问题在训练相对较小的模型时也可能发生。我们的实现已于 2022 年初由 Cruise 上游贡献到 PyTorch。

分层 SGD 如何工作？

顾名思义，分层 SGD 将所有进程按不同层级组织成组，并按照以下规则运行同步：

• 同一层级的所有组拥有相同数量的进程，并且这些组中的进程以相同的频率并行同步，其中同步周期由用户预先定义。
• 组的层级越高，使用的同步周期越大，因为同步成本更高。
• 当多个重叠的组根据其周期需要同步时，为了减少冗余同步并避免组间的数据竞争，只有最高层级的组执行同步。

下图展示了一个包含 16 个进程（分布在 8 台机器上，每台机器有 2 个 GPU）的 4 层分层 SGD 示例：

1. 第 1 层： 每个进程在本地运行 mini-batch SGD；
2. 第 2 层： 跨 2 台机器的每个 4 进程组每隔 2 步运行同步；
3. 第 3 层： 跨 4 台机器的每个 8 进程组每隔 4 步运行同步；
4. 第 4 层： 跨 8 台机器的所有 16 个进程的全局进程组每隔 8 步运行同步。

特别地，当步数可以被 8 整除时，只执行第 3 层的同步；当步数可以被 4 整除但不能被 8 整除时，只执行第 2 层的同步。

直观上，分层 SGD 可以看作是本地 SGD 的扩展，本地 SGD 只有两层结构——每个进程在本地运行 mini-batch SGD，然后以一定的频率进行全局同步。这也有助于解释，就像本地 SGD 一样，分层 SGD 同步的是模型参数而不是梯度。否则，当频率大于 1 时，梯度下降将是数学上不正确的。

为什么分层 SGD 可以缓解掉队者问题？

这里的关键在于，当出现随机掉队者时，它只直接减慢相对较小的一组进程的速度，而不是所有进程。下次出现另一个随机掉队者时，很可能会减慢另一组较小的进程，因此分层结构有助于平滑掉队者效应。

下面的例子假设在总共 8 个进程中，每一步都有一个随机掉队者。经过 4 步后，运行同步 SGD 的普通 DDP 将被掉队者减慢 4 次，因为它在每一步都运行全局同步。相比之下，分层 SGD 在前两步与 4 个进程的组运行同步，然后在接下来的两步运行一次全局同步。我们可以看到，前两次和后两次掉队者有很大的重叠，因此性能损失可以得到缓解。

本质上，这个分层 SGD 示例的缓解效果实际上介于每 2 步同步一次的本地 SGD 和每 4 步同步一次的本地 SGD 之间。分层 SGD 相对于本地 SGD 的主要优势在于，在相同的全局同步频率下具有更好的收敛效率，因为分层 SGD 允许更多低层级同步。此外，对于分层 SGD 而言，有可能在模型性能相当的情况下提供比本地 SGD 更低的全局同步频率，从而带来更高的训练性能，特别是在大规模分布式训练中。

易用性

掉队者缓解在分布式训练中并非一项全新的研究。已经提出了多种方法，例如 gossip SGD、数据编码、梯度编码，以及一些专门为参数服务器架构设计的方法，包括备用工作节点和陈旧同步并行（SSP）。然而，据我们所知，在此次工作之前，我们尚未找到一个优秀的开源 PyTorch 实现来缓解掉队者问题，并能像插件一样集成到我们在 Cruise 的训练系统中。相比之下，我们的实现只需要最小的改动——无需修改现有代码或调整任何现有超参数。这对工业用户来说是一个非常有吸引力的优势。

如下面的代码示例所示，只需在 DDP 模型的设置中添加几行代码，训练循环代码可以保持不变。如前所述，分层 SGD 是本地 SGD 的扩展形式，因此其启用方式与本地 SGD 非常相似（参见 PyTorch PostLocalSGDOptimizer 文档）。

1. 注册一个后本地 SGD 通信钩子，运行完全同步 SGD 的预热阶段，然后推迟分层 SGD。
2. 创建一个后本地 SGD 优化器，它封装了一个现有的本地优化器和一个分层 SGD 配置。

import torch.distributed.algorithms.model_averaging.hierarchical_model_averager as hierarchicalSGD
from torch.distributed.algorithms.ddp_comm_hooks.post_localSGD_hook import (
    PostLocalSGDState,
    post_localSGD_hook,
)
from torch.distributed.optim import PostLocalSGDOptimizer

ddp_model = nn.parallel.DistributedDataParallel(
    module=model,
    device_ids=[rank],
)

# Register a post-local SGD communication hook for the warmup.
subgroup, _ = torch.distributed.new_subgroups()
state = PostLocalSGDState(subgroup=subgroup, start_localSGD_iter=1_000)
ddp_model.register_comm_hook(state, post_localSGD_hook)

# Wraps the existing (local) optimizer to run hierarchical model averaging.
optim = PostLocalSGDOptimizer(
  optim=optim,
  averager=hierarchicalSGD.HierarchicalModelAverager(
    # The config runs a 4-level hierarchy SGD among 128 processes:
    # 1) Each process runs mini-batch SGD locally;
    # 2) Each 8-process group synchronize every 2 steps;
    # 3) Each 32-process group synchronize every 4 steps;
    # 4) All 128 processes synchronize every 8 steps.
    period_group_size_dict=OrderedDict([(2, 8), (4, 32), (8, 128)]),
    # Do not run hierarchical SGD until 1K steps for model parity.
    warmup_steps=1_000)
)

算法超参数

分层 SGD 有两个主要超参数：period_group_size_dict 和 warmup_steps。

• period_group_size_dict 是一个有序字典，映射同步周期到进程组大小，用于在分层结构中初始化不同大小的进程组，以便并行同步参数。较大的组预期使用较大的同步周期。
• warmup_steps 指定作为预热阶段的步数，用于在分层 SGD 之前运行同步 SGD。与后本地 SGD 算法类似，通常建议设置预热阶段以获得更高的精度。该值应与注册 post_localSGD_hook 时 PostLocalSGDState 中使用的 start_localSGD_iter 参数相同。通常，预热阶段至少应覆盖训练开始时损失急剧下降的阶段。

PyTorch 实现与相关论文提出的初始设计之间的一个细微区别是，预热阶段之后，默认情况下，每个主机内的进程仍会在每一步运行主机内梯度同步。这是因为

1. 主机内通信成本相对较低，并且通常可以显著加速收敛；
2. 主机内组（对于大多数工业用户而言大小为 4 或 8）通常是分层 SGD 中同步最频繁的最小进程组的良好选择。如果同步周期为 1，则梯度同步比模型参数同步（即模型平均）更快，因为 DDP 会自动重叠梯度同步和反向传播。

可以通过在 PostLocalSGDState 中取消设置 post_local_gradient_allreduce 参数来禁用此类主机内梯度同步。

演示

现在我们将演示分层 SGD 如何通过缓解掉队者问题来加速分布式训练。

实验设置

我们比较了分层 SGD、本地 SGD 和同步 SGD 在 ResNet18（模型大小：45MB）上的性能。由于模型非常小，训练不会受到同步期间数据传输成本的瓶颈限制。为了避免从远程存储加载数据引起的干扰，输入数据是从内存中随机模拟的。我们改变了训练使用的 GPU 数量，范围从 64 到 256。每个工作节点的批量大小为 32，训练迭代次数为 1,000。由于我们在本组实验中不评估收敛效率，因此未启用预热阶段。

我们还在 128 和 256 个 GPU 上以 1% 的比率、在 64 个 GPU 上以 2% 的比率模拟了掉队者，以确保平均每一步至少有一个掉队者。这些掉队者随机出现在不同的 CUDA 设备上。每个掉队者除了正常的每步训练时间（在我们的设置中约为 55 毫秒）外，会暂停 1 秒。这可以被视为一种实际场景，即训练期间 1% 或 2% 的输入数据在数据预处理成本（I/O 和/或动态数据转换）方面是异常值，并且此类成本比平均成本高出 20 倍以上。

下面的代码片段展示了如何在训练循环中模拟掉队者。我们将其应用于 ResNet 模型，它也可以轻松应用于其他模型。

     loss = loss_fn(y_pred, y)
     # Emulate a straggler that lags for 1 second at a rate of 1%.
     if random.randint(1, 100) == 1:
         time.sleep(1)
     loss.backward()
     optimizer.step()

实验是在 us-central1 GCP 集群上进行的。每台机器配备 4 个 NVIDIA Tesla T4 GPU，每个 GPU 有 16 GB 显存，通过 32 Gbit/s 以太网连接。每个实例还配备 96 个 vCPU，360 GB RAM。

我们对本地 SGD 和分层 SGD 使用了多种配置。本地 SGD 分别每隔 2、4 和 8 步运行一次全局同步。

我们使用以下配置运行了分层 SGD：

1. 在 64 个 GPU 上
   1. 每个 8 进程组、32 进程组和全局 64 进程组分别每隔 2、4 和 8 步进行同步。表示为“HSGD 2-8,4-32,8-64”。
   2. 每个 32 进程组和全局 64 进程组分别每隔 4 和 8 步进行同步。表示为“HSGD 4-32,8-64”。
2. 在 128 个 GPU 上
   1. 每个 8 进程组、32 进程组和全局 128 进程组分别每隔 2、4 和 8 步进行同步。表示为“HSGD 2-8,4-32,8-128”。
   2. 每个 32 进程组和全局 128 进程组分别每隔 4 和 8 步进行同步。表示为“HSGD 4-32,8-128”。
3. 在 256 个 GPU 上
   1. 每个 4 进程组、16 进程组、64 进程组和全局 256 进程组分别每隔 1、2、4 和 8 步进行同步。表示为“HSGD 1-4,2-16,4-64,8-256”。
   2. 每个 8 进程组、64 进程组和全局 256 进程组分别每隔 2、4 和 8 步进行同步。表示为“HSGD 2-8,4-64,8-256”。
   3. 每个 16 进程组和全局 256 进程组分别每隔 4 和 8 步进行同步。表示为“HSGD 4-16,8-256”。

实验结果

下图显示了在模拟掉队者存在的情况下，不同通信方案相对于同步 SGD 基线的加速比。我们可以得出以下观察结果：

1. 正如预期，我们可以看到分层 SGD 和本地 SGD 都可以通过降低同步频率实现更高的加速比。
2. 分层 SGD 方案在 64 个 GPU 上的加速比分别为 2.08X-2.45X，在 128 个 GPU 上为 2.57X-2.68X，在 256 个 GPU 上为 2.63X-3.25X。这表明分层 SGD 可以显著缓解掉队者问题，并且在更大规模下，这种缓解效果可以更有效。
3. 同步周期为 2 步和 8 步的本地 SGD 的性能可以分别被视为实验的分层 SGD 方案的下限和上限。这是因为分层 SGD 方案的全局同步频率低于每 2 步一次，但与每 8 步一次的全局同步相比，它们在小规模组上的低层级同步是额外的开销。

总的来说，分层 SGD 可以提供比本地 SGD 更精细的通信成本和模型质量之间的权衡。因此，当同步周期相对较大的本地 SGD（如 8 或 4）无法提供令人满意的收敛效率时，分层 SGD 更有可能实现良好的加速比和模型性能相当。

由于实验中只使用了模拟数据，我们在此未展示模型性能相当性，但在实践中可以通过两种方式实现：

1. 调整超参数，包括分层结构和预热步数；
2. 在某些情况下，对于相同数量的训练步数，分层 SGD 可能会导致模型质量略低于原始模型（即收敛速度较慢），但每训练步的加速比达到 2 倍以上，仍然可以通过增加步数来达到模型性能相当，同时总训练时间更少。

局限性

在将分层 SGD 应用于掉队者缓解之前，用户应该注意此方法的一些局限性：

1. 此方法只能缓解非持续性掉队者，这些掉队者在不同时间出现在不同的工作节点上。然而，对于持续性掉队者的情况，这可能是由特定主机上的硬件故障或网络问题引起的，这些掉队者每次都会减慢同一个低层级子组的速度，几乎无法实现掉队者缓解。
2. 此方法只能缓解低频率的掉队者。例如，如果每一步有 30% 的工作节点可能随机成为掉队者，那么大多数低层级同步仍然会因掉队者而变慢。因此，分层 SGD 可能无法显示出相对于同步 SGD 的明显性能优势。
3. 由于分层 SGD 应用的是模型平均，它不像普通 DDP 使用的梯度平均那样与反向传播重叠，因此其掉队者缓解带来的性能增益必须大于通信与反向传播之间没有重叠造成的性能损失。因此，如果掉队者仅导致训练速度降低不到 10%，分层 SGD 可能无法带来多少加速比。未来可以通过重叠优化器步骤和反向传播来解决这一局限性。
4. 由于对分层 SGD 的研究不如本地 SGD 充分，因此无法保证具有更细粒度同步的分层 SGD 能比某些高级形式的本地 SGD 收敛得更快，例如SlowMo，它可以通过慢动量提高收敛效率。然而，据我们所知，这些高级算法目前还不能像分层 SGD 那样原生支持作为 PyTorch DDP 插件。

致谢

我们衷心感谢 Cruise 团队成员 Bo Tian, Sergei Vorobev, Eugene Selivonchyk, Tsugn-Hsien Lee, Dan Ring, Ian Ackerman, Lei Chen, Maegan Chew, Viet Anh To, Xiaohui Long, Zeyu Chen, Alexander Sidorov, Igor Tsvetkov, Xin Hu, Manav Kataria, Marina Rubtsova, 和 Mohamed Fawzy，以及 Meta 团队成员 Shen Li, Yanli Zhao, Suraj Subramanian, Hamid Shojanzeri, Anjali Sridhar 和 Bernard Nguyen 的支持。