基于分层 SGD 的 PyTorch DDP 的流浪者缓解

由王毅（Cruise AI）、罗汉·瓦拉（Meta AI）撰写

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

需要解决拖油瓶问题

在 DDP（分布式数据并行）设置中，当其中一个或多个进程运行速度明显慢于其他进程（称为“落后者”）时，可能会出现落后者问题。当这种情况发生时，所有进程都必须等待落后者完成同步梯度并完成通信，这实际上将分布式性能瓶颈限制在速度最慢的工人上。因此，即使是训练相对较小的模型，通信成本仍然可能成为性能瓶颈。

落后者的潜在原因

严重的落后者问题通常是由同步前的负载不平衡引起的，许多因素可能导致这种不平衡。例如，在分布式环境中，一些数据加载器工人可能会成为落后者，因为某些输入示例在数据大小方面可能是异常值，或者由于不稳定的网络 I/O，某些示例的数据传输可能会大幅减慢，或者实时数据转换的成本可能会有很大的差异。

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

逃兵的出现

如果我们对有逃兵的 DDP 训练作业进行性能分析，我们可以发现，在某些步骤中，一些进程的梯度同步成本（即 allreduce 梯度）可能比其他进程高得多。因此，即使模型规模很小，分布式性能也可能被通信成本所主导。在这种情况下，一些进程在某个步骤中比逃兵运行得更快，因此它们必须等待逃兵，并在 allreduce 上花费更长的时间。

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

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

第 1 步和第 3 步因掉队者而减慢

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

无掉队者的正常情况

PyTorch 中的分层 SGD

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

层次化 SGD 是如何工作的？

如其名所示，层次化 SGD 将所有进程组织成不同级别的组，并按照以下规则进行同步：

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

下图展示了在 8 台机器上 16 个进程的 4 级层次 SGD 示例，每台机器有 2 个 GPU：

1. 第 1 级：每个进程在本地运行小批量 SGD；
2. 级别 2：每 4 进程组跨 2 台机器每 2 步进行同步；
3. 级别 3：每 8 进程组跨 4 台机器每 4 步进行同步；
4. 级别 4：所有 16 个进程在 8 台机器上的全局进程组每 8 步进行同步。

特别地，当步数能被 8 整除时，仅执行 3）中的同步，当步数能被 4 整除但不能被 8 整除时，仅执行 2）中的同步。

直观地看，层次化 SGD 可以看作是局部 SGD 的扩展，它只有两级层次——每个进程都本地运行小批量 SGD，然后在一定频率下全局同步。这也可以解释为什么，就像局部 SGD 一样，层次化 SGD 同步的是模型参数而不是梯度。否则，当频率大于 1 时，梯度下降在数学上将是错误的。

为什么层次化 SGD 可以减轻拖沓现象？

这里的关键洞察是，当存在随机拖沓者时，它只会直接减缓相对较小的一组进程，而不是所有进程。下次另一个随机拖沓者很可能减缓另一组不同的进程，因此层次结构可以帮助平滑拖沓效应。

下面这个例子假设在每一步中，总共有 8 个进程中存在一个随机落后的进程。经过 4 步后，运行同步 SGD 的 vanilla DDP 会因为落后进程而慢下来 4 次，因为它在每一步都进行全局同步。相比之下，分层 SGD 在前两步后对 4 个进程的组进行同步，然后在另外两步后进行全局同步。我们可以看到，前两步和最后两步的落后进程有较大的重叠，因此性能损失可以得到缓解。

实际上，这个分层 SGD 示例的缓解效果实际上是在每 2 步和每 4 步的本地 SGD 频率之间。分层 SGD 相对于本地 SGD 的主要优势是相同的全局同步频率下更好的收敛效率，因为分层 SGD 允许更多的低级同步。此外，分层 SGD 还可以提供低于本地 SGD 的全局同步频率，同时保持模型等价，从而提高训练性能，尤其是在大规模分布式训练中。

易用性

追踪者缓解并不是分布式训练中的新研究。已经提出了多种方法，例如八卦 SGD、数据编码、梯度编码，以及一些特别为参数服务器架构设计的，包括备份工作者和陈旧的同步并行。然而，据我们所知，在我们这项工作之前，我们没有找到一种好的开源 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 arg 相同。通常，预热阶段至少应覆盖训练开始时损失急剧下降的阶段。

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 设备上。除了正常的每步训练时间（在我们的设置中约为 55ms）外，每个落后者还会停滞 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 有 16GB 内存，通过 32 Gbit/s 以太网连接。每个实例还具备 96 个 vCPU 和 360GB 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 更细粒度的权衡。因此，当在相对较大的同步周期（如 8 或 4）下，本地 SGD 无法提供令人满意的收敛效率时，分层 SGD 有更大的机会实现良好的加速和模型等价。

由于实验中只使用了模拟数据，我们没有在这里展示模型等价，但在实际中可以通过以下两种方式实现：

1. 调整包括层次结构和预热步骤在内的超参数；
2. 对于某些情况，分层 SGD 可能导致与相同训练步数相比，原始模型的质量略有下降（即较低的收敛速度），但每步训练的加速达到 2X+，仍然可以通过更多的步骤实现模型等价，但总的训练时间仍然更少。

局限性

在应用分层 SGD 进行慢速节点缓解之前，用户应该了解这种方法的几个局限性：

1. 这种方法只能缓解非持续性的慢速节点，这些慢速节点在不同时间会出现在不同的工作者上。然而，对于由硬件退化或特定主机上的网络问题引起的持续性慢速节点，这些慢速节点会每次都减慢相同的低级子组，导致几乎无法缓解慢速节点。
2. 这种方法只能缓解低频慢速节点。例如，如果 30%的工作者在每一步都可能随机成为慢速节点，那么大多数低级同步仍然会被慢速节点减慢。因此，分层 SGD 可能不会显示出比同步 SGD 明显的性能优势。
3. 由于分层 SGD 应用了与 vanilla 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 的支持。