Zero Redundancy Optimizer：内存节约型优化器

Zero Redundancy Optimizer：内存节约型优化器

在当前大模型参数规模动辄上百亿、上千亿的背景下，训练这些庞然大物早已不再是单卡甚至单机能够胜任的任务。显存瓶颈成了横亘在每一个开发者面前的一道高墙——哪怕你拥有 A100 或 H100 这样的顶级 GPU，面对 Llama3-70B 或 Qwen-72B 这类模型时，依然可能连一次前向传播都跑不起来。

正是在这种“算力追不上模型膨胀速度”的现实压力下，Zero Redundancy Optimizer（ZeRO）应运而生。它不是某种魔法，而是一种系统性的显存“瘦身”策略，由微软 DeepSpeed 团队提出并实现，核心思想简单却极具颠覆性：既然每张 GPU 都保存完整的优化器状态、梯度和参数是浪费，那就把它们拆开，分着存。

这一思路直接打破了传统数据并行中“每个设备全量复制”的固有模式，使得原本需要数十张高端 GPU 才能完成的训练任务，在更少设备上成为可能。更重要的是，随着 ms-swift 等现代训练框架对 ZeRO 的原生集成，用户不再需要深入底层通信机制，只需一个配置文件或几行命令，就能启用这套强大的显存优化体系。

显存去哪了？传统数据并行的“冗余之痛”

要理解 ZeRO 的价值，得先看清楚问题出在哪。

假设我们有一个 70B 参数的模型，使用 Adam 优化器进行 FP16 训练。那么仅在单卡上：

• 模型参数本身占用约 140GB 显存（70B × 2 bytes）
• 梯度再占 140GB
• Adam 的动量和方差各占 140GB，合计 280GB

三项加起来就是560GB——这还只是单卡！如果采用传统的数据并行方式，在 4 卡环境下，每个设备都要保存完整副本，总显存需求接近2.2TB，即便所有数据都能压缩到 GPU 上，也远远超出了任何现有硬件的能力。

但关键在于：这么多重复的数据真的有必要吗？

答案是否定的。训练的本质是通过分布式 batch 计算梯度，最终聚合更新参数。每个 GPU 只需负责自己那份数据对应的参数更新即可，并不需要持有全部状态。正是基于这个洞察，ZeRO 提出了分阶段消除冗余的设计哲学。

ZeRO 的三级进阶：从状态分片到参数卸载

ZeRO 并非一蹴而就的技术，而是分为三个演进阶段，逐级削减显存占用：

第一阶段：ZeRO-1 —— 优化器状态分片

最开始的冗余来自优化器。比如 Adam 中每个参数都有两个辅助状态（动量 $m$ 和方差 $v$），这部分通常比模型参数还大。ZeRO-1 的做法很简单：把这些状态按数据并行维度切分成 N 份，每张卡只保留属于自己那份 batch 所需的状态。

前向和反向传播仍使用完整模型，但在参数更新时，只更新本地负责的部分。其他参数则通过AllGather动态获取。

显存节省：约减少 $1/2 \sim 2/3$，具体取决于优化器类型。

第二阶段：ZeRO-2 —— 梯度也分片

ZeRO-2 在前者基础上进一步将梯度进行分片。反向传播完成后，各卡上的梯度经过ReduceScatter被拆分，每张卡仅保留对应分片的梯度和优化器状态。

此时，每张卡维护：
- 完整模型参数（仍需全量加载）
- 局部梯度
- 局部优化器状态

显存再降约 50%。对于大模型而言，这意味着原本只能用 8 卡跑的任务，现在 4 卡也能扛住。

第三阶段：ZeRO-3 —— 参数分片登场

这才是真正的“杀手锏”。ZeRO-3 不再要求每张卡保存完整模型参数。相反，参数也被分片存储，每个设备仅缓存当前计算所需的那一小部分，其余参数通过运行时AllGather实时拉取，用完即释放。

这就意味着：即使单卡装不下整个模型，也能参与训练。

