模型训练之dp、ddp与ZeRO

参考：https://zhuanlan.zhihu.com/p/513449071
https://zhuanlan.zhihu.com/p/617133971
一个经典数据并行的过程如下：

若干块计算GPU，如图中GPU0~GPU2；1块梯度收集GPU，如图中AllReduce操作所在GPU。
在每块计算GPU上都拷贝一份完整的模型参数。
把一份数据X（例如一个batch）均匀分给不同的计算GPU。
每块计算GPU做一轮FWD和BWD后，算得一份梯度G。
每块计算GPU将自己的梯度push给梯度收集GPU，做聚合操作。这里的聚合操作一般指梯度累加。当然也支持用户自定义。
梯度收集GPU聚合完毕后，计算GPU从它那pull下完整的梯度结果，用于更新模型参数W。更新完毕后，计算GPU上的模型参数依然保持一致。
聚合再下发梯度的操作，称为AllReduce。
前文我们说过，聚合梯度 + 下发梯度这一轮操作，称为AllReduce。接下来我们介绍目前最通用的AllReduce方法：Ring-AllReduce。它由百度最先提出，非常有效地解决了数据并行中通讯负载不均的问题，使得DDP得以实现。

dp：单进程多线程进行数据并行化。GPU:0忙前忙后去分发数据、汇总梯度，更新参数的时候，其他GPU就静静躺着。这种效率不高的数据并行方法，注定要被淘汰。是的，我们迎来了DDP(DistributedDataParallel)。
ddp：多进程。每个进程都独立优化所有的步骤，在每个进程中他们在自己GPU上计算loss，反向计算梯度。
运行程序的时候命令：

CUDA_VISIBLE_DEVICES="0,1" python -m torch.distributed.launch --nproc_per_node 2 main.py

DDP消除了DP的低效率，没有了主GPU，每个GPU执行相同的任务，他们自力更生，不再需要主GPU分发数据，他们每一个都需要从磁盘中加载数据，体验着模型训练的每一步。当然，他们也需要通信，首先，DistributedSampler负责分发数据，所有GPU计算得到梯度之后会通过通信后端进行通信，ring all-reduce的方式每个GPU都会得到其他所有GPU的参数梯度，进而reduce parameters gradients，去更新参数。
针对DP来说，Dataloder的batch_size是针对所有卡训练batch_size的和，例如10卡，每张卡batch_size是20，那么就要指定batch_size为200。针对DDP来说，batch_size就是一张卡所训练使用的batch_size的值。

ddp比dp训练速度快很多。

ZeRO：零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式，旨在降低存储。