[源码解析] PyTorch 分布式(8) -------- DistributedDataParallel之论文篇

[源码解析] PyTorch 分布式(8) -------- DistributedDataParallel之论文篇
- 0x00 摘要
- 0x01 原文摘要
- 0x02 引论
  - 2.1 挑战
  - 2.2 实现和评估
- 0x03 背景
  - 3.1 PyTorch
  - 3.2 数据并行
  - 3.3 AllReduce
- 0x04 系统设计
  - 4.1 API
  - 4.2 梯度规约
    - 4.2.1 A Naive Solution
    - 4.2.2 Gradient Bucketing
    - 4.2.3 Overlap Computation with Communication
    - 4.2.4 Gradient Accumulation
  - 4.3 Collective Communication
- 0x05 实施
  - 5.1 Python前端
  - 5.2 Core Gradient Reduction
- 0xFF 参考

0x00 摘要

工欲善其事，必先利其器，为了更好的分析代码，我们先来学习一下相关论文。

PyTorch 开发者在实现的同时，发布了一篇论文：[ PyTorch Distributed: Experiences on Accelerating Data Parallel Training ] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, Soumith Chintal。

其地址为：http://www.vldb.org/pvldb/vol13/p3005-li.pdf

因为论文较长，所以本文翻译其思路和实现之中的部分内容，在后文之中将以这篇论文为基础，结合源码来进行分析。本文不完全按照原论文的顺序进行翻译，笔者会对其重点做标注，也会按照自己的理解进行调整，另外，原文是基于 PyTorch 1.5，与最新 PyTorch 有部分出入。

0x01 原文摘要

深度学习的最新进展证明了大型数据集和大型模型的价值，这就需要将模型训练扩展到更多计算资源之上。由于其简单的原理和广泛的适用性，数据并行已成为分布式培训的一种流行解决方案。通常，分布式数据并行技术在每个计算源上复制模型以在每个worker之上独立地生成梯度，然后在每次迭代中通信这些梯度以保持模型副本的一致性。尽管该技术概念简单，但计算和通信之间的微妙依赖性使得优化分布式训练效率非常重要。从1.5版开始，Pytorch 提供了几种加速分布式数据并行的技术，包括bucketing梯度、通信重叠计算和跳过梯度同步。评估表明，当适当配置时，Pyrotch分布式数据并行模块可使用256 GPU实现近似线性的可扩展性。

0x02 引论

训练DNN模型通常重复执行以下三个步骤：

向前传递以计算损失。
向后传播以计算梯度。
以及优化器步骤以更新参数。

数据并行性的概念普遍适用于此类框架：应用程序可以创建一个模型的多个副本，每个模型副本处理一部分训练数据，并独立执行向前和向后传播。之后，模型副本可以根据算法同步其梯度或更新的参数。

2.1 挑战

看起来，完全在应用程序端构建数据并行的工作版本是可能的，因为它只需要在每次迭代中插入适当的通信。然而，挤出最后一点性能需要在设计和调整方面付出巨大的努力。在平台端提供本机分布式数据并行API将帮助应用程序开发人员专注于优化其模型，而平台开发团队可以持续透明地提高训练速度。

要提供一个通用的分布式数据并行包，有三个方面的挑战。

数学等价：数据并行的目的是加速对大型数据集的训练。应用程序希望获得相同的结果模型，就好像所有培训都是在本地进行，没有模型复制一样。这就要求尽管它是分布式训练，但是应该数学等价于本地训练。
非侵入式和拦截式API：应用程序开发通常从本地模型开始，然后在必要时扩展。所以需要有一个从本地模型开始，修改代码以适应分布式的过程。
- 为了避免这个从本地模型到分布式模型的过渡期间太过麻烦，API在应用程序代码中必须是非侵入性的。
- 另一方面，API也需要允许一个内部实现来及时截获各种信号，以便执行通信和系统优化。
高性能：数据并行训练受制于计算和通信之间微妙的依赖关系。设计和实现必须探索解决方案空间，以有效地将更多资源转换为更高的训练吞吐量。

