分布式训练基础

深度学习

对于算法工程师而言，单卡跑不起来或者跑起来慢是很常见的问题，分布式训练各个框架都有支持。有时候训练模型一张卡放得下但误以为多卡一定就快，训练的时候使用多卡效率还不如单卡。也没有深入理解背后的原理，写这篇笔记以理解分布式训练中的 parameter server 算法与 Ring AllReduce 算法, 理解各自的工作原理即可

分布式训练，主要可以分为以下三种：

  • 并行方式: 数据并行 VS 模型并行
  • 参数更新方式: 同步更新 VS 异步更新
  • 对于参数的同步更新方式，主流分为: parameter server VS AllReduce

并行方式

  • 模型并行: 不同的GPU输入相同的数据，运行模型的不同部分，比如多层网络的不同层。比如网络太大，单卡已经装不小模型了，这种情况通讯开销大，各卡之间相当于前后依赖，毕竟模型各层之间是串行。
  • 数据并行: 不同的GPU 运行相同的完整的模型，输入不同的数据。一般分布式都是这种方式。

更新方式

当采用数据并行时，由于每块GPU上运行一部分数据，需要考虑将每个GPU上的梯度如何进行汇总来更新模型参数。

  • 同步更新:每个输入batch的数据会先按GPU划分到各个卡上，等所有GPU计算完成后，再统一计算新权值，然后所有GPU同步新值后，再进行下一轮计算。存在木桶效应，即同步更新会有等待，速度取决于最慢的那块GPU。但是实际一般采用同步更新，实现简单。
  • 异步更新:每个GPU计算完梯度后，无需等待其他更新，立即更新整体权值并同步。异步更新没有等待，但是其实现考虑的因素太多，这儿也不深究了。

parameter server VS AllReduce

parameter server

假如有5块GPU， 0号GPU为master，也是reducer, GPU 0将数据分成五份分到各个卡上，每张卡负责自己的那一份mini-batch的训练，得到grad后，返回给GPU 0上做累积，得到更新的权重参数后，再分发给各个卡. 导致木桶效应严重，所有的卡都要和master通信进行数据、梯度、参数的通信，时间开销大.

有N个GPU，通信一次完整的参数所需时间为K，那么使用PS架构，花费的通信成本为: $T = 2(N-1) K$,通信时间与GPU数量线性相关

在pytorch中 data parallel 即采用 ps 算法，整个流程如下图:

ring AllReduce

该算法是百度在 《Bringing HPC techniques to deep learning 》提出:

在 ring Allreduce 中，所有的GPU形成一个环。每个GPU有一个左邻居和一个右邻居；它只会向它的右邻居发送数据，并从它的左邻居接收数据，该算法分两步进行：第一步，scatter-reduce, 第二步：allgather.

第一步: 参数分为N份，相邻的GPU传递不同的参数，在传递N-1次之后，可以得到每一份参数的累积（在不同的GPU上）

在第一次迭代时，

GPU Send Receive
0 Chunk 0 Chunk 4
1 Chunk 1 Chunk 0
2 Chunk 2 Chunk 1
3 Chunk 3 Chunk 2
4 Chunk 4 Chunk 3

经过4轮迭代

第二步：The Allgather

Allgather 也需要 N-1 次迭代，使每块GPU得到所有参数

在第一次迭代时:

GPU Send Receive
0 Chunk 1 Chunk 0
1 Chunk 2 Chunk 1
2 Chunk 3 Chunk 2
3 Chunk 4 Chunk 3
4 Chunk 0 Chunk 4

在pytorch中，DDP 即采用allreduce 算法，整体流程如下:

ref

https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/