ZeRO Optimization（零冗余优化）

ZeRO（Zero Redundancy Optimizer）是专为超大规模语言模型训练设计的显存优化框架。

• 核心动机：标准数据并行要求每张 GPU 保存完整的模型状态副本，当模型参数达到百亿级时，单卡显存远不够用
• 解决思路：将冗余的三类模型状态（优化器状态、梯度、参数）改为分片存储，每个进程只持有自己负责的那一份
• 效果：显存占用随 GPU 数量近线性降低，通信开销可控，与数据并行在代码层几乎无缝集成
• 归属：由微软 DeepSpeed 团队于 2019 年提出，已成为 HuggingFace、PyTorch FSDP 等主流框架的基础

ZeRO 的核心操作是分片（Partitioning）与按需聚合（AllGather）。

• 前向传播：各进程通过 AllGather 拼回完整参数，完成前向计算后释放非本分片的参数
• 反向传播：梯度通过 ReduceScatter 聚合后，每个进程只保留属于自己分片的梯度
• 参数更新：各进程使用自己持有的优化器状态和梯度更新对应的参数分片
• 通信原语：主要使用 AllGather 和 ReduceScatter，ZeRO-3 通信量约为标准数据并行的 1.5×，ZeRO-1/2 通信量与数据并行持平
• 关键区别：参数在需要时才被拼回，不使用时即丢弃，从根本上打破「每卡必须存全量」的限制

ZeRO 按分片粒度分为三个递进阶段，显存节省逐级增大。

• ZeRO-1：仅分片优化器状态（如 Adam 的 32-bit 权重副本、一阶/二阶矩），显存降至标准数据并行的约 1/4，通信量不变
• ZeRO-2：在 ZeRO-1 基础上同时分片梯度，显存降至约 1/8，通信量仍与数据并行相同，实用性最强
• ZeRO-3：对模型参数也做分片，显存理论上随 GPU 数线性扩展，通信量略增至 1.5×；PyTorch 的 FSDP（Fully Sharded Data Parallel）即参考此阶段
• ZeRO-Infinity：基于 ZeRO-3，将分片进一步卸载至 CPU 内存或 NVMe SSD，可在消费级显卡上训练千亿参数模型
• ZeRO-R：补充性优化，针对激活值（Activation）和临时缓冲区进一步压缩显存

ZeRO 广泛应用于需要大显存的大模型训练与微调场景。

• 预训练超大模型：如 GPT-3（175B）、Megatron-Turing NLG（530B）等均使用 DeepSpeed ZeRO 训练
• 大模型全量微调（Full Fine-tuning）：显存受限时以 ZeRO-2/3 替代 LoRA 等参数高效方法
• 多机多卡集群：配合 InfiniBand 高速互联，ZeRO-3 在数百张 A100/H100 上线性扩展
• 消费级硬件训练：ZeRO-Infinity + CPU offload 让 RTX 3090 也能微调 30B 模型
• 与 3D 并行结合：ZeRO（数据维度） + Tensor Parallelism（算子维度） + Pipeline Parallelism（层维度）组成工业级训练方案

ZeRO 常与模型并行、混合精度等技术混淆，需要明确区分。

• vs 数据并行（Data Parallel）：ZeRO 本质是增强版数据并行，但去除了每卡存全量副本的冗余；传统 DDP 每卡必须存完整模型状态
• vs 模型并行（Tensor/Pipeline Parallelism）：模型并行将不同层或算子切分到不同 GPU，计算图本身被拆分；ZeRO 不改变计算图，仅在存储层分片
• vs 混合精度（AMP）：AMP 通过 FP16 计算减少显存，ZeRO 通过分片减少冗余；两者可叠加使用
• vs Gradient Checkpointing：梯度检查点用重计算换显存，ZeRO 用通信换显存，侧重点不同，同样可组合

ZeRO 并非万能，在特定场景下存在明显局限。

• 通信带宽敏感：ZeRO-3 需要频繁 AllGather，低带宽互联（如千兆以太网）会成为严重瓶颈，吞吐量可能不如 ZeRO-2
• 单卡场景无效：ZeRO 依赖多进程分片，单 GPU 场景没有收益
• 显存并非无限：ZeRO-3 显存节省是相对的，激活值（Activation Memory）随 batch size 增大不受 ZeRO 控制，仍需梯度检查点配合
• 常见误解：认为 ZeRO 等同于「模型并行」；实际上 ZeRO 是数据并行的优化，计算仍在每张卡上完整进行（ZeRO-3 通过 AllGather 临时拼回参数）
• CPU offload 代价：ZeRO-Infinity 的 NVMe 卸载会引入大量 I/O 延迟，训练速度可能大幅下降

ZeRO 从提出到演进经历了多个里程碑。

• 2019年10月：微软 Rajbhandari 等人发布 arXiv 预印本（arXiv:1910.02054），提出 ZeRO-1/2/3 三阶段框架
• 2020年2月：微软开源 DeepSpeed 库，ZeRO 作为核心功能随之发布，迅速被学术界与工业界采用
• 2020年（SC20）：ZeRO 论文正式发表于国际超算大会（SC20），获广泛引用
• 2021年：发布 ZeRO-Infinity（arXiv:2104.07857），支持 NVMe 卸载，训练万亿参数模型
• 2021年（PyTorch）：PyTorch 参考 ZeRO-3 推出 FSDP（Fully Sharded Data Parallel），成为官方标准 API
• 2022年至今：ZeRO 思想被广泛移植，Megatron-LM、ColossalAI、MindSpore 等框架均实现类似机制，成为大模型训练的事实标准