初识分布式训练

news/2025/11/9 18:39:40/文章来源:https://www.cnblogs.com/rifo/p/19204672

假设有N块GPU,模型有ψ个参数。

前提知识：每个参数对应一个梯度值，且SGD每个参数对应一个一阶动量，Adam每个参数对应一个一阶、一个二阶动量

DP（data parallel）

数据并行（单进程，多线程，只用一个cpu核），每个GPU上都保存完整的模型参数（param,or p）与优化器状态(optimizer state,or os)。磁盘读取数据，通过一个cpu进程将数据分成多份，每个GPU一份，反向传播结束后，得到各自完整的梯度，其他GPU将各自计算的梯度同步到GPU0上，由GPU0的优化器平均梯度后并更新GPU0的参数，再将更新后的参数广播给其他GPU。

各GPU的通信量如下：（其中N-1为其余GPU数，ψ为梯度更新一次的通信量）

DP的问题在于：（1）单进程多线程，只能利用一个CPU核（2）GPU0的通信与计算压力大于其他GPU

DDP（distributed data parallel）

分布式数据并行（多进程），每个GPU上也都保存完整的模型参数（param,or p）与优化器状态(optimizer state,or os)。相比DP而言，用各自的优化器计算ring-allreduce后的梯度并更新参数，网络状态和优化器状态始终保持同步

Ring-AllReduce两个阶段：scatter-reduce和allgather

scatter-reduce示意图如下，每个GPU存有a、b、c，共3个梯度，经过3(模型参数量、梯度量即ψ)个通信回合，每个GPU上都有一个梯度值是完整的。（这个阶段梯度传递并相加）

allgather示意图如下，（梯度传递并直接覆盖）。

此外，为了提升效率，将梯度分组，先计算出来的那组梯度开始ring-allreduce同步，此时下一组梯度也正在被计算。

边计算边同步梯度

DDP通信量如下：（其中N-1为通信次数，ψ/N为单次通信量）

Deepspeed ZeRO

DP和DDP每个模型都需要存储完整的模型参数、梯度、优化器状态。使用Adam优化器的显存占用如图

而Deepspeed ZeRO-1、-2、-3分别对这三者做了处理，每张卡的显存占用如下图（baseline为标准DDP）

DDP的显存为fp16（2B）参数+fp16（2B）梯度+优化器（12B=4B一阶动量+4B二阶动量+4B参数）

ds-zero-1:相比DDP优化了优化器显存

ds-zero-2:相比DDP优化了优化器显存+梯度显存

ds-zero-3:相比DDP优化了优化器显存+梯度显存+参数显存

Deepspeed ZeRO-1

以3个GPU为例(Aadm+AWP)，每个GPU上的梯度和参数是完整的，但优化器不是，整个过程如下：

同步：每个GPU从后往前计算梯度，当最后1/3的梯度计算完成，GPU0和GPU1将其同步给GPU2，计算呢后1/3的梯度均值；与此同时，3个GPU计算计算中间1/3的梯度，并同步给GPU1；前1/3的梯度，同步给GPU0；到此各GPU都拿到了自己负责那部分的梯度均值（fp16）。
更新：接着将fp16转为fp32（不存储）->进行缩放->更新优化器的fp32一阶动量和二阶动量->更新fp32的参数（master-weights）

->更新各自fp16的参数。
广播：广播各自那部分的fp16参数给其他GPU。

通信量分析：（N-1表示GPU数量，ψ/N为其余每个GPU同步的梯度量）