理解 PyTorch DDP 分布式数据并行模式

PyTorch 使用 DDP(Distributed Data Parallel) 实现了真正的分布式数据并行,在下面的两个场景下都可以使用 DDP 实现模型的分布式训练: 单机、多 GPU(单进程多线程的伪分布式) 多机、多 GPU(多机多进程的真正分布式) 上面第一种方式,就是类似使用简单的 DP 数据并行模式,但是 DP 使用的单进程、多线程的范式来实现的;而 DDP 完全使用了多进程的方式,包括单机多进程、多机多进程,如果是多机的情形则对应着物理上的分布式多进程模式。为了获得更好的性能,最好是使用 DDP 模式来训练模型,即使是在单机、多 GPU 的情况下,也建议使用 DDP 模式来实现基于数据并行的模型训练,使用单机 DDP 模式训练模型的性能要比 DP 模式好很多。 DDP 基于集合通信(Collective Communications)来实现分布式训练过程中的梯度同步。在反向传播过程中,DDP 使用 AllReduce 来实现分布式梯度计算和同步。 下面,我们从与 DDP 相关的几个方面来理解 DDP 的设计与实现,包括: 集合通信(Collective Communication) 通信后端(Communication Backend) DDP 内部实现概览