ZeRO 优化器的三个阶段深度对比：显存节省与通信开销的量化分析

ZeRO 优化器的三个阶段深度对比显存节省与通信开销的量化分析一、分布式训练的内存瓶颈与 ZeRO 的设计思路大规模模型训练面临的核心约束是单卡显存容量。以GPT-3规模的模型175B参数为例仅FP16参数就占用350GB显存加上Adam优化器状态FP32的参数副本、一阶动量m和二阶动量v约2.8TB再加上梯度存储和激活值单卡A100-80G完全无法承载。数据并行虽然能将计算分散到多卡但每张卡仍需保存完整的模型状态副本显存瓶颈未被消除。DeepSpeed提出的ZeROZero Redundancy Optimizer通过对模型状态进行分片来消除数据并行中的冗余存储。模型状态包括三类数据优化器状态Optimizer States、梯度Gradients和模型参数Parameters。ZeRO的三个阶段分别对这三类数据进行不同程度的分片从而在显存节省与通信开销之间取得不同的平衡点。具体而言ZeRO-1 仅对优化器状态进行分片可实现约 4 倍的显存节省且通信开销与传统数据并行保持一致ZeRO-2 在此基础上进一步分片梯度显存节省提升至约 8 倍通信开销依然不变而 ZeRO-3 则对优化器状态、梯度及模型参数三者全部分片显存节省与数据并行度呈线性关系但通信开销会增加约 50%。二、ZeRO-1 与 ZeRO-2优化器状态与梯度的分片ZeRO-1仅对优化器状态进行分片。在标准数据并行中每个rank都持有完整的优化器状态对于Adam而言是12N字节N为参数量。ZeRO-1将优化器状态均匀划分到所有数据并行进程每个rank只持有总优化器状态的1/NdNd为数据并行度。在反向传播完成后每个rank对其持有的参数子集执行优化器步骤即更新参数然后通过all-gather将更新后的参数广播给所有rank。ZeRO-2在ZeRO-1的基础上进一步分片梯度。在反向传播过程中每个rank计算完整梯度后通过reduce-scatter操作将梯度归约并分散到对应rank——每个rank只保留其负责的那部分参数的归约后梯度其余梯度直接丢弃。由于梯度存储从完整N变为N/Nd显存进一步降低。两者的通信量与传统数据并行完全相同ZeRO-1在优化器更新后需要一次all-gather通信量约为2NZeRO-2在梯度计算后需要一次reduce-scatter通信量约为2N。通信开销不增加是这两个阶段的显著优势。 模拟ZeRO-1和ZeRO-2的显存占用对比不同模型规模下的节省效果 import numpy as np from dataclasses import dataclass dataclass class MemoryEstimation: 模型状态显存估算结果单位GB optimizer_states: float gradients: float parameters: float total: float def estimate_memory_per_gpu( num_params: int, # 模型参数量 dp_degree: int, # 数据并行度 bytes_per_param: int 2, # FP16参数2字节 bytes_per_opt_state: int 12, # AdamFP32参数副本mv 4*312字节 ) - dict[str, MemoryEstimation]: 估算不同ZeRO阶段下的单卡显存占用。 标准数据并行DP每卡存储全部模型状态 ZeRO-1仅分片优化器状态 ZeRO-2分片优化器状态 梯度 ZeRO-3分片全部模型状态 N num_params Nd dp_degree # 基础存储量字节 param_bytes N * bytes_per_param grad_bytes N * bytes_per_param opt_bytes N * bytes_per_opt_state def to_gb(b): return b / (1024 ** 3) return { DP: MemoryEstimation( optimizer_statesto_gb(opt_bytes), gradientsto_gb(grad_bytes), parametersto_gb(param_bytes), totalto_gb(opt_bytes grad_bytes param_bytes) ), ZeRO-1: MemoryEstimation( optimizer_statesto_gb(opt_bytes / Nd), # 仅分片OS gradientsto_gb(grad_bytes), # 梯度完整保存 parametersto_gb(param_bytes), # 参数完整保存 totalto_gb(opt_bytes / Nd grad_bytes param_bytes) ), ZeRO-2: MemoryEstimation( optimizer_statesto_gb(opt_bytes / Nd), # 分片OS gradientsto_gb(grad_bytes / Nd), # 分片梯度 parametersto_gb(param_bytes), # 参数完整保存 totalto_gb(opt_bytes / Nd grad_bytes / Nd param_bytes) ), ZeRO-3: MemoryEstimation( optimizer_statesto_gb(opt_bytes / Nd), gradientsto_gb(grad_bytes / Nd), parametersto_gb(param_bytes / Nd), # 分片参数 totalto_gb((opt_bytes grad_bytes param_bytes) / Nd) ), } # 示例估算175B参数模型在64卡A100上的显存占用 # results estimate_memory_per_gpu(num_params175e9, dp_degree64) # for stage, mem in results.items(): # print(f{stage}: 总计 {mem.total:.2f} GB) ## 三、ZeRO-3参数分片的通信代价量化 ZeRO-3 将模型参数也进行分片每个 rank 只持久化保存 1/Nd 的参数。这带来了近线性的显存节省——理想情况下Nd 张卡上的总参数存储量等于单卡的存储量。然而这种极致的显存节省以增加通信为代价。 在 ZeRO-3 的前向传播中当某个层需要其参数时对应的 rank 通过 all-gather 从所有 rank 收集该层的完整参数。具体而言假设 Rank 0 需要执行某层计算它会向其他所有 Rank 发起 all-gather 请求持有该层参数分片的 Rank 1、Rank 2 及 Rank 3 随即发送各自持有的分片数据。Rank 0 收到所有分片后重组为完整参数并执行前向计算计算完成后立即释放非本地参数分片仅保留自己负责的部分。反向传播同样需要 all-gather 参数和 all-gather 梯度。额外的通信量约为前向传播中各层参数的 all-gather 操作总和在典型的 Transformer 模型中约为 1.5 倍标准数据并行的通信量。 通信开销的接受度取决于训练规模对于 175B 级别的模型在 64 卡以上训练ZeRO-3 几乎是唯一可行的方案对于 7B 级别的小模型ZeRO-2 的显存已经足够增加参数分片带来的通信开销往往得不偿失。 ## 四、阶段选择的决策框架 ZeRO 阶段的选择并非越高越好而是需要在显存约束、通信带宽和实现复杂度之间进行权衡。以下决策树可以作为选择依据 首先判断单卡能否容纳完整模型。如果可以例如 7B 模型在 A100-80G 上使用 ZeRO-1 即可仅对 32GB 的优化器状态进行分片通信开销为零。 如果单卡无法容纳完整模型但可以容纳参数梯度例如 13B 模型在 A100-80G 上使用 ZeRO-2 是最优选择——它不会引入额外通信但将梯度也从单卡中卸载。 如果单卡连参数都无法容纳例如 70B 模型ZeRO-3 是唯一选择。此时建议配合使用**offload**策略ZeRO-Offload 或 ZeRO-Infinity将优化器状态进一步卸载到 CPU 内存甚至 NVMe 存储。 此外ZeRO-3可以与**模型并行**Tensor Parallelism、Pipeline Parallelism组合使用形成3D并行策略。典型实践是在节点内使用张量并行利用NVLink高带宽在节点间使用ZeRO-3数据并行利用Infiniband流水线并行用于跨越更多节点时的通信优化。 ## 五、总结 ZeRO-1、ZeRO-2、ZeRO-3构成了一个从保守到激进的显存优化谱系。ZeRO-1以零额外通信代价节省优化器状态存储适用面最广ZeRO-2在梯度和优化器状态上同时分片将显存节省翻倍而不增加通信是大多数中等规模训练的最佳实践ZeRO-3通过参数分片实现近线性的显存扩展代价是每层前反向各增加一次all-gather通信适用于超大模型的多机多卡训练。三者之间的选择本质是显存约束 vs 通信带宽的工程权衡不存在普遍最优解——最优方案总是针对特定模型规模、硬件拓扑和训练预算定制的结果。