Pytorch 分布式训练DDP(torch.distributed)详解-原理-代码

她的我

已于 2024-01-15 13:07:28 修改

阅读量8.7k

点赞数 28

文章标签： pytorch 分布式深度学习计算机视觉

于 2024-01-12 12:33:53 首次发布

本文链接：https://blog.csdn.net/weixin_42503655/article/details/135537053

版权

本文详细介绍了PyTorch中的分布式训练方法DP和DDP，包括它们的原理、流程和代码示例，重点阐述了数据并行和模型并行的区别，以及如何通过DistributedDataParallel进行DDP训练。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

0. 引言

目前已经有很多分布式训练的方法，一直在用，这里详细学习总结一下各种分布式训练的原理，本文主要介绍DP以及DDP，后面会介绍流水线并行、张量并行以及Zero优化等，用于学习记录，错误请指出。

Pytorch老版本的数据并行DP已经被弃用，这里还是分别介绍一下两者的过程，下面的图需要仔细看一下，方面更好的理解。
在这里插入图片描述

DP采用多线程实现，速度很慢，只适合单机多卡，目前已经被弃用。这里作为对比还是说明一下整个流程。DP需要各个GPU之间的通信量很大，主卡性能和通信开销容易成为瓶颈。
在这里插入图片描述

将模型权重从主卡（默认GPU:0）broadcast到所有GPU上，输入的bacth大小的数据会scatter到各个GPU上(每张卡上输入batch / 4，假设有四张卡)。
之后独自在各张卡上执行前向过程，每张卡的输出结果会被gather到主卡上，然后在主卡上计算总的损失，总损失是一个标量，一般可以求平均值，再通过损失函数求导，得到损失的梯度，分发到各个GPU上。
在每个GPU根据链式法则继续计算各个参数的梯度，再将所有设备上的梯度回传到主卡上进行梯度累加并求梯度均值。在主卡上通过平均梯度更新模型权重。最后将更新后的模型参数 broadcast 到其余 GPU 中进行下一轮的前向传播。

DDP可适用于单机多卡和多机多卡的数据并行训练，本质是采用了多进程控制。通常是一个进程控制一个GPU，每一个GPU上都有一份相同的模型权重，且每个进程在训练时都有独自的优化器，执行自己的更新过程，DDP进程间通信只传递梯度，通信量较少。下面还是说明一下整个流程。
在这里插入图片描述

首先将主卡的模型参数broadcast到其他GPU上，然后每个 DDP 进程都会创建一个 local Reducer 来负责梯度同步。每个进程从磁盘加载 batch 数据，并将它们传递到其 GPU，这一点在下面外面会展开讲解。
每个 GPU 都有自己的前向过程，分别得到相应的预测结果。在DDP中，损失函数的计算是在每个设备上独立进行的，每个设备计算的损失都是基于其处理的数据子集和相应的预测结果。
根据损失函数可以计算出参数梯度，之后将梯度信息在各个 GPUs 间进行 All-Reduce，每个 GPU 都收到其他 GPU 的梯度进行梯度累加后求平均值，然后可以独自进行反向传播和参数更新。
由于初始权重参数相同ÿ