通过分层SGD缓解 PyTorch DDP 中的慢节点问题 – PyTorch

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

慢节点缓解的需求

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

慢节点的潜在原因

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

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

慢节点的出现

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

下图显示了 PyTorch 分析器在某个用例中输出的两个跟踪文件的屏幕截图。每个屏幕截图都分析了 3 个步骤。

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

第一步和第三步都因慢节点而变慢

第二个屏幕截图显示了没有慢节点的正常情况。在这种情况下，所有梯度同步都相对较短。

chart showing normal case without stragglers

没有慢节点的正常情况

PyTorch 中的分层 SGD

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

分层 SGD 如何工作？

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

同一级别的所有组具有相同数量的进程，并且这些组中的进程以相同的频率并发同步，其中同步周期由用户预定义。
组级别越高，使用的同步周期越长，因为同步变得越昂贵。
当多个重叠组应根据其周期进行同步时，为减少冗余同步并避免组间数据竞争，仅运行最高级别组的同步。

下图展示了 8 台机器上的 16 个进程（每台机器有 2 个 GPU）的 4 级分层 SGD 示例

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

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

An example of 4-level hierarchy SGD among 16 processes on 8 machines, each of which has 2 GPUs

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

为什么分层 SGD 可以缓解慢节点问题？

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

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

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

易用性

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

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

注册一个局部 SGD 后通信钩子，以运行完全同步 SGD 的预热阶段并推迟分层 SGD。
创建一个局部 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 实现与相关论文中提出的最初设计之间存在一个细微的差异，即在预热阶段之后，默认情况下，每个主机内的进程仍会在每一步运行主机内梯度同步。这是因为：

主机内通信相对便宜，通常可以显著加速收敛；
主机内组（对于大多数行业用户来说，大小为 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。

架构	ResNet18 (45MB)
工作节点	64, 128, 256
后端	NCCL
GPU	Tesla T4，16 GB 内存
批次大小	32 x ## 工作节点数
慢节点持续时间	1 秒
慢节点率	128 和 256 个 GPU 上为 1%，64 个 GPU 上为 2%

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

我们运行了以下配置的分层 SGD

在 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”。
在 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”。
在 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 基线（在模拟慢节点下）的加速比。我们可以得出以下观察结果：

正如预期的那样，我们可以看到分层 SGD 和局部 SGD 都可以通过较低的同步频率实现更高的加速。
在 64 个 GPU 上，分层 SGD 方案的加速比分别为 2.08X-2.45X；在 128 个 GPU 上为 2.57X-2.68X；在 256 个 GPU 上为 2.63X-3.25X。这表明分层 SGD 可以显著缓解慢节点问题，并且这种缓解在大规模情况下可能更有效。
同步周期为 2 步和 8 步的局部 SGD 的性能可以被视为实验分层 SGD 方案的下限和上限。这是因为分层 SGD 方案的全局同步频率低于每 2 步，但其小规模组的低级同步与每 8 步全局同步相比是额外的开销。

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

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

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

局限性

在将分层 SGD 应用于慢节点缓解之前，用户应该注意这种方法的一些局限性：

此方法只能缓解非持久性慢节点，即在不同时间发生在不同工作节点上的慢节点。但是，对于持久性慢节点，这可能是由特定主机上的硬件退化或网络问题引起的，这些慢节点将每次都拖慢相同的低级子组，导致几乎无法缓解慢节点问题。
这种方法只能缓解低频慢节点。例如，如果 30% 的工作节点在每一步都可能随机成为慢节点，那么大多数低级同步仍然会被慢节点拖慢。因此，分层 SGD 可能不会显示出比同步 SGD 明显的性能优势。
由于分层 SGD 采用的模型平均不与反向传播重叠，而香草 DDP 使用的梯度平均则会重叠，因此其慢节点缓解的性能提升必须超过通信和反向传播之间没有重叠所造成的性能损失。因此，如果慢节点只导致训练速度下降不到 10%，分层 SGD 可能无法带来多少加速。这个限制可以通过未来优化器步进和反向传播的重叠来解决。
由于分层 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 的支持。

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