DeepSpeed之ZeRO系列

date

May 15, 2025

slug

deepspeed_zero

TODO

🥰总结

切分效果

切分策略	切分部分	显存占用	额外的通信
baseline	—-		无
zero1	优化器状态		无
zero2	梯度		无
zero3	权重		1.5x

zero-offload将优化器状态放在CPU, GPU计算出梯度, 传给CPU更新参数, 然后将更新后的参数传回GPU更新GPU的参数用于后续计算

一般使用zero1或zero2

⁉️问题

zero节省了哪部分显存?, 能节省多少?

zero主要是节省模型状态(包括权重, 梯度, 优化器状态)占用的显存, 这一部分一共, zero1切分优化器状态, 变成, zero2继续切分梯度, 变成, zero3继续切分权重, 变成, 还有zero-offload, 这是把显存卸载的内存, 会带来负优化效果, 一般没有实用价值

zero是否带来了额外的通信?

zero1和zero2不会带来额外的通信, 但是zero3会带来, 所以一般用zero1或zero2

🧐内容

目前主流的大模型训练技术路线为: GPU + PyTorch + Megatron + DeepSpeed,

其中DeepSpeed的核心是ZeRO(Zero Redundancy Optimizer), 简单来说，它是一种显存优化的数据并行(data parallelism, DP)方案。

ZeRO: 一种去除冗余的数据并行方法

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models ，DeepSpeed项目最初就是论文中ZeRO方法的官方实现。

显存分类

ZeRO将模型训练阶段，每张卡中显存内容分为两类：

模型状态（model states）: 模型参数（fp16）、模型梯度（fp16）和Adam状态（fp32的模型参数备份，fp32的momentum和fp32的variance）。假设模型参数量为，则存储字节数为:

可以看到，Adam状态占比75%。

剩余状态（residual states）: 除了模型状态之外的显存占用，包括激活值（activation）、各种临时缓冲区（buffer）以及无法使用的显存碎片（fragmentation）。

相比之下，激活值可以用来大大减少，所以模型状态就成了头号显存杀手，它也是ZeRO的重点优化对象。而其中Adam状态又是第一个要被优化的。

zero-stage1

针对模型状态的存储优化（去除冗余），ZeRO使用的方法是分片（partition），假设一共N张卡, 每张卡只存的模型状态，这样系统内只维护一份模型状态。

首先进行分片操作的是模型状态中的optimizer states，也就是上图中的。模型参数（parameters）和梯度（gradients）仍旧是每张卡保持一份，此时，每张卡的模型状态所需显存字节为:

当比较大时，趋向于字节，也就是原来的1/4

通行量分析

在单独使用的情况下，单个显卡会保存完整的模型参数和梯度。随后使用reduce-scatter将梯度reduce至不同的显卡上(此时不同显卡仅拥有完整平均梯度的一部分)，该步骤的通信量是。各个显卡使用”部分最终梯度”更新对应的优化器状态，然后再更新对应的参数(此时每个显卡上的模型都更新了一部分参数)。最后，使用all-gather将分布在各个显卡上的更新后参数分发自所有显卡上(此时所有显卡上都有了完整的更新后参数)，该步骤的通信量是。总的来说，各个显卡仅需要持有部分优化器状态即可，且总的通信量仍然是2。

而且我们会发现每张显卡实际上拥有完整的Gradient，但是由于 Optimizer States 的分区，每个显卡只需要一部分 Gradient。这里实际上每张卡都冗余存储了 Gradient。这也是为什么的通信量等于

zero-stage2

如果继续对模型梯度进行分片，也就是上图中的，模型参数仍旧是每张卡保持一份，此时，每张卡的模型状态所需显存字节为:

当比较大时，趋向于字节，也即是原来1/8

通信量分析

由于对 Gradient 进行了分区，在更新参数前，先通过 ReduceScatter 把每个 Optimizer States 所需要的 Gradient 发送到对应的设备上。所以，这个操作的通信量仍然是（请反复理解下）。

每个分区好的 Optimizer Sates 获得对应的 Gradient 更新其参数后，只需要执行一次 AllGather，把自己更新的模型参数发出去并收集别人更新的模型参数，基于附录中的分析这同样是的通信量。因此每个训练步骤的总通信量为 + =2 ，与标准的 DP 完全相同。

zero-stage3

如果继续对模型参数进行分片，也就是下图中的，此时每张卡的模型状态所需显存字节为:

当 比较大时，趋向于0

通行量分析

由于对参数进行了分区，那么一个显而易见的事实是在前向和反向传播阶段都需要一次 AllGather 来保证计算的正确。值得说明的是，模型参数的分发是按神经网络的计算顺序去流水线分发的，因为如果不考虑计算的过程直接分发会导致一些参数被多次的分发并直接丢弃。从宏观的角度来看，相当于每次训练需要多两次 AllGather 操作。那么可能有人就会有疑问：那为什么通信量不是 4× 呢？

