ZeRO（Zero Redundancy Optimizer，零冗余优化器）

在将分布式并行DDP推向百亿、千亿参数大模型预训练的工业级舞台时传统的 DDP 因为“每张卡都要全量复制一份模型参数和优化器状态”的硬伤导致显存极易炸裂。为了打破这个僵局微软在 DeepSpeed 框架中提出了ZeROZero Redundancy Optimizer零冗余优化器。它是一种极其天才的、介于数据并行DDP和模型并行MP之间的分布式优化技术。它的核心思想是从宏观上看它依然是跑着高效的数据并行逻辑但在微观的显存存储上它把冗余的参数、梯度和优化器状态像切蛋糕一样“分片Shard”给不同的 GPU。一、 显存分析真正的“显存隐形杀手”是什么在训练大模型时单卡显存主要被两大部分吃掉激活值Activations和模型状态Model States。在 FP16 / BF16 混合精度训练中模型状态包含了三大块业内统称OSG 内存占用O (Optimizer States - 优化器状态)比如最常用的 AdamW每 1 个参数都需要额外维护一份 FP32 的 Master Weights、一份一阶动量Momentum和一份二阶方差Variance。光这一项每个参数就要吃掉12 字节显存G (Gradients - 梯度)FP16 格式每个参数吃掉2 字节。P (Parameters - 模型参数)FP16 格式每个参数吃掉2 字节。对于一个 7B70亿参数的模型光是这三项死数据不含前向传播动态产生的激活值单卡全量复制就需要7B×(1222) Bytes≈112 GB7\text{B} \times (12 2 2)\text{ Bytes} \approx 112\text{ GB}7B×(1222)Bytes≈112GB这已经远远超过了单张 A100/H10080GB的物理显存极限。二、 ZeRO 的三阶段切分ZeRO-1, ZeRO-2, ZeRO-3ZeRO 采用了层层递进的策略将这三部分数据逐步拆分到不同的 GPU 进程中。1. ZeRO-Stage 1: 优化器状态分片Optimizer State Sharding做法每张卡依然保留全量的 FP16 参数P和 FP16 梯度G。但是最庞大的 FP32Adam 优化器状态O被均匀切分成NNN份NNN为 GPU 总数/DDP 物理卡数每张卡只存其中的1/N1/N1/N。通信逻辑反向传播结束后各卡算出全量梯度。每张卡把自己负责的那1/N1/N1/N梯度通过Reduce-Scatter收集起来并在本地更新自己负责的那1/N1/N1/N优化器状态和 Master Weights。更新完后再通过一次All-Gather把更新后的 FP16 参数同步给所有卡。收益砍掉了显存中最大的一块砖显存需求直接暴跌至原来的约四分之一而几乎没有带来额外的通信开销因为其通信量和标准 DDP 的全量 All-Reduce 理论上是完全等价的。2. ZeRO-Stage 2: 梯度分片Gradient Sharding做法在 Stage 1 的基础上连FP16 梯度G也切碎。每张卡只负责存储和维护与它那部分优化器状态相对应的1/N1/N1/N梯度。通信逻辑反向传播时当某个 Transformer 层的梯度被算出来后各卡立刻通过Reduce-Scatter把梯度聚合到负责该分片的特定 GPU 上其余 GPU 立刻释放该层的梯度显存。收益显存进一步释放。Stage 1 Stage 2 已经是目前工业界微调Fine-tuning大模型最常用的标配组合既省显存通信效率又极高。3. ZeRO-Stage 3: 模型参数分片Parameter Sharding做法到了最激进的阶段连FP16 模型参数P也切碎分摊。任何一张卡在静止状态下手里的 P、G、O 都只有1/N1/N1/N。通信逻辑*前向传播时算到第lll层所有卡通过一次All-Gather把这一层存在别人那里的参数临时拉过来拼成完整的层把这一个 Batch 的前向算完算完立马把拉过来的参数擦除只留自己的1/N1/N1/N。反向传播时同样算到第lll层再通过 All-Gather 把该层参数拉过来配合梯度算完反向算完再次擦除。收益显存占用直接和 GPU 数量NNN成反比这意味着只要你的集群卡足够多你就能无限容纳下任意大的模型。代价带来了约 1.5 倍的额外跨卡通信开销。在大规模千卡集群上极度依赖高性能的机间网络InfiniBand 或者是高带宽的 RoCEv2 拓扑。三、 ZeRO-Offload拿内存换显存的降维打击如果你的手里只有单机多卡或者消费级显卡比如 4 张 RTX 4090哪怕开了 ZeRO-3 还是可能 OOM微软还附赠了另一个大杀器Offload 机制。核心思想GPU 显存昂贵且有限但宿主机机箱里的CPU 内存Host Memory和NVMe 固态硬盘空间非常大且便宜。工作流* 把暂时用不到的优化器状态O甚至模型参数P借用 PCIe 通道换出Offload到 CPU 内存里。甚至可以直接把优化器更新参数的计算逻辑AdamW 的矩阵加减乘除交由 CPU 去算算完再把更新后的参数塞回 GPU。最新演进现在的 ZeRO-Infinity 技术甚至可以 offload 到 NVMe 固态硬盘理论上在一台普通的单机服务器上就能加载万亿参数的模型。四、 工业界落地PyTorch 原生 FSDP因为 ZeRO 实在是太好用了PyTorch 官方在 1.11 之后直接将其“抄”进了核心源码并进行了原生重构命名为FSDPFully Sharded Data Parallel完全分片数据并行。对应关系PyTorch FSDP 的默认全分片模式FULL_SHARD完美等价于 DeepSpeed 的ZeRO-3其SHARD_GRAD_OP模式则等价于ZeRO-2。技术优势FSDP 完美嵌入了 PyTorch 的nn.Module原生树状结构中能更精细地控制每一个子模块Sub-module在反向传播时的通信时机通过wrapping_policy在与 PyTorch 原生的自动混合精度AMP以及编译加速torch.compile联动时比外挂的 DeepSpeed 拥有更好的软硬件融合度与更少的 CPU Launch 延迟。 总结与性能调优思维在具体做分布式基础设施部署时选择哪一个 Stage 往往取决于你手头的硬件网络拓扑这也是Amdahl 定律的体现单机内部 / NVLink 环境放心大胆地开 ZeRO-2 或 ZeRO-3。因为单机内部卡间带宽极高如 NVLink 可达数百 GB/sAll-Gather 和 Reduce-Scatter 的通信延迟可以被极致的计算完全掩盖。跨机千卡集群 / 走普通千兆/万兆网卡尽量只开ZeRO-1 或 ZeRO-2。如果强行上 ZeRO-3你会发现由于网卡带宽太低GPU 绝大部分时间都在等网络把参数拉过来Communication-Bound整体的硬件利用率MFU会跌到惨不忍睹。这时候应该改用Megatron-LM 的 3D 混合并行TPPP去裁剪模型而不是纯依赖 ZeRO-3 的完全切分。