想象一下，你要读一本十万页的书，但手边只有一个能放十页的小桌子。传统做法是你得先把整本书搬进来；而 ZeRO-3 的逻辑是：只把当前要看的那几页拿过来，看完换下一批。虽然翻页多了点，但至少你能看完这本书。

实测表明，配合 CPU Offload 后，ZeRO-3 可将 BLOOM-176B 的训练显存从数千 GB 压缩至数百 GB，真正实现了“万亿参数可训”。

当然，天下没有免费的午餐。越高级别的分片意味着越多的通信操作。尤其是 ZeRO-3，频繁的AllGather会带来显著延迟。因此工程实践中必须辅以通信优化手段，否则性能反而下降。

如何用？配置即生效

好在 DeepSpeed 把这一切封装得足够友好。你不需要手动写 NCCL 通信逻辑，也不必重构模型结构，只需要一个 JSON 配置文件，就能激活 ZeRO 的全部能力。

{ "train_batch_size": "auto", "train_micro_batch_size_per_gpu": "auto", "gradient_accumulation_steps": "auto", "optimizer": { "type": "AdamW", "params": { "lr": "auto", "weight_decay": "auto" } }, "fp16": { "enabled": "auto" }, "bf16": { "enabled": "auto" }, "zero_optimization": { "stage": 3, "offload_optimizer": { "device": "cpu", "pin_memory": true }, "offload_param": { "device": "cpu", "pin_memory": true }, "overlap_comm": true, "contiguous_gradients": true, "sub_group_size": 1e9 }, "steps_per_print": 2000, "wall_clock_breakdown": false }

几个关键字段值得细说：

• "stage": 3：启用 ZeRO-3，三重分片齐上阵。
• offload_*：开启 CPU 卸载，把优化器状态或参数暂存到主机内存，进一步缓解 GPU 压力（这就是所谓的ZeRO-Infinity）。
• overlap_comm: 允许通信与计算重叠，利用 GPU 空闲周期传输数据，有效掩盖延迟。
• contiguous_gradients: 将梯度连续存储，提升ReduceScatter效率。

保存为ds_config.json后，配合如下命令即可启动训练：

deepspeed --num_gpus=4 train.py \ --deepspeed ds_config.json \ --model_name_or_path bigscience/bloom-7b1

如果你用的是 ms-swift 这类高层框架，体验更接近“一键微调”——选模型、勾选项、点运行，背后的 ZeRO 策略自动匹配最优配置。

工程实践中的权衡艺术

尽管 ZeRO 强大，但并非“Stage 越高越好”。实际部署时，需要根据硬件条件和任务目标做精细权衡。

什么时候该用 ZeRO-3？

当你面对的是>30B 参数的大模型，且单卡显存无法容纳完整参数时，ZeRO-3 几乎是唯一选择。例如在 4×A100（80GB）上微调 Llama3-70B，若不用参数分片，根本无法启动训练。

但代价也很明显：通信密集。特别是在千兆以太网环境下，AllGather成为性能瓶颈。推荐搭配 InfiniBand 或 RoCE 网络使用。

何时退回到 ZeRO-2？

如果模型在 10–30B 范围内，且 GPU 显存尚可接受（如 A100×2 可勉强放下 Qwen-14B），建议优先使用 ZeRO-2。它保留了完整的模型参数副本，避免了频繁通信，整体吞吐更高。

我见过不少团队盲目开启 ZeRO-3 导致训练速度下降 40% 的案例——省了显存，却丢了效率。

混合精度不可少

无论哪个 stage，都应结合 bf16 或 fp16 使用。不仅显存减半，还能提升计算效率。注意开启 loss scaling 防止梯度下溢，尤其在低精度+小 batch 场景下。

监控通信占比

DeepSpeed 提供了wall_clock_breakdown: true配置项，可用于分析训练时间分布。理想情况下，通信时间不应超过总耗时的 20–30%。一旦超过，说明网络已成为瓶颈，需调整 micro-batch size、拓扑结构或升级网络设备。

与其他技术的协同效应

ZeRO 的真正威力，体现在它能无缝融合于现代大模型训练生态。

+ LoRA / QLoRA：双重压缩