2.2 实现和评估

PyTorch以nn.Module类的形式提供分布式数据并行，其中应用程序在构建时以子模块的形式提供其模型。为了保证数学等效性，所有副本都从相同的模型参数初始值开始，并同步梯度，以便在整个训练迭代中保持参数一致。为了最大限度地降低集成度，该实现（分布式数据并行模型）暴露了与用户模型相同的forward API，这允许应用程序无缝地用分布式数据并行模型对象替换之前出现的用户模型，而无需额外的代码更改。设计中集成了多种技术，以提供高性能培训，包括bucketing gradients，与计算的重叠通信和跳过同步。

评估是在一个专用的32 GPU集群和一个更大的共享权限中的256 GPU上进行的。我们开发了基准程序来评估不同规模的分布式包，以深入了解不同优化技术和配置的性能影响。实验还包括NCCL和Gloo通信库之间的比较。结果表明：

通信是影响训练延迟的主要因素，其影响随模型尺寸的增大而增大；
存储桶大小对通信效率有很大影响，如果配置正确，可能会导致2倍以上的加速；
适当跳过同步将显著减少分摊的通信开销，而不会显著降低收敛速度。

0x03 背景

3.1 PyTorch

PyTorch将值组织成张量，张量是具有丰富数据操作集的通用n维数组。模块定义了从输入值到输出值的转换，其正向传递期间的行为由其 forward 成员函数指定。模块可以包含张量作为参数。例如，线性模块包含权重参数和偏差参数，其正向函数通过将输入乘以权重并添加偏差来生成输出。

应用程序通过将本机模块（如线性、卷积等）和自定义forward函数中的Function（如relu、pool等）粘合在一起，构成自己的模块。典型的训练迭代包括使用输入和标签生成损失的前向传递，计算参数梯度的后向传递，以及使用梯度更新参数的优化器步骤。更具体地说，在向前传播过程中，PyTorch构建了一个autograd图来记录所执行的动作。然后，在后向过程中，使用autograd图进行反向传播以生成梯度。最后，优化器应用梯度来更新参数。训练过程重复这三个步骤，直到模型收敛。

3.2 数据并行

PyTorch 提供了多种工具来促进分布式训练，包括：

DataParallel，用于在同一台机器上使用多个GPU的单进程多线程进行数据并行训练。
DistributedDataParallel，用于跨GPU和机器的多进程数据并行训练。
RPC，用于一般分布式模型并行训练（例如，参数服务器）。

论文的其余部分主要关注分布式数据并行。数据并行通过在操作优化步骤之前进行梯度通信来实现分布式训练，这样可以确保使用完全相同的梯度集来更新所有模型副本的参数，因此模型副本可以在迭代中保持一致。

参数平均是扩展模型训练的另一种流行技术。类似地，它可以跨多台机器启动多个过程，但不是同步梯度，而是直接计算所有模型参数的平均值。这发生在本地优化器步骤之后，这意味着参数平均可以完全作为一个辅助步骤实现，完全不需要与本地训练步骤交互，这很有吸引力，因为它可以轻松、干净地解耦分布式训练和本地迭代的代码。但是参数平均有几个注意事项。

与局部训练相比，参数平均可产生截然不同的结果，这有时会对模型精度造成不利影响。根本原因是，参数平均在数学上并不等同于本地处理所有输入数据，尤其是当优化器依赖于过去的本地梯度值（如动量）时。由于不同的模型副本可能会看到不同的梯度，因此optimizers中的状态可能会逐渐发散，从而导致梯度下降方向冲突。当从局部优化模型切换到大规模部署模型时，这可能会导致性能上莫名其妙的差异。
参数平均的结构将计算（即反向传递）和通信（即计算平均值）协调到非重叠阶段，使用optimizer step() 函数作为硬分离点。无论我们如何大力优化计算或通信，一种类型的资源在任何给定时间都将处于空闲状态，从而放弃大量性能优化机会。

