当前位置：首页 > news >正文

大模型显示优化之ZeRO-1/ZeRO-2/ZeRO-3

news 2026/6/3 5:46:15

1. 简介

zero-1、zero-2、zero-3 是deepspeed的配置方法，对应megatron也有相应的方法，Megatron-LM 的实现方式：Distributed Optimizer（分布式优化器）。等效于 ZeRO-1，Megatron 的 Distributed Optimizer 默认行为就是将优化器状态（Optimizer States）均匀地切分并分布在数据并行（DP）组的所有 GPU 上。等效于 ZeRO-2，由于 Megatron 通常结合混合精度训练，它在计算完梯度后，会通过Reduce-Scatter操作直接将梯度同步并切分到各卡上，不再保留全量梯度。这在效果上完全等同于 ZeRO-2。zero-3 将参数也拆分卡来存，但后续实际反向梯度更新时操作时还是需要all-gather，参数显存还是会全量缓存，再一个Megatron针对参数拆分，更多使用的是TP/PP拆分，所以业界megatron架构使用zero-3不多, 所以本文不做重点分析。

Zero架构说的是DP并行域GPU之间。

阶段	优化对象	核心原理	效果
ZeRO-1	优化器状态 (OS)	将优化器状态切分并分布到各个 GPU 上，每个 GPU 只负责更新自己那一块。	显存占用降低约为原来的 1/4（以 Adam 为例）。
ZeRO-2	OS + 梯度 (G)	在 ZeRO-1 基础上，进一步将梯度也进行切分。每个 GPU 只保留对应参数的梯度。	进一步降低显存占用，是目前最常用的平衡配置。
ZeRO-3	OS + G + 参数 (W)	最彻底的切分。模型参数在平时也分布在不同 GPU 上，只有在正向/反向传播需要时才临时同步。	显存占用理论上随 GPU 数量线性下降，支持训练超大规模模型。

实际官方Megatron实现中，ZeRO-2 反向不只是对梯度进行切分，还对参数在back阶段进行了小段时间的切分，后面AllGather回收，是一个技术操作。这样好处：

1. 节省显存

2. 避免冗余计算

3. 最后的AllGather可以和后续的layer forward 做overlap