LoRA 的思路是从参数角度做减法——冻结主干，只训练低秩适配矩阵。而 ZeRO 是从内存管理角度优化布局。两者结合堪称“黄金搭档”：

• LoRA 减少了可训练参数量；
• ZeRO 进一步压缩优化器状态和梯度存储；
• 最终可在消费级多卡环境完成百亿模型微调。

实测显示，QLoRA + ZeRO-3 组合可在 2×A100 上微调 65B 模型，显存峰值控制在 70GB 以内。

+ FSDP / Megatron-LM：混合并行扩展

ZeRO 主打数据并行内的显存优化，而 FSDP 和 Megatron-LM 支持张量并行、流水线并行。三者可以组合形成3D 并行训练架构，适用于千亿级以上模型。

例如：
- 流水线并行划分模型层；
- 张量并行切分注意力头或 FFN 层；
- ZeRO 处理数据并行中的状态分片；

这种多层次拆解策略，是目前训练最大模型的标准范式。

+ vLLM 推理加速：闭环落地

训练之后是推理。ms-swift 等平台已支持从 ZeRO 训练 → 权重合并 → vLLM 加速推理的完整链路。特别是当使用 CPU Offload 时，记得在训练结束后执行zero_to_fp32.py工具将分片状态合并回单一 checkpoint，否则无法直接加载。

解决了哪些真实痛点？

痛点一：显存爆炸，根本跑不起来

这是最常见的问题。很多开源模型 FP16 加载即超限。解决方案就是ZeRO-3 + CPU Offload。把大部分参数“扔”到内存里，GPU 只留活跃分片。虽然带宽低了些，但至少能跑。

我在某次实验中用 4×A100 微调 Baichuan2-70B，原始方案 OOM，启用 ZeRO-Infinity 后顺利收敛，显存稳定在 75GB 左右。

痛点二：成本太高，公司扛不住

企业往往不愿投入几十张 V100/A100。而 ZeRO 能让资源利用率翻倍。原来需要 64 卡的任务，现在 16 卡 + 更高利用率即可完成。云上按小时计费的情况下，节省可达数万元。

痛点三：实验迭代慢，调参像玄学

RLHF 阶段常需尝试 DPO、KTO、ORPO 等多种算法。ms-swift 提供图形化界面，内置 ZeRO 模板，用户只需勾选“启用 DeepSpeed”，系统自动生成最优配置，极大缩短验证周期。

架构视角：ZeRO 在系统中的位置

在一个典型的训练平台上（如基于 ms-swift 构建的系统），ZeRO 处于分布式训练引擎的核心层：

+----------------------------+ | Application Layer | | (Training Script, | | LoRA/QLoRA Config) | +------------+---------------+ | +--------v--------+ | Framework Layer | | (ms-swift / | | DeepSpeed) | +--------+--------+ | +--------v--------+ | ZeRO Optimization| | Engine (Stage 1/2/3)| +--------+--------+ | +--------v--------+ | Communication | | Backend (NCCL) | +-----------------+ | +--------v--------+ | Hardware Layer | | (A100/H100 GPUs,| | InfiniBand) | +-----------------+

它介于高层训练逻辑与底层通信库之间，既不影响模型定义，又能深度介入显存调度。这种“透明性”正是其广受欢迎的关键。

结语：显存优化的未来方向

ZeRO 的出现，标志着大模型训练从“拼硬件”转向“拼效率”的时代。它让我们意识到：提升资源利用率，有时比增加资源本身更重要。

未来，这一理念还将继续演进：

• MoE + ZeRO：稀疏激活特性天然适合分片管理；
• 异构卸载：GPU + CPU + NVMe 分层存储，实现更大规模 offload；
• 动态分片：根据计算图自动识别活跃参数，按需加载；
• 编译器级优化：结合 TorchDynamo 或 Triton，实现更智能的内存规划。

可以预见，随着模型持续增大，ZeRO 类技术不会被淘汰，反而会更深地融入训练栈底层，成为像“内存回收”一样的默认机制。

而对于开发者来说，最好的消息或许是：你不必成为分布式专家，也能训练超大模型。