本章节参考UP主 ↗

分布式训练是为了解决大规模深度学习模型在单设备上无法容纳或训练过慢的问题，通过将计算和存储负载分散到多个 GPU 或机器上协同完成训练。

其中，DP (DataParallel) 是 PyTorch 早期的单进程多线程方案，将一个 batch 的数据分发到多个 GPU 并行前向/反向，但梯度汇总和参数更新由主 GPU（如 cuda:0）完成，存在通信瓶颈且不支持多机，现已基本弃用。 DDP (DistributedDataParallel) 是当前主流的多进程分布式训练方法，每个 GPU 独立运行一个进程，拥有完整的模型副本，通过高效的 All-Reduce 操作同步梯度，支持单机多卡和多机多卡，扩展性好、性能高。 而 DeepSpeed ZeRO（Zero Redundancy Optimizer）则从内存优化角度出发，在数据并行基础上对优化器状态、梯度和模型参数进行分片存储，显著降低显存占用，使百亿甚至千亿参数模型能在普通 GPU 集群上训练。

三者中，DDP 侧重通信效率，ZeRO 侧重内存节省，实践中常结合使用（如 DDP + ZeRO）以兼顾速度与规模。

Data Parallel (DP)#

DP是从硬盘读取数据，通过一个CPU进程将数据集划分成多份，然后分发给每个GPU独立进行前向和反向传播，计算出各自的梯度，然后各GPU将各自的梯度传到GPU0上。GPU0通过对这些梯度求和取平均更新自己模型的参数，然后广播给其他GPU。 如图： 在分布式训练里，可能一半的时间都花在多卡之间的通信上，下面分析一下DP的通信量：

假设模型参数量为$\psi$，GPU节点数为$N$。 则对于GPU0，传入梯度为：$(N-1)\psi$，传出参数为$(N-1)\psi$ 对于其他GPU：传出梯度为$\psi$，传入参数为$\psi$。

存在的问题#

Distributed Data Parallel（DDP）#

介绍DDP之前，先了解一下Ring-AllReduce：Ring-AllReduce 是一种高效的分布式梯度同步算法，广泛应用于大规模深度学习训练中（如 PyTorch DDP、Horovod 等）。其核心思想是将参与训练的多个 GPU（或节点）组织成一个逻辑环形拓扑（ring），每个设备只与左右两个邻居通信，通过多轮分块传输逐步完成全局梯度的聚合。

整个过程分为两个阶段：Scatter-Reduce（各设备将本地梯度分块，沿环传递并累加，最终每块由一个设备持有完整和）和 All-Gather（将累加后的梯度块沿环广播，使所有设备获得完整的全局梯度）。

相比传统的中心化聚合（如 Parameter Server 或 DP 的主 GPU 汇总），Ring-AllReduce 消除了单点瓶颈，通信量与设备数无关（带宽最优），且能充分利用多设备间的并行通信能力。尤其在高速互联网络（如 NVLink、InfiniBand）下，它能实现接近线性的扩展效率，是现代大模型分布式训练的通信基石。

首先将GPU0的参数$a_0$传给GPU1，GPU1的参数$b_1$传给GPU2，GPU2的参数$c_2$传给GPU0，形成一个闭环，然后GPU与各自传入的参数相加，得到下面的结果。

然后将GPU0的参数$c_0+c_2$传给GPU1，GPU1的参数$a_0+a_1$传给GPU2，GPU2的参数$b_1+b_2$传给GPU0，得到下面的结果。

以此类推，最终得到：

从而实现了所有GPU上都有聚合后的模型参数，在这个过程中，每个GPU同时发送和接收，可以最大限度利用每个显卡的上下行带宽。