这是因为，由于模型参数恰好也分区了，所以我们不需要更新完参数后通过 AllGather 去共享模型参数了。我们只需要在更新参数前，使用 ReduceScatter 把对应的 Gradient 进行分发。

综上，总的通信量为 2×(forward + backward)+(gradient average)=3×，为标准 DP 通信量的 1.5 倍。

通信量可视化

优化效果

上图中Memory Consumption 第二列给出了一个示例： K=12,Φ=7.5B,N=64 ，可以看到显存优化相当明显。

在DeepSpeed中，对应ZeRO-1，对应ZeRO-2，对应ZeRO-3;

预训练中一般只使用ZeRO-1。因为zero3会增加额外通信, 本来通信就是瓶颈, 所以得不偿失; 而如果使用了GA(Gradient Accumulation), 由于GA也需要通信, 所以zero2也会额外带来通信需求;而需要一般都会开GA, 所以zero2也不常用

计算流程图解

官方给出了一个五分钟的解释视频，我们一张张截取看一下：

首先我们有一个16个Transformer块构成的模型，每一个块都是一个Transformer块。

有一个很大的数据集和四个GPU。

我们使用三阶段策略，将OPG和数据都进行拆分放在四张卡上。

每个模块下的格子代表模块占用的显存。第一行是FP16版本的模型权重参数，第二行是FP16的梯度，用来反向传播时更新权重，剩下的大部分绿色部分是优化器使用的显存部分，包含（FP32梯度，FP32方差，FP32动量，FP32参数）它只有在FP16梯度计算后才会被使用。ZeRO3使用了混合精度，因此前向传播中使用了半精度的参数。

每个模块还需要一部分空间用于存放激活值，也就是上面蓝色的部分。

每个GPU都会负责模型的一部分，也就是图中的

现在进入ZeRO3的一个分布式训练流程：

首先，GPU_0将自身已经有的模型部分权重通过broadcast发送到其他GPU。

当所有GPU都有了权重后，除了GPU_0以外的GPU会将他们存储在一个临时缓存中。

进行前向传播，每个GPU都会使用的参数在自己的进程的数据上进行前向传播，只有每个层的激活值会被保留。

计算完成后，其他GPU删除这部分的模型参数。

接下来，GPU_1将自己的模型权重参数广播发送到其他GPU。所有GPU上使用进行前向传播。

计算完成后，其他GPU删除这部分的模型参数。

依次类推，将每个GPU上的各自的模型权重都训练完。

前向传播结束后，每个GPU都根据自己数据集计算一个损失。

开始反向传播。首先所有GPU都会拿到最后一个模型分块（也就是）的损失。反向传播会在这块模型上进行，的激活值会从保存好的激活值上进行计算。

其他GPU将自己计算的的梯度发送给GPU_3进行梯度累积，最后在GPU_3上更新并保存最终的权重参数。

备注：梯度累积之前讲过，将几个小批次的数据的梯度累积，累加够一个大批次后更新模型权重。

其他GPU删除临时存储的权重参数和梯度，所有GPU都删除的激活值。

GPU_2发送参数到其他GPU，以便它们进行反向传播并计算梯度。

依次类推，直到每个GPU上自己部分的模型参数都更新完。

现在每个GPU都有自己的梯度了，开始计算参数更新。

优化器部分在每个GPU上开始并行。

优化器会生成FP32精度的模型权重，然后转换至FP16精度。

FP16精度的权重成为了下一个迭代开始时的模型参数，至此一个训练迭代完成。

总结一下，基本上就是把模型拆的更细了，原先的模型并行只是拆模型，现在不光拆模型，还把内部的优化器给拆了，并且只有在使用到的时候才会占据显存。

剩余状态优化

再来看剩余状态，也就是激活值（activation）、临时缓冲区（buffer）以及显存碎片（fragmentation）。

激活值同样使用分片方法，并且配合checkpointing

模型训练过程中经常会创建一些大小不等的临时缓冲区，比如对梯度进行AllReduce啥的，解决办法就是预先创建一个固定的缓冲区，训练过程中不再动态创建，如果要传输的数据较小，则多组数据bucket后再一次性传输，提高效率

显存出现碎片的一大原因是时候gradient checkpointing后，不断地创建和销毁那些不保存的激活值，解决方法是预先分配一块连续的显存，将常驻显存的模型状态和checkpointed activation存在里面，剩余显存用于动态创建和销毁discarded activation

通信开销分析

和 的通信量和传统数据并行相同，会增加通信量。

ZeRO-offload