纯 DP	ZeRO-2
Forward	各 rank 用完整 W	各 rank 用完整 W（相同）
Backward 后通信	AllReduce梯度（每人拿完整梯度）	ReduceScatter梯度（每人只拿 1/DP 梯度，显存也只存1/DP属于自己的梯度）
Optimizer step	各自完整更新 W（结果一致，冗余计算）	各自只更新 W 的 1/DP 段(此更新过程比较复杂）
Step 后	无需额外通信（W 天然一致）	需要AllGather W恢复完整参数
显存节省	无	梯度 + 优化器状态各节省 1/DP

注意：

AdamW全局grad_norm

路径	通信方式	时机
标准路径	`all_reduce`on model parallel group（TP × DP）	optimizer.step() 内，clip grad 前
PP bypass 路径	TP 内`all_reduce`+ PP 间`send/recv`逐 stage 累加	pre_step 阶段，流水线化减少同步 barrier

AdamW 的step()中确实有一次全局 grad norm 的all_reduce通信，用于计算全局 L2 norm 以确定clip_coeff（梯度裁剪系数）。这是每一步更新都必须做的集合通信，会引入跨所有 model parallel rank 的同步点。

2. 显存与通信量分析

为了让 ZeRO-1 和 ZeRO-2 的区别更加直观，我把之前流程图里的抽象内容，具体化成了4 张 GPU 卡在不同阶段的显存状态。

这样你可以像看“快照”一样，清晰地看到每张卡上到底存了什么。

设定：假设模型有4个参数块：[P0, P1, P2, P3]。4 张 GPU 卡训练。FP16训练的模型为例，参数量为：
参数 (Weights):字节。
梯度 (Grads):字节。
优化器Adam 状态:
FP32 权重副本（为了精度）：。
Momentum（动量）：。
Variance（方差）：。

场景一：ZeRO-1 (只切分优化器状态)

核心特征：每张卡都有完整的参数和完整的梯度，但只负责更新1/4的优化器状态。

GPU 卡	前向/反向计算时	梯度通信后 (All-Reduce)	参数更新后
GPU 0	参数:`[P0, P1, P2, P3]` 梯度:`[G0, G1, G2, G3]` 优化器状态:`[O0]`(只负责P0)	梯度:`[G_avg0, G_avg1, G_avg2, G_avg3]` (已同步为平均梯度)	*用`G_avg0`更新`O0`，计算出`P0_new` 然后拼出完整参数`[P0_new, P1_new, P2_new....]`
GPU 1	参数:`[P0, P1, P2, P3]` 梯度:`[G0, G1, G2, G3]` 优化器状态:`[O1]`(只负责P1)	梯度:`[G_avg0, G_avg1, G_avg2, G_avg3]`	*用`G_avg1`更新`O1`，计算出`P1_new` 然后拼出完整参数`[P0_new, P1_new, P2_new....]`

显存占用：

高。因为每张卡都要存下4份参数 + 4份梯度。
冗余度高。P0被同时存在了 4 张卡上。

场景二：ZeRO-2 (切分梯度 + 优化器状态)

核心特征：每张卡有完整的参数，但只保留1/4的梯度，并只更新对应的1/4优化器状态。

GPU 卡	前向/反向计算时 (初始)	梯度通信后 (Reduce-Scatter)	参数更新后
GPU 0	参数:`[P0, P1, P2, P3]` 梯度(原始):`[G0, G1, G2, G3]` 优化器状态:`[O0]`	梯度(保留):`[G_avg0]` 梯度(丢弃):`[G_avg1, G_avg2, G_avg3]`✔️ 丢弃	用`G_avg0`更新`O0`，计算出`P0_new`。然后通过 All-Gather 从其他卡获取 P1~P3 的更新。
GPU 1	参数:`[P0, P1, P2, P3]` 梯度(原始):`[G0, G1, G2, G3]` 优化器状态:`[O1]`	梯度(保留):`[G_avg1]` 梯度(丢弃):`[G_avg0, G_avg2, G_avg3]`✔️ 丢弃	用`G_avg1`更新`O1`，计算出`P1_new`。然后通过 All-Gather 从其他卡获取 P0, P2, P3 的更新。

显存占用：
中等。每张卡存4份参数 + 1份梯度。
显存优化：相比 ZeRO-1，节省了 3 份梯度的存储空间。

两张图的对比总结

特征	ZeRO-1 (图里场景)	ZeRO-2 (图里场景)
每张卡上的参数	全部`[P0, P1, P2, P3]`	全部`[P0, P1, P2, P3]`
每张卡上的梯度	全部`[G_avg0...G_avg3]`(All-Reduce后)	只有1块`[G_avg0]`(Reduce-Scatter后)
优化器状态	分片`[O0]`	分片`[O0]`
参数更新方式	各卡独立计算出完整参数	各卡计算部分参数，再互相广播合并
主要节省	不节省梯度	节省了3/4的梯度显存

通过这两张“快照”，你应该能清晰地看到：ZeRO-2 的本质，就是用梯度通信后的一个“丢弃”动作，换来了大量的显存空间。

通信量总结

维度	ZeRO-1	ZeRO-2	ZeRO-3
参数存储	完整 (每卡都有)	完整 (每卡都有)	切分(每卡1/DP)
梯度存储	完整 (每卡都有)	切分(每卡1/DP)	切分(每卡1/DP)
优化器状态	切分 (每卡1/DP)	切分 (每卡1/DP)	切分 (每卡1/DP)
单卡模型状态显存	`2Ψ + 2Ψ + 12Ψ/DP`	`2Ψ + 2Ψ/DP + 12Ψ/DP`	`(2Ψ+2Ψ+12Ψ)/DP`
主要通信	All-Reduce (梯度)	Reduce-Scatter + All-Gather	All-Gather ×2 + Reduce-Scatter
通信量	2×Ψ(最小)	2×Ψ	3×Ψ(最大)
显存节省	仅优化器状态	优化器+梯度	全部

3. Megatron ZeRO配置

Stage	分片内容	Megatron对应参数
ZeRO-1	优化器状态分片（m,v)	--user-distributed-optimizer
ZeRO-2	优化器分片+梯度分片	--user-distributed-optimizer+ --overlap-grad-reduce
ZeRO-3	优化器分片+梯度+参数	需要单独搞