与DP不同，DDP是多进程的，每个进程为自己的GPU准备数据并和其他GPU通信，每个GPU用自己的模型进行前向和反向传播，因为每个GPU的数据都不同，所以梯度也不同，最后通过Ring-AllReduce实现梯度的同步与聚合，每个GPU通过更新梯度来更新模型： 可以更深入理解DDP，在反向传播过程中，越后面的层数，梯度最先得到，但是每计算一个梯度就进行Ring-AllReduce会大幅增加开销，所以会使用桶，计算的梯度会保存在桶中。具体来说，会对每个层注册监听器（为了让DDP框架知道哪些梯度计算好了），当所有GPU的同一个桶都装满了，会执行一次Ring-AllReduce同步，在这个过程中GPU还在计算其他的梯度，当所有的桶都梯度同步了，每个GPU会调用自己的优化器进行更新： 下面分析一下DDP通信量： 假设模型参数量为$\psi$，GPU节点数为$N$。 对于每一个GPU进程： Scatter-Reduce阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ AllGather阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ 总传入/传出：$2\psi$ 与集群大小无关

总结#

DeepSpeed ZeRO#

DeepSpeed ZeRO（Zero Redundancy Optimizer）是微软提出的一种显存优化技术，旨在消除传统数据并行训练中的内存冗余。它通过将优化器状态、梯度和模型参数在多个 GPU 之间进行分片存储，显著降低每个设备的显存占用。ZeRO 分为三个阶段：ZeRO-1 分片优化器状态，ZeRO-2 增加梯度分片，ZeRO-3 进一步分片模型参数，使得百亿甚至千亿参数的大模型能在普通 GPU 集群上高效训练。该技术可与分布式训练框架（如 PyTorch DDP）结合使用，在几乎不损失训练速度的前提下，大幅提升模型规模的可扩展性。

正常情况下每个GPU上要有Data、模型参数和梯度以及优化器的参数和梯度，下面使用的优化器是Adam，所以要存储一阶动量和二阶动量：

ZeRO-1#

由上图可知，占用显存最多的就是优化器状态，因为它是FP32精度，并且每个GPU都需要存储，所以ZeRO-1对优化器状态进行了分片。 这里假设有一个9层的网络结构，每个分片优化器负责管理3层网络结构。当反向传播完成了倒数前3层的梯度计算时，此时GPU0和GPU1并没有计算这些层的优化器，所以需要传给GPU2的优化器进行梯度的求和取平均，当所有的GPU都拿到对应层的梯度时，将梯度值转换为FP32梯度（防止下溢），然后更新对应层的参数。最后，每个GPU将自己更新的那部分层的参数广播给其他GPU，这样就完成了一个epoch的训练。

通讯量分析#

假设模型参数量为$\psi$，GPU节点数为$N$。 对于每一个GPU进程： 梯度收集阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ 参数广播阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ 总传入/传出：$2\psi$ 与DDP通讯量相同。

ZeRO-2#

ZeRO-2的思路更ZeRO-1一样，只是将梯度进行了分片。 在ZeRO-1的基础上，每个GPU只维护了自己负责层的梯度值，其他操作都是一致的，进一步节省了显存。

通讯量分析#

假设模型参数量为$\psi$，GPU节点数为$N$。 对于每一个GPU进程： 梯度收集阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ 参数广播阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ 总传入/传出：$2\psi$ 与DDP通讯量相同。

ZeRO-3#

ZeRO-3 进一步分片模型参数。在前向传播过程中，由于每个GPU维护不同层的参数，前向传播开始时，GPU1和GPU2没有前3层的参数，所以需要通过GPU0广播分发参数给他们进行计算，GPU1和GPU2计算完之后就把这些参数丢弃，从而节省了显存。 前向传播结束后，进行反向传播。同样的，反向传播开始时，GPU0和GPU1没有最后3层的参数，所以需要通过GPU2广播分发参数给他们进行计算，计算完之后就丢弃，节省缓存。

通讯量分析#

假设模型参数量为$\psi$，GPU节点数为$N$。 对于每一个GPU进程： 梯度收集阶段传入/传出：$(N - 1) \frac{\psi}{N} \approx \psi$ 参数广播阶段传入/传出：$2*(N - 1) \frac{\psi}{N} \approx 2\psi$ 总传入/传出：$3\psi$ 是DDP传输量的1.5倍。

对比#