pytorch 分布式训练

分布式训练介绍

数据并行当任务数据的量非常大，但是模型本身规模不大（显存可以放得下）的时候，可以把数据按照一定的规则分配到不同的机器/卡上分工合作。同时，每个机器/卡都加载同一个模型，各自单独地根据输入数据计算梯度，并对梯度按照一定规则进行通信，这即梯度的汇总更新。接下来就是对模型参数的更新，确保每个iter结束的时候，所有卡上的参数一致。当然，这也要求初始化的时候每张卡的参数一致。数据并行的目的，就是希望在同样的时间内：更多的数据同时训练，加快了同等量级数据的训练完成时间， 相当于增大 batch_size，梯度下降方向会更贴合数据集整体的实际情况。

数据并行流程图如下：

All Reduce 架构：
每个GPU计算完结果后，需要汇总更新，就是图中红色框的All-Reduce就是达成数据并行的重要操作。

Allreduce操作是常用通信模式的一种。在说Allreduce之前，先看下reduce。reduce称为规约运算，是一系列运算操作的统称，细分来说包括SUM、MIN、MAX、PROD、LOR等。reduce意为减少/精简，因为其操作在每个进程上获取一个输入元素数组，通过执行操作后，将得到精简的更少的元素。例如下面的Reduce sum：

可理解为多个GPU的结果汇总到一块，再通过某种方式去得到一个最终的梯度回传的结果（多GPU共享一个结果）

从图中可以看出，all reduce操作可通过单节点上reduce+broadcast操作完成。

Allreduce方式也分多种，目前Parrots和Pytorch框架里用的都是Ring AllReduce 这种方式，而在进程间传输的就是模型参数的梯度。

在训练过程的表现可以如下图所示。

这个闭环的Ring可以 这张图表示。多卡训练的每张卡对应一个worker，所有worker形成一个闭环，接受上家的梯度，再把累加好的梯度传给下家。最终计算完成后更新整个环上worker的梯度（这样所有worker上的梯度就相等了），然后每张卡各自求梯度反向传播。

第二种架构： Server-Work架构
这里的server和worker就对应了我们训练用到的多卡，通常DP接口中输入参数的第一张卡被用作「server」，其他卡号就是对应的workers。

Server 负责：

• 将数据切成等量大小
• 从 device0分发到其他卡上并把模型复制到各个卡上
• 前向计算损失和反向计算梯度
• 更新参数
• 在下一个iter开始的时候把参数更新到worker上

worker负责：

• 载入server分发过来的数据和模型
• 模型前向计算

PS架构在训练中的劣势显而易见：通信的成本随着卡数的增加而增加； server负担过重，容易显存爆炸。
还有的设计是只把server用作梯度更新，进一步减轻负担，比如下面这个过程：

在Pytorch支持的分布式方式中：
DP使用的是Parameter Server(PS)架构；
DDP官网推荐，性能优于DP,
All-reduce架构；apex是英伟达出品的工具库，封装了DistributedDataParallel，All-Reduce架构。

作者：九点澡堂子
链接：https://zhuanlan.zhihu.com/p/366253646
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。