PyTorch多机多卡训练：DDP实战与技巧
点上⽅蓝字计算机视觉联盟获取更多⼲货
在右上⽅ ··· 设为星标 ★，与你不见不散
仅作学术分享，不代表本公众号⽴场，侵权联系删除
转载于：作者⼁996黄⾦⼀代@知乎（已授权）
AI博⼠笔记系列推荐
零. 概览
想要让你的PyTorch神经⽹络在多卡环境上跑得⼜快⼜好？那你definitely需要这⼀篇！
No one knows DDP better than I do!
– – magic_frog（⼿动狗头）
本⽂是DDP系列三篇中的第三篇。本系列⼒求深⼊浅出，简单易懂，猴⼦都能看得懂（误）。
在过去的两篇⽂章⾥，我们已经对DDP的理论、代码进⾏了充分、详细的介绍，相信⼤家都已经了然在胸。但是，实践也是很重要的。正所谓理论联系实践，如果只掌握理论⽽不进⾏实践，⽆疑是纸上谈兵。
在这篇⽂章⾥，我们通过⼏个实战例⼦，来给⼤家介绍⼀下DDP在实际⽣产中的应⽤。希望能对⼤家有所帮助！
1. 在DDP中引⼊SyncBN
2. DDP下的Gradient Accumulation的进⼀步加速
3. 多机多卡环境下的inference加速
4. 保证DDP性能：确保数据的⼀致性
5. 和DDP有关的⼩技巧如何学习催眠术
6. 控制不同进程的执⾏顺序
7. 避免DDP带来的冗余输出
请欢快地开始阅读吧！
依赖：pytorch(gpu)>=1.5，python>=3.6
⼀. 在DDP中引⼊SyncBN
什么是SyncBN？
SyncBN就是Batch Normalization(BN)。其跟⼀般所说的普通BN的不同在于⼯程实现⽅式：SyncBN能够完美⽀持多卡训练，⽽普通BN 在多卡模式下实际上就是单卡模式。
我们知道，BN中有moving mean和moving variance这两个buffer，这两个buffer的更新依赖于当前训练轮次的batch数据的计算结果。但是在普通多卡DP模式下，各个模型只能拿到⾃⼰的那部分计算结果，所以在DP模式下的普通BN被设计为只利⽤主卡上的计算结果来计算moving mean和moving variance，之后再⼴播给其他卡。这样，实际上BN的batch size就只是主卡上的batch size那么⼤。当模型很⼤、batch size很⼩时，这样的BN⽆疑会限制模型的性能。
为了解决这个问题，PyTorch新引⼊了⼀个叫SyncBN的结构，利⽤DDP的分布式计算接⼝来实现真正的多卡BN。
SyncBN的原理
SyncBN的原理很简单：SyncBN利⽤分布式通讯接⼝在各卡间进⾏通讯，从⽽能利⽤所有数据进⾏BN计算。为了尽可能地减少跨卡传输量，SyncBN做了⼀个关键的优化，即只传输各⾃进程的各⾃的 ⼩batch mean和 ⼩batch variance，⽽不是所有数据。具体流程请见下⾯：团干培训心得
1. 前向传播
2. 在各进程上计算各⾃的 ⼩batch mean和⼩batch variance
3. 各⾃的进程对各⾃的 ⼩batch mean和⼩batch variance进⾏all_gather操作，每个进程都得到s的全局量。
4. 注释：只传递mean和variance，⽽不是整体数据，可以⼤⼤减少通讯量，提⾼速度。
5. 每个进程分别计算总体mean和总体variance，得到⼀样的结果
6. 注释：在数学上是可⾏的，有兴趣的同学可以⾃⼰推导⼀下。
7. 接下来，延续正常的BN计算。
8. 注释：因为从前向传播的计算数据中得到的batch mean和batch variance在各卡间保持⼀致，所
以，running_mean和running_variance就能保持⼀致，不需要显式地同步了！
9. 后向传播：和正常的⼀样
SyncBN与DDP的关系
⼀句话总结，当前PyTorch SyncBN只在DDP单进程单卡模式中⽀持。SyncBN⽤到 all_gather这个分布式计算接⼝，⽽使⽤这个接⼝需要先初始化DDP环境。
复习⼀下DDP的伪代码中的准备阶段中的DDP初始化阶段
d. 创建管理器reducer，给每个parameter注册梯度平均的hook。
i. 注释：这⼀步的具体实现是在C++代码⾥⾯的，即reducer.h⽂件。
e. （可能）为可能的SyncBN层做准备
这⾥有三个点需要注意：
分配工资会计分录这⾥的为可能的SyncBN层做准备，实际上就是检测当前是否是DDP单进程单卡模式，如果不是，会直接停⽌。
这告诉我们，SyncBN需要在DDP环境初始化后初始化，但是要在DDP模型前就准备好。
为什么当前PyTorch SyncBN只⽀持DDP单进程单卡模式？
从SyncBN原理中我们可以看到，其强依赖了all_gather计算，⽽这个分布式接⼝当前是不⽀持单进程
红头船公园
多卡或者DP模式的。当然，不排除未来也是有可能⽀持的。
怎么⽤SyncBN？
怎么样才能在我们的代码引⼊SyncBN呢？很简单：
# DDP init
dist.init_process_group(backend='nccl')
# 按照原来的⽅式定义模型，这⾥的BN都使⽤普通BN就⾏了。
model = MyModel()
# 引⼊SyncBN，这句代码，会将普通BN替换成SyncBN。
model = vert_sync_batchnorm(model).to(device)
# 构造DDP模型
荤段子大全
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
⼜是熟悉的模样，像DDP⼀样，⼀句代码就解决了问题。这是怎么做到的呢？
convert_sync_batchnorm的原理：
了dules.batchnorm._BatchNorm类，就把它替换成SyncBN。也就是说，如果你的Normalization层是⾃⼰定义的特殊类，没有继承过_BatchNorm类，那么convert_sync_batchnorm是不⽀持的，需要你⾃⼰实现⼀个新的SyncBN！
@classmethod
def convert_sync_batchnorm(cls, module, process_group=None):
r"""Helper function to convert all :attr:`BatchNorm*D` layers in the model to
:class:`SyncBatchNorm` layers.
"""
module_output = module
if isinstance(module, dules.batchnorm._BatchNorm):
module_output = SyncBatchNorm(module.num_features,
module.eps, um,
module.affine,
process_group)
if module.affine:
_grad():
module_output.weight = module.weight
module_output.bias = module.bias
module_output.running_mean = module.running_mean
module_output.running_var = module.running_var
module_output.num_batches_tracked = module.num_batches_tracked
for name, child in module.named_children():
module_output.add_module(name, vert_sync_batchnorm(child, process_group))
del module
return module_output
⼆. DDP下的Gradient Accumulation的进⼀步加速
什么是Gradient Accmulation？
为什么还能进⼀步加速？
我们仔细思考⼀下DDP下的gradient accumulation。
# 单卡模式，即普通情况下的梯度累加
for 每次梯度累加循环
<_grad()
for 每个⼩step
prediction = model(data)
loss_fn(prediction, label).backward() # 积累梯度，不应⽤梯度改变
optimizer.step() # 应⽤梯度改变
我们知道，DDP的gradient all_reduce阶段发⽣在loss_fn(prediction, label).backward()。这意味着，在梯度累加的情况下，假设⼀次梯度累加循环有K个step，每次梯度累加循环会进⾏K次 all_reduce！但事实上，每次梯度累加循环只会有⼀次 optimizer.step()，即只应⽤⼀次参数修改，这意味着在每⼀次梯度累加循环中，我们其实只要进⾏⼀次gradient all_reduce即可满⾜要求，有K-1次 all_reduce被浪费了！⽽每次 all_reduce的时间成本是很⾼的！
如何加速
所以，我们可以这样实现加速：
model = DDP(model)
for 每次梯度累加循环
<_grad()
# 前K-1个step，不进⾏梯度同步，累积梯度。
for K-1个⼩step:
坚韧不拔意思_sync():
prediction = model(data)序号函数
loss_fn(prediction, label).backward()
# 第K个step，进⾏梯度同步
prediction = model(data)
loss_fn(prediction, label).backward()
optimizer.step()
给⼀个优雅写法（同时兼容单卡、DDP模式哦）：
from contextlib import nullcontext
# 如果你的python版本⼩于3.7，请注释掉上⾯⼀⾏，使⽤下⾯这个：
# from contextlib import suppress as nullcontext
if local_rank != -1:
model = DDP(model)
<_grad()
for i, (data, label) in enumerate(dataloader):
# 只在DDP模式下，轮数不是K整数倍的时候使⽤no_sync
my_context = _sync if local_rank != -1 and i % K != 0 el nullcontext
with my_context():
prediction = model(data)
loss_fn(prediction, label).backward()
if i % K == 0:
optimizer.step()
<_grad()
是不是很漂亮！
三. 多机多卡环境下的inference加速
问题
有⼀些⾮常现实的需求，相信⼤家肯定碰到过：
1. ⼀般，训练中每⼏个epoch我们会跑⼀下inference、测试⼀下模型性能。在DDP多卡训练环境下，能不能利⽤多卡来加速inference
速度呢？
2. 我有⼀堆数据要跑⼀些⽹络推理，拿到inference结果。DP下多卡加速⽐太低，能不能利⽤DDP多卡来加速呢？
解法
这两个问题实际是同⼀个问题。答案肯定是可以的，但是，没有现成、省⼒的⽅法。
测试和训练的不同在于：
1. 测试的时候不需要进⾏梯度反向传播，inference过程中各进程之间不需要通讯。
2. 测试的时候，不同模型的inference结果、性能指标的类型多种多样，没有统⼀的形式。
3. 我们很难定义⼀个统⼀的框架，像训练时model=DDP(model)那样⽅便地应⽤DDP多卡加速。
解决问题的思路很简单，就是各个进程中各⾃进⾏单卡的inference，然后把结果收集到⼀起。单卡inference很简单，我们甚⾄可以直接⽤DDP包装前的模型。问题其实只有两个：
我们要如何把数据split到各个进程中
我们要如何把结果合并到⼀起
如何把数据split到各个进程中：新的data sampler
想⼤家肯定还记得，在训练的时候，我们⽤的 torch.utils.data.distributed.DistributedSampler帮助我们把数据不重复地分到各个进程上去。但是，其分的⽅法是：每段连续的N个数据，拆成⼀个⼀个，分给N个进程，所以每个进程拿到的数据不是连续的。这样，不利于我们在inference结束的时候将结果合并到⼀起。
所以，这⾥我们需要实现⼀个新的data sampler。它的功能，是能够连续地划分数据块，不重复地分到各个进程上去。直接给代码：