鉴于上述基本缺陷，我们决定使用数据并行性来同步梯度而不是参数来实施分布式训练。请注意，应用程序仍然可以使用PyTorch轻松构建参数平均值。事实上，后文中描述的集合通信特性是该用例的合适解决方案。应用程序只需要显式地启动AllReduce操作来相应地计算平均参数。

3.3 AllReduce

AllReduce是一个基础通信API，其被 DistributedDataParallel 用于计算所有进程的梯度求和。

多个通信库都提供了AllReduce ，包括NCCL、Gloo和MPI。AllReduce操作要求每个参与进程都提供一个大小相等的张量，然后将给定的算术运算（如sum、prod、min、max）应用于所有进程的输入张量，并向每个参与者返回相同的结果张量。

一个 AllReduce 简单的实现可以简单地让每个进程向所有对等进程广播其输入张量，然后独立地应用算术运算。然而，由于AllReduce对分布式训练速度有显著影响，通信库实现了更复杂、更高效的算法，如基于环的AllReduce和基于树的AllReduce。由于一个AllReduce操作在所有进程加入之前无法启动，因此它被认为是一种同步通信，而不是参数服务器中使用的P2P通信。

0x04 系统设计

PyTorch 提供了分布式数据并行（DDP）模块，这有助于轻松地跨多个进程和机器来进行并行化训练。在分布式培训期间，每个流程都有自己的本地模型副本和本地优化器。就正确性而言，分布式数据并行训练和本地训练必须在数学上等价。DDP可以通过如下来确保训练正确性：

所有模型副本从完全相同的模型状态开始，并在每次向后传播之后，得到相同的参数梯度。
因此，即使来自不同流程的优化器都是独立的，它们也应该能够在每次迭代结束时将其本地模型副本置于相同的状态

下图示出了DDP的构建块，它包含Python API前端、C 梯度归并核心算法，并使用 c10d 集合通信库。以下部分按此堆栈图的自上而下顺序显示。

第4.1节介绍了推动DDP API设计的一般原则。第4.2节介绍Pyrotch分布式数据并行包中使用的扩展梯度归并技术。最后，第4.3节讨论了DDP的集合通信后端选项。

4.1 API

在设计API时，我们定义了两个设计目标来实现必要的功能。

非侵入性：API必须对应用程序是非侵入的。应用程序开发人员通常从编写本地培训脚本开始，并在单个计算机上达到资源限制时扩展。在这一点上，要求开发人员重写整个应用程序以支持分布式数据并行训练是不可接受的。相反，开发人员应该能够通过最少的修改来重用本地训练脚本。
拦截：API需要允许实现拦截各种信号以便及时触发适当的算法。分布式数据并行旨在通过使用更多的计算资源来加速训练。这一过程需要在计算和通信方面进行微妙的优化，以实现最佳性能。因此，API必须对内部实现提供尽可能多的优化机会。

鉴于上述要求，我们将分布式数据并行实现为一个nn 模块，该模块将本地模型作为构造函数参数，并透明地同步反向过程中的数据。下面的代码片段显示了使用DDP模块的示例。

本例使用nn.Linear层在第10行创建局部模型。
然后，它在第11行将本地模型转换为分布式训练模型，并在第12行设置优化器。
第14行到第23行是典型的前向传播、后向传播和优化器步骤实现。

在这个玩具分布式培训示例中，第11行是将本地培训应用程序转换为分布式应用程序的唯一区别，它满足了非侵入性需求，还满足交互要求。构造器允许DDP检查模型结构和参数。构造完成后，本地模型将被分布式模型替换，然后分布式模型可以很容易地拦截forward()调用以执行相应的必要操作。对于向后传播，DDP依靠向后钩子触发梯度规约，即，在损失张量上执行backward()时，autograd引擎将执行梯度规约。

4.2 梯度规约