一张卡训不了大模型，根因是显存不足，ZeRO-Offload的想法很简单：显存不足，内存来补。

ZeRO-Offload并不希望为了最小化显存占用而让系统的计算效率下降，否则的话，我们只用CPU和内存不就得了。但是将部分GPU的计算和存储下放到CPU和内存，必然涉及CPU和GPU之间的通信增加，不能让通信成为瓶颈，此外GPU的计算效率相比于CPU也是数量级上的优势，也不能让CPU参与过多计算，避免成为系统瓶颈，只有前两条满足的前提下，再考虑最小化显存的占用。

计算流程

现在的计算流程是，在GPU上面进行前向和后向计算，将梯度传给CPU，进行参数更新，再将更新后的参数传给GPU。为了提高效率，可以将计算和通信并行起来，GPU在反向传播阶段，可以待梯度值填满bucket后，一遍计算新的梯度一遍将bucket传输给CPU，当反向传播结束，CPU基本上已经有最新的梯度值了，同样的，CPU在参数更新时也同步将已经计算好的参数传给GPU，如下图所示。

offload策略

为了找到最优的offload策略，作者将模型训练过程看作数据流图（data-flow graph）。

圆形节点表示模型状态，比如参数、梯度和优化器状态

矩形节点表示计算操作，比如前向计算、后向计算和参数更新

边表示数据流向

下图是某一层的一次迭代过程（iteration/step），使用了混合精读训练，前向计算（FWD）需要用到上一次的激活值（activation）和本层的参数（parameter），反向传播（BWD）也需要用到激活值和参数计算梯度，

如果用Adam优化器进行参数更新（Param update），流程如下：

下面我们为边添加权重，物理含义是数据量大小（单位是字节），假设模型参数量是 M ，在混合精度训练的前提下，边的权重要么是2M（fp16），要么是4M（fp32），

我们现在要做的就是沿着边把数据流图切分为两部分，分布对应GPU和CPU，计算节点（矩形节点）落在哪个设备，哪个设备就执行计算，数据节点（圆形）落在哪个设备，哪个设备就负责存储，将被切分的边权重加起来，就是CPU和GPU的通信数据量。

ZeRO-Offload的切分思路是：

图中有四个计算类节点：FWD、BWD、Param update和float2half，前两个计算复杂度大致是 O(MB) ， B 是batch size，后两个计算复杂度是 O(M) 。为了不降低计算效率，将前两个节点放在GPU，后两个节点不但计算量小还需要和Adam状态打交道，所以放在CPU上，Adam状态自然也放在内存中，为了简化数据图，将前两个节点融合成一个节点FWD-BWD Super Node，将后两个节点融合成一个节点Update Super Node。如下图右边所示，沿着gradient 16和parameter 16两条边切分。

DeepSpeed ZeRO 分片与 Offload 策略对比

为了更好地理解 DeepSpeed 的 ZeRO 策略，以下对各阶段及 Offload 方案进行对比：

ZeRO Stage	描述	显存占用	训练速度
ZeRO-0	纯数据并行，不进行任何分片，所有状态在每个 GPU 上完全复制。	最高	最快
ZeRO-1	仅分片优化器状态，梯度和参数仍复制。	较高	略慢于 ZeRO-0
ZeRO-2	分片优化器状态和梯度。	中等	慢于 ZeRO-1
ZeRO-3	分片优化器状态、梯度和模型参数。	最低	明显慢于 ZeRO-2，受模型规模和网络带宽影响

Offload 类型	描述	显存占用	训练速度
ZeRO-1 + CPU Offload	在 ZeRO-1 基础上，将优化器状态卸载到 CPU 内存，降低 GPU 显存占用，但依赖 PCIe 带宽且占用 CPU 内存。	中偏低	慢于 ZeRO-1
ZeRO-2 + CPU Offload	在 ZeRO-2 基础上，将优化器状态卸载到 CPU 内存，对大模型进一步降低 GPU 显存，但增加 CPU–GPU 数据传输。	较低	慢于 ZeRO-2
ZeRO-3 + CPU Offload	在 ZeRO-3 基础上，将优化器状态和模型参数卸载到 CPU，GPU 显存占用最低，但 CPU–GPU 通信开销极大。	极低	非常慢
ZeRO-Infinity(NVMe Offload)	基于 ZeRO-3，将状态卸载到 NVMe 设备，突破 CPU 内存限制，适合超大模型；性能高度依赖 NVMe 并行读写速度。	极低需 NVMe 支持	慢于 ZeRO-3，但通常优于 CPU Offload 方案

参考链接

DeepSpeed之ZeRO系列：将显存优化进行到底 https://zhuanlan.zhihu.com/p/513571706

ZeRO: Zero Redundancy Optimizer，一篇就够了 https://zhuanlan.zhihu.com/p/663517415