Skip to content

Latest commit

 

History

History
131 lines (77 loc) · 7.86 KB

PyTorch Distributed-Data Parallel Training.md

File metadata and controls

131 lines (77 loc) · 7.86 KB

前言

本文介绍 PyTorch 1.5 里的数据并行。

训练主要包含三部分:

  1. forward pass: 计算loss
  2. backward pass:计算梯度
  3. optimizer step : 更新参数

可以创建模型的多个副本,每个工作在训练数据的一部分上,单独执行forward和backward。之后多个副本间可以同步梯度或者更新参数。

提供一个通用的分布式数据并行的package,包含三部分:

  1. 数学上相等:虽然是分布式,但是跟本地训练一样
  2. 非侵入和截断的 API:一般都是从本地开发然后扩展到分布式。API 需要允许内部实现可以即时拦截信号来做通信和系统优化
  3. 高性能:计算和通信

为了保证数学上的相等,所有的副本的模型参数都从同样的初始值开始,在训练的每次迭代种进行梯度同步来保证参数是一致的。

使用到一些技术:

  1. 梯度分桶(gradients bucketing) # 那我们应该配置成什么样比较好?
  2. 通信和计算叠加
  3. 跳过同步(synchronization)

也覆盖了在 NCCL和Gloo通信库。结论显示:

  1. 通信是训练中延迟的大头,尤其是当模型增大的时候
  2. 桶的大小会显著影响通信效率,当配置得当可以加速2倍
  3. 正确跳过同步可以显著减少分摊之后的通信开销,而不显著降低收敛速度

本论文三块贡献:

  1. 揭示了大范围使用的工业领先的分布式训练的实现
  2. 真实世界里的一些优化点
  3. 展示了性能优化的一些经验

问题:

  1. reduce 梯度而非参数,而且是每一层各自去做?对,reduce 是挂在每个 parameter 上的 grad 函数里。可以在backward时逐步进行

2. PyTorch 和数据并行

2.1 PyTorch

在 forward pass里,PyTorch 构建一个自动求梯度的图,记录需要执行的操作。然后在backward 时,使用 autograd 图来指导反向转播过程产生 梯度。最终优化器利用梯度更新参数。训练过程就是不断迭代这三步骤直到模型收敛。

2.2 数据并行

PyTorch 提供几个工具来做分布式训练:

  1. DataParalle:单进程,多线程数据并行训练,使用同一个机器上的多块卡
  2. DistributedDataParallel:在多 GPU和机器上进行多进程的数据并行训练
  3. RPC:给通用的分布式模型并行训练(比如PS)

本论文主要集中在 DistributedDataParallel,数据并行是通过在优化器步骤前集群里同步梯度来确保所有模型副本的参数都是使用同一套梯度来更新的,因此模型副本都能在iteration期间保持一致。可以选择同步梯度或权重,这里选择了同步梯度。保持数据并行 和不并行的效果一样