DDP中的梯度规约算法在过去的版本中有所发展。为了介绍当前实现的结构，让我们从一个简单的解决方案开始，逐步引入更多的复杂性，并在PyTorch v1.5.0中使用当前版本。这还将解释第 4.1节中描述的相同简单API如何允许我们安装各种性能优化算法。

4.2.1 A Naive Solution

如第4节开头所述，DDP通过让所有训练过程（1）从相同的模型状态开始，以及（2）在每次迭代中使用相同的梯度，来保证正确性。前者可以通过在DDP构建时将模型状态从一个进程广播到所有其他进程来实现。为了实现后者，一个简单的解决方案是：可以在本地向后传播之后和更新本地参数之前插入梯度同步阶段。但是，第4.1节中显示的API没有为此阶段提供明确的入口点，因为backward()和step()之间没有任何内容。幸运的是，PyTorch autograd引擎接受定制的后向hook。DDP可以注册autograd钩子，以在每次向后传播后触发计算。钩子函数被激发时，每个钩子扫描所有局部模型参数，并从每个参数检索梯度张量。然后，它使用AllReduce 集合通信操作来计算所有进程中每个参数的平均梯度，并将结果写回梯度张量。

Naive Solution 工作正常，但存在两个性能问题：

集合通信在小张量上表现不佳，这在具有大量小参数的大型模型上尤为突出。
将梯度计算和同步分离会因为两者之间的硬边界而丧失计算与通信重叠的机会。

4.2.2 Gradient Bucketing

梯度bucketing的思想是基于这样一个观察，即集合通信在大张量上更有效。下图（a）和（b）提供了定量视图，显示了AllReduce 60M torch.float32的总执行时间。每个AllReduce的参数数量不同。

为了最大限度地提高带宽利用率，所有reduce操作都是异步启动的，并同时阻塞等待所有操作，以便模仿DDP的梯度归并算法。实验在一台支持NVLink[3]的服务器上进行，该服务器带有两个NVIDIA Quadro GP100 GPU。NCCL AllReduce直接在CUDA输入张量上运行，而Gloo AllReduce则在CPU输入张量上运行，以便消除在使用Gloo后端时将CUDA内存复制到CPU内存的开销。对于NCCL和Gloo，当使用较大的输入张量时，总通信时间明显减少。Gloo在每个输入张量约500K参数时达到最高速度，而NVLink上的NCCL甚至没有20M参数GPU张量的明显饱和信号。

这些实验表明，如果DDP在短时间内等待并将多个梯度存储到一个AllReduce操作中，它可以实现更高的吞吐量和更低的延迟，而不是在每个梯度存储可用时立即启动专用的AllReduce。这对于具有许多小参数的模型尤其有用。但是，DDP不应在一个AllReduce中传输所有数据，否则，在计算结束之前无法启动任何通信。上图（c）和（d）显示了包含大约60M参数的ResNet152 的GPU和CPU反向计算时间。X轴是准备好的梯度数量，Y轴是自向后传播开始以来经过的时间。GPU上的后向传播大约需要250毫秒才能完成，这与NVLink上的NCCL的数量级相同。这一结论也适用于Gloo和CPU后向传播。这些测量预示着，对于相对较小的bucket大小，DDP可以在向后传播的同时启动AllReduce操作，以使通信与计算重叠，这将改变每次迭代的延迟。

4.2.3 Overlap Computation with Communication

在梯度上的AllReduce操作可以在本地向后传播完成之前开始。使用bucketing，DDP需要等待同一个bucket中的所有内容，然后开始启动通信。

在这种设置下，只是在向后传播结束时触发AllReduce不再足够。我们需要对更频繁的信号作出反应，并更迅速地启动 AllReduce。因此，DDP为每个梯度累加器注册一个autograd hook。Hook 在其相应的累加器更新梯度之后被触发，并将检查其所属的bucket。如果相同桶中所有梯度的钩子都已触发，则最后一个钩子将触发该桶上的异步AllReduce。

有两点需要注意。

