PyTorch单机多卡从入门到入土（坑点记录）

在 Vision 里用上 Transformer 之后，单卡训练连两位数的 batchsize 都开不了，必须得学学单机多卡的使用了。

PyTorch 中，多卡训练有两种方案：

• DataParallel：只支持单机多卡，代码很方便，只需要添加一行，但是效率比较低，不推荐使用
• DistributedDataParallel：支持多机多卡，效率高，但是要折腾一下代码

基于性能考虑，一般我们会选择第二种 DDP. 其基本概念与流程请参阅网上的相关资料（见参考资料），此不赘述，这里只记录一些要点和坑点。

坑点记录

1. 保存模型 checkpoint、记录 tensorboard、输出准确率、甚至是 tqdm 都只让主进程（rank0）执行，其他进程不执行。

2. 在用 DDP 封装模型后，模型的本体（我们定义的那个类）是 model.module；所以保存模型时，最好保存 model.module.state_dict()，否则存下来的参数的 key 前面会多一个 module.，不便再次 load 模型。

3. 我们经常会在训练的每个 epoch 后进行 evaluate / inference，为避免有些进程测试完之后开始了下一轮训练，但其他进程还在测试，最好在每个 epoch 开始训练前（或者每个 epoch 完成训练后）用 dist.barrier() 同步一下。

4. 单卡 evaluate / inference 用的模型最好是本地模型 model.module 而非 DDP 包装的模型，否则非主进程会在第 3 条的 dist.barrier() 处卡死，推测原因是 DDP 包装的模型会在一些地方同步 bucket^[8][9]。

5. 多卡训练时报错：

   RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.FloatTensor [1]] is at version 3; expected version 2 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

   但单卡训练没有问题。这是因为对同一个模型连续 forward 了多次而没有 backward^[6]（比如在 GAN 的训练中，要分别将真实图像和生成图像通过判别器），DDP 会在每次 forward 中广播主进程模型的 buffers，但这个广播操作是 inplace 的^[7]，从而引起上述报错。

   解决办法：在 DDP 定义时指定参数 broadcast_buffers=False，取消广播。

6. 【尚未遇到过】对于依赖于模型初始 parameter 的一些 optimizer，要将其定义在 model=DDP(model) 之后，来保证各进程 optimizer 的初始状态是一致的。这是因为 optimizer 本身是和 DDP 无关的，DDP 不会同步 optimizer 的状态^[1]。当然，一般我们用的 optimizer 不会依赖于模型的 parameter，可以不管这一点。

7. 多卡训练时在 backward() 处卡死，没有报错，单卡训练正常：

   训练有逻辑分支，其中一条分支中部分模型参数不会被更新；如果没有设置 find_unused_parameters=True，那么进入这条分支后反向传播卡住。

   解决办法：在 DDP 定义时指定参数 find_unused_parameters=True.

何时会同步？

1. 在前向传播前，各进程同步模型的 parameter 和 buffer；
2. 反向传播时，各进程同步模型的梯度；
3. 调用 all_reduce() 时；
4. 调用 dist.barrier() 时；
5. 调用 torch.cuda.synchronize() 时【尚不清楚 torch.cuda.synchronize() 和 dist.barrier() 的区别】。

参考资料

1. [原创][深度][PyTorch] DDP系列第一篇：入门教程 - 996黄金一代的文章 - 知乎 https://zhuanlan.zhihu.com/p/178402798 ↩︎
2. [原创][深度][PyTorch] DDP系列第二篇：实现原理与源代码解析 - 996黄金一代的文章 - 知乎 https://zhuanlan.zhihu.com/p/187610959 ↩︎
3. [原创][深度][PyTorch] DDP系列第三篇：实战与技巧 - 996黄金一代的文章 - 知乎 https://zhuanlan.zhihu.com/p/250471767 ↩︎
4. 【分布式训练】单机多卡的正确打开方式（三）：PyTorch - Nicolas的文章 - 知乎 https://zhuanlan.zhihu.com/p/74792767 ↩︎
5. Pytorch 分布式训练 - 会飞的闲鱼的文章 - 知乎 https://zhuanlan.zhihu.com/p/76638962 ↩︎
6. 【PyTorch踩坑】一个排查了一下午的坑 - Tramac的文章 - 知乎 https://zhuanlan.zhihu.com/p/409117481 ↩︎
7. https://github.com/pytorch/pytorch/issues/22095 ↩︎
8. https://github.com/pytorch/pytorch/issues/54059 ↩︎
9. https://discuss.pytorch.org/t/torch-distributed-barrier-hangs-in-ddp/114522 ↩︎