还有一个参数平均(Parameter averaging),可以用来扩展模型训练。不是同步梯度,而是直接对所有模型副本的参数进行平均。这一步发生在 optimizer step 之后,所以参数平均(parameter averaging 可以完全作为额外步骤实现,不需要参与到本地训练步骤里,虽然可以轻松、完全和本地迭代解耦合,但是有缺点:

  • 数学上并不等价于在本地把所有数据都处理了,最终结果是不同的
  • 每个模型的分片看到的梯度不一样,所以最终可能优化器里的状态会分散,导致出现不同的梯度下降方向
  • 这种结构把计算(比如 backward)和通信(计算平均)编排为不重叠。所以无论如何优化,计算或通信有一个是空闲的。

问题:为啥?不是也可以重叠嘛? 不太好重叠,因为目前典型的都是 第三部单独一个 optimizer.Step(),一次计算出所有参数的更新。

由于上述问题,最终决定 ddp 使用数据并行来同步梯度而非参数。

2.3 AllReduce

AllReduce 是被 DistributedDataParallel 使用的通信原语 API,用来在所有进程间计算梯度的求和。很多苦都支持,比如 NCCL,Gloo,MPI。AllReduce 操作期待每个参与的进程提供一个大小相同的 Tensor,集体地提供一个算数操作(sum,prod,min,max)作用在所有进程的输入 Tensor上,返回同样的tensor结果给每个参与方。最简单的实现是每个进程都广播输入的tensor到集群中其他人,然后各自独立进行算数运算。然后 AllReduce 在分布式训练速度上有显著影响,通讯库实现了很多复杂高效的算法,比如基于环的 AllReduce, 基于树的。由于一个 AllReduce 操作只能等所有进程 ready后才能开始,所以它被当作一个同步的通信,跟 PS 中使用的p2p是不同的。

3. 设计

在分布式训练里,每个进程有自己本地的模型副本和本地的优化器。为了保证正确性,分布式数据并行训练和本地训练必须数学上相等。DDP 保证所有模型副本从同一个模型状态开始,在每个backward之后,同样的梯度。

3.1 API

有两个设计目标:

  • 非侵入式(non-intrusive)
  • 可拦截(Interceptive) :API 里可以拦截多种信号

使用构造函数,能够拦截,观察到模型结构和参数。构造之后,本地模型被分布式代替,可以拦截 forward() 函数来执行一些需要的动作。为了backward,DDP 依赖于 backward hook来出发梯度合并,会在执行 backward 时被 autograd 引擎调用。

3.2 梯度合并

DDP 里的梯度合并算法。

3.2.1 简单版本

在 augograd 的hook里,可以通过 AllReduce 来同步。但是这样有两个问题:

  1. 集合通信在小的tensor上性能很差,尤其在大模型,有大量小参数的情况。为啥?
  2. 梯度计算和同步分离开,就没有机会把计算和通信重叠到一起

3. 系统设计

为了保证正确性,ddp 训练和本地训练必须在数学上等价

问题:AllReduce 如何保证各台主机上修改的都是同一份tensor?万一期间又有更新?

3.2.2 梯度分桶(Gradient Bucketing)

60M torch.float32 的参数。在更大的tensor上,集合通信效率更高。与其每一层梯度tensor 有效后就立马发起 AllReduce 操作,不如等一小会儿然后把多个梯度桶合并到一个 AllReduce 操作里,这样吞吐更高,延迟更低。(为啥延迟更低?)这尤其对有很多小参数的模型有效。

问题:怎么衡量模型是“小参数”?比如特别深,而又窄的模型?

又提到相对小的bucket 大小,DDP 可以把 AllReduce 和 backward 结合起来,用计算掩盖通信开销,可以在每个 iteration 延迟显著降低。

3.2.3 计算和通信的重叠

当本地 backward 结束后,梯度上的 AllReduce 就可以开始了。使用了 bucketing 后,DDP 只需要等待同一个桶内的所有内容到位后,开始通信。DDP 给每个梯度累加器注册了一个 autograd hook。当对应的累加器更新了梯度后,这个狗子就会执行,扽带同一个 bucket 里的所有梯度都到位,最后一个狗子的执行会触发本桶上一个异步的 AllReduce。

有两个警告:

  1. 在所有进程上的 reducing order 必须一致,否则乱序。解决方法是按照 model.parameters() 作为顺序来划分到桶里。而且按照桶的次序(下标)来reduce,而不是按照 bucket ready的时间
  2. 由于pytorch的动态图,有一些梯度可能在一些迭代过程中跳过。所以在 forward 时,会记录有哪些参数被用到了,这些是backward时需要对梯度做 reduce 的

问题:在 pipeline 并行时,难道不能用吗?

3.2.4 梯度累加

当一个batch太大,可以分成好多 microbatch,等大家都跑完后,再做累加。

3.3 集合通信

NCCL 的实现里,ProcessGroup 维护了特定的 cuda stream 来做通信,所以不会阻塞默认的流。

4. 实现和评估

6. 经验

7 相关工作

[30] An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pages 8024–8035. Curran Associates, Inc., 2019.