首先，所有进程的归并顺序必须相同，否则，AllReduce内容可能不匹配，导致不正确的归并结果或程序崩溃。然而，PyTorch在每次向前传播时都会动态地构建autograd图，不同的过程可能在梯度就绪顺序上不一致。下图（a）给出了一个示例，其中两个垂直轴表示时间，虚线表示梯度准备就绪的时间。在过程1中，四个梯度按顺序计算，但梯度g2在过程2的g3和g4之后计算。在这种情况下，如果所有进程都在准备就绪后立即AllReduce bucket，则AllReduce内容将不匹配。因此，所有流程必须使用相同的bucketing顺序，并且没有流程可以在装载bucket i之前就在bucket i 1上启动AllReduce。如果bucket 0是最后一个准备就绪的bucket，那么通信就不可能与计算重叠【笔者：因为 bucket 0 是最后就绪的，所以其他bucket在这之前都不会被执行计算，就不能与通信重叠了】。PyTorch v1.5.0通过使用model.parameters()的相反顺序作为bucketing顺序来解决此问题，我们做如下假设：层（layers）可能按照正向过程中调用的相同顺序进行注册。因此，其反向顺序就是反向过程中的梯度计算顺序的近似表示。诚然，这并不是一个完美的解决方案，但它是一个近似方案，我们可以用最少的工程开销来实现它。
其次，一次训练迭代可能只涉及模型中的一个子图，并且子图在每次迭代中可能不同，这意味着在某些迭代中可能会跳过某些Gradient。然而，由于 gradient-to-bucket 的映射是在构建时确定的，这些缺少的梯度将使一些bucket 永远看不到最终的自动装载hook，并且无法将bucket标记为就绪。因此，向后传播可能会暂停。下图（b）示出了一个示例，其中在一次迭代中跳过了与梯度g3相对应的参数，导致g3缺少就绪信号。为了解决这个问题，DDP从前向传播的输出张量遍历autograd图，以找到所有参与的参数。这些参与张量的准备就绪是结束向后传播完成的有效信号。因此，DDP可以通过在向前传播结束时主动标记剩余的参数梯度来避免等待。请注意，此更改并不妨碍我们开发非侵入式API，因为应用程序可以直接调用DDP上的forward函数，并且DDP可以轻松地将此步骤插入其成员函数中。

下面算法给出了DDP的伪码。

Constructor包含两个主要步骤，广播模型状态和安装autograd挂钩。DDP的 forwad 函数是本地模型 forwad 函数的简单包装器。它遍历autograd图以相应地标记未使用的参数。autograd钩子将内部参数索引作为输入，这有助于找到参数张量及其所属范围。它将局部梯度写入bucket中的正确偏移量，然后启动异步AllReduce操作。伪代码中省略了一个附加的结束步骤，它等待AllReduce操作，并在反向过程结束时将值写回梯度。

下图阐明了DDP在向前和向后传播期间如何与局部模型交互。

上述解决方案适用于大多数用例。但是，由于DDP总是计算所有梯度的平均值，并将它们写回parameter.grad字段，因此优化器无法区分梯度是否参与了最后一次向后传播。由于DDP和优化器的解耦设计，DDP没有旁侧通道向优化器暗示该信息。如果没有这些信息，训练过程可能会受到模型精度回归的影响，例如，当优化器使用梯度感知信息跳过动量值更新时。为了解决这个问题，DDP应该只接触哪些确实涉及向后传播的梯度。然而，由于在对等DDP过程中，前向/后向过程可能仍然涉及到局部缺失梯度，因此无法仅从局部autograd图中提取该信息。因此，DDP使用位图跟踪本地参数参与者，并启动另外一个AllReduce来收集全局未使用的参数。不幸的是，由于元素类型可能不匹配，DDP无法将此位图合并到其他梯度AllReduce操作中。只有当应用程序显式地告诉DDP查找未使用的参数时，这种额外的开销才会出现，因此只有在必要时才会支付代价。

4.2.4 Gradient Accumulation

加速分布式数据并行训练的一种常用技术是降低梯度同步频率。在全局同步梯度之前，应用程序可以执行n次局部训练迭代，而不是在每次迭代中启动AllReduce。如果输入批次太大而无法装入设备，这也很有帮助，因为应用程序可以将一个输入批次拆分为多个微批次，在每个微批次上运行局部向前和向后传播，并且仅在大批次的边界处启动梯度同步。理论上，这应该产生与大批量数据一次性处理相同的结果，因为梯度将简单地累积到相同的张量。然而，这在某种程度上与第 4.2.3节中讨论的梯度归并算法相冲突。该算法将在每次向前传递结束时将未使用的参数标记为就绪，而一次迭代中未使用的参数仍可以参与后续迭代。此外，DDP无法区分应用程序是否应该在向后或通过多次迭代累积梯度后立即调用optimizer.step()。因此，我们需要为这个用例引入一个额外的接口（即，no_sync ）。

在内部，no_sync 的实现非常简单。上下文管理器只是在进入和退出上下文时切换一个标志，该标志在DDP的forward 功能中使用。在 no_sync 。全局未使用参数的信息也会累积在位图中，并在下次通信发生时使用。下面是一个示例代码段。

4.3 Collective Communication

分布式数据并行训练使用一种特殊的通信模式：每个参与者提供一个相同尺寸的张量，并收集所有参与者的全局和（global sum）。这可以通过如下方式来实现：首先是一个gather操作，然后使用点对点通信对每个参与者进行局部归并（local reductions），但这将丧失性能优化的机会。

DDP构建在集合通信库之上，包括三个选项：NCCL、Gloo和MPI。DDPs从三个库中获取API，并将它们包装到同一个ProcessGroup API中。该名称预示着ProcessGroup希望多个进程作为一个组一起工作。

所有ProcessGroup实例通过使用集合服务（rendezvous service）同时构造，其中第一个实例将进行阻塞，一直等待，直到最后一个实例加入。对于NCCL后端，ProcessGroup为通信维护一组专用的CUDA流，以便通信不会阻止默认流中的计算。由于所有通信都是集合操作，因此所有ProcessGroup实例上的后续操作必须在大小和类型上匹配，并遵循相同的顺序。对所有库使用相同的ProcessGroup API允许我们使用相同的DDP实现来试验不同的通信算法。例如，PyTorch v1.5提供了一个round-robin ProcessGroup实现，它获取ProcessGroup实例列表，并以循环方式向这些ProcessGroup实例发送集合通信。通过使用round-robin ProcessGroup，在单个NCCL、Gloo或MPI处理组无法饱和链路容量的情况下，DDP可以获得更高的带宽利用率。

0x05 实施

在过去的几个版本中，DDP的实现已经演进了好几次。本节重点介绍PyTorch v1.5.0的当前状态。DDP实现同时存在于 Python和C 文件，Python 部分包括公开API和非性能关键的组件，C 提供核心梯度归并算法。Python API 通过Pybind11来调用C 核心。

5.1 Python前端

DDP nn.module在distributed.py中实现，它包含面向用户的组件。组件包括构造函数、forward 函数和 no_sync 上下文管理器。除了在第4节中强调的一般思想外，Python前端中还有几个塑造DDP行为的实现细节。

DDP构造器API中公开了Configurable Knobs，包括

process_group，用于指定DDP运行AllReduce的流程组实例，这有助于避免和默认流程组混淆；
bucket_cap_mb，用于控制AllReduce bucket大小，应用程序应调整此以优化训练速度，
find_unused_parameters，来切换DDP是否应检测未使用的参数，DDP是通过遍历autograd图来完成检测的。

本地模型中的模型设备关联性（Model Device Affinity ）也控制DDP的行为，特别是当模型跨越多个设备时，这在模型太大而无法装入单个设备时很常见。对于大型模型，应用程序可以将模型的不同层放置在不同的设备上，并使用Tensor.to(device) API将中间输出从一个设备移动到另一个设备。DDP也适用于多设备模型。只要将 device_ids参数设置为None或空列表，DDP就会检查模型，执行健全性检查并相应地应用配置。然后，将多设备模型视为一个整体。

当一个层需要跟踪运行方差和运行平均值（例如BatchNorm）等状态时，模型缓冲区（Model Buffers）是必要的。DDP通过让rank 0 的进程获得支持模型缓冲区的权限。如果模型包含缓冲区，DDP在本地模型上开始前向传递之前，将缓冲区值从rank 0进程广播到所有其他进程。此行为也与no_sync模式兼容。当启用no_sync模式时，它会在正向过程中正确设置一个标志，以指示它是否期望在下一个反向过程中执行梯度规约。如果通信发生，DDP将在随后的前向传递之前广播缓冲区。

5.2 Core Gradient Reduction

主要的开发工作花费在gradient reduction上，因为这是DDP中与性能最相关的步骤。该实现存在于reducer.cpp中，它由四个主要组件组成，即：

构建参数到桶的映射。
安装autograd hook。
启动bucket AllReduce
检测全局未使用的参数。

我们接下来阐述这四个组成部分。

参数到桶映射（Parameter-to-Bucket Mapping）对DDP速度有相当大的影响。在每次向后传播中，将所有参数梯度中的张量复制到桶中，并在AllReduce之后将平均梯度复制回桶中。为了加速复制操作，存储桶始终与参数在同一设备上创建。如果模型跨越多个设备，DDP会考虑设备关联性，以确保同一存储桶中的所有参数都位于同一设备上。AllReduce的顺序也会对结果产生影响，因为它决定了多少通信可以与计算重叠。DDP按model.parameters()的相反顺序启动AllReduce。

Autograd Hook是DDP在后向传播中的切入点。在构建过程中，DDP遍历模型中的所有参数，在每个参数上找到梯度累加器，并为每个梯度累加器安装相同的post hook函数。梯度累加器将在相应的梯度准备就绪时，会触发post hooks，DDP将计算出整个桶何时全部就绪，这样可以启动AllReduce操作。然而，由于无法保证梯度准备的顺序，DDP不能选择性地选择安装挂钩的参数。在当前的实现中，每个bucket都保留一个挂起的梯度计数。每个post-hook函数都会递减计数，当计数为零时，DDP会将一个桶标记为就绪。在下一次向前传播中，DDP会为每个桶补齐待定的累积计数。

Bucket AllReduce是DDP中通信开销的主要来源。一方面，在同一个桶中装入更多的梯度将减少通信开销的摊销系统。另一方面，由于每个桶需要等待更多的梯度，因此使用较大的桶尺寸将导致更长的归并等待时间。因此，桶大小是关键的权衡。默认情况下，每个存储桶的大小为25MB。应用程序应该根据经验测量其影响，并将其设置为其用例的最佳值。

全局未使用参数（Globally Unused Parameters）的梯度在向前和向后过程中应保持不变。检测未使用的参数需要全局信息，因为在一个DDP过程中，一个参数可能在一次操作中不存在，但可能在另一个过程的同一次迭代中参与训练。因此DDP在位图中维护本地未使用的参数信息，并启动额外的AllReduce以收集全局位图。由于位图比张量尺寸小得多，因此模型中的所有参数共享同一位图，而不是创建每桶位图（per-bucket bitmaps）。位图位于CPU上，以避免为每次更新启动专用CUDA内核。但是，某些ProcessGroup后端可能无法在CPU 张量上运行AllReduce。例如，ProcessGroupNCCL仅支持CUDA张量。此外，由于DDP应该与任何定制的ProcessGroup后端一起工作，它不能假设所有后端都支持CPU张量。为了解决这个问题，DDP在同一设备上维护另一个位图作为第一个模型参数，并调用非阻塞拷贝操作（non-blocking copy）将CPU位图移动到设备位图以进行集合通信。

0xFF 参考

http://www.vldb.org/pvldb/vol13/p3005-li.pdf

mapreduce 分布式 api pytorch

0 人点赞