目录DeepSpeed ZeRO 的设计动机ZeRO-1优化器状态分片ZeRO-2梯度分片ZeRO-3全参数分片ZeRO-Offload 与卸载DeepSpeed ZeRO 的边界与失效模式摘要DeepSpeed ZeROZero Redundancy Optimizer通过分阶段消除数据并行中的冗余存储将显存占用降低到原来的 1/N。ZeRO-1 分片优化器状态ZeRO-2 分片梯度ZeRO-3 分片全部参数。本文从 ZeRO 的设计动机出发分析三阶段的分片原理、通信模式和卸载策略。1. DeepSpeed ZeRO 的设计动机数据并行训练中每个 GPU 持有完整的模型参数、梯度和优化器状态副本。这些副本是冗余的——每个 GPU 上的参数值完全相同。ZeRO 的核心思想是消除冗余存储只在需要时收集完整数据。1.1 数据并行的冗余分析存储内容每个 GPU 存储实际需要冗余度模型参数完整14GB for 7B分片14GB/NN梯度完整14GB for 7B分片14GB/NN优化器状态完整28GB for 7B, Adam分片28GB/NN总计56GB56GB/NN1.2 ZeRO 的核心思想ZeRO 的核心思想是分阶段消除冗余DDP 冗余存储ZeRO-1: 分片优化器状态ZeRO-2: 分片梯度ZeRO-3: 分片参数显存节省: 4x (Adam)显存节省: 8x显存节省: 12x1.3 DeepSpeed ZeRO 的历史演进ZeRO 论文2019→ ZeRO-1/2 实现2020→ ZeRO-3 全分片2020→ ZeRO-Offload2021→ ZeRO-Infinity2022。1.4 DeepSpeed ZeRO 的产业应用模型规模ZeRO 阶段GPU 数BERT-Large340MZeRO-264GPT-3175BZeRO-310,000LLaMA 65B65BZeRO-32,048BLOOM 176B176BZeRO-33841.5 DeepSpeed ZeRO 的局限性ZeRO 的局限性包括通信量增加分片越多通信量越大、实现复杂度高需要手动管理分片以及小模型收益有限小模型下 ZeRO 的收益不如 DDP。2. ZeRO-1优化器状态分片2.1 ZeRO-1 的原理ZeRO-1 只分片优化器状态模型参数和梯度保持完整。优化器状态如 Adam 的动量和方差占显存最大通常是模型参数量的 2 倍分片后显存节省显著。2.2 ZeRO-1 的显存节省分片内容未分片7B, FP16分片后8 GPU节省模型参数14GB14GB0%梯度14GB14GB0%优化器状态28GB3.5GB87.5%总计56GB31.5GB43.75%2.3 ZeRO-1 的通信ZeRO-1 在优化器更新时需要通信每个 GPU 只更新自己的分片然后通过 All-Gather 收集完整更新后的参数。2.4 ZeRO-1 的实现importdeepspeed# ZeRO-1 配置zero_config{zero_optimization:{stage:1,# ZeRO-1reduce_bucket_size:5e8,allgather_bucket_size:5e8}}model_engine,optimizer,_,_deepspeed.initialize(modelmodel,optimizeroptimizer,config_paramszero_config)3. ZeRO-2梯度分片3.1 ZeRO-2 的原理ZeRO-2 在 ZeRO-1 的基础上进一步分片梯度。每个 GPU 只存储本分片参数的梯度不存储完整梯度。3.2 ZeRO-2 的显存节省分片内容未分片7B, FP16分片后8 GPU节省模型参数14GB14GB0%梯度14GB1.75GB87.5%优化器状态28GB3.5GB87.5%总计56GB19.25GB65.6%3.3 ZeRO-2 的通信ZeRO-2 在反向传播时使用 Reduce-Scatter 分发梯度在优化器更新后使用 All-Gather 收集参数。3.4 ZeRO-2 的实现# ZeRO-2 配置zero_config{zero_optimization:{stage:2,# ZeRO-2reduce_bucket_size:5e8,allgather_bucket_size:5e8,contiguous_gradients:True,overlap_comm:True# 通信重叠}}4. ZeRO-3全参数分片4.1 ZeRO-3 的原理ZeRO-3 在 ZeRO-2 的基础上进一步分片模型参数。每个 GPU 只存储本分片参数不存储完整参数。4.2 ZeRO-3 的显存节省分片内容未分片7B, FP16分片后8 GPU节省模型参数14GB1.75GB87.5%梯度14GB1.75GB87.5%优化器状态28GB3.5GB87.5%总计56GB7GB87.5%4.3 ZeRO-3 的通信ZeRO-3 在前向和反向传播时都需要 All-Gather 收集完整参数计算后丢弃非本分片参数。4.4 ZeRO-3 的实现# ZeRO-3 配置zero_config{zero_optimization:{stage:3,# ZeRO-3reduce_bucket_size:5e8,allgather_bucket_size:5e8,contiguous_gradients:True,overlap_comm:True,stage3_max_live_parameters:1e9,stage3_prefetch_bucket_size:5e8,stage3_param_persistence_threshold:1e6}}4.5 ZeRO 三阶段对比阶段参数分片梯度分片优化器分片显存节省通信量ZeRO-1否否是4x2 × ModelZeRO-2否是是8x2 × ModelZeRO-3是是是Nx3 × Model5. ZeRO-Offload 与卸载5.1 ZeRO-Offload 的原理ZeRO-Offload 将部分计算和存储卸载到 CPU 内存进一步减少 GPU 显存占用。5.2 卸载策略卸载内容卸载到显存节省速度影响优化器状态CPU减少 50% GPU 显存慢 10-20%参数CPU减少 33% GPU 显存慢 20-30%梯度CPU减少 33% GPU 显存慢 20-30%5.3 ZeRO-Offload 的实现# ZeRO-3 Offload 配置zero_config{zero_optimization:{stage:3,offload_optimizer:{device:cpu,# 优化器卸载到 CPUpin_memory:True},offload_param:{device:cpu,# 参数卸载到 CPUpin_memory:True}}}5.4 ZeRO-InfinityZeRO-Infinity 将卸载扩展到 NVMe 存储支持千亿参数模型的训练存储层级容量带宽延迟存储内容GPU 显存80GB2 TB/s纳秒当前活跃参数CPU 内存1TB100 GB/s微秒预取参数NVMe 存储10TB10 GB/s毫秒不活跃参数6. DeepSpeed ZeRO 的边界与失效模式6.1 通信瓶颈问题表现解决方案通信量大训练速度慢增加 GPU 数量通信延迟高同步等待时间长使用更高速网络通信不平衡某些 GPU 负载高优化通信拓扑6.2 卸载瓶颈问题表现解决方案CPU 带宽不足卸载等待时间长使用更高速 CPU 内存CPU 内存不足卸载失败增加 CPU 内存NVMe 带宽不足卸载速度慢使用 NVMe RAID6.3 DeepSpeed ZeRO 的优缺点总结优点缺点显存节省显著通信量增加支持超大模型实现复杂度高灵活的分阶段选择小模型收益有限支持卸载到 CPU/NVMe卸载速度慢7. DeepSpeed ZeRO 的工程实践7.1 ZeRO 阶段选择指南模型规模推荐阶段原因1BDDPZeRO-0显存足够通信少1B-10BZeRO-2梯度分片节省显存10B-100BZeRO-3全参数分片100BZeRO-3 Offload卸载到 CPU/NVMe7.2 性能优化优化策略描述效果通信重叠通信与计算重叠减少 20% 训练时间梯度累积模拟大 batch提高 GPU 利用率混合精度BF16 训练减少 50% 显存参数预取预取下一个模块的参数减少通信等待7.3 监控与调试指标描述告警阈值通信时间通信占总时间比例30%显存使用各 GPU 显存使用率90%卸载速度CPU/NVMe 卸载速度低于预期 50%8. ZeRO 的通信模式详解8.1 ZeRO-1 通信ZeRO-1 只在优化器更新时需要通信defzero1_communication(model,world_size,rank):ZeRO-1 通信模式# 前向传播无需通信lossmodel.forward(batch)# 反向传播All-Reduce 梯度与 DDP 相同model.backward()# 优化器更新只更新本分片shard_sizelen(model.parameters())//world_size param_shardlist(model.parameters())[rank*shard_size:(rank1)*shard_size]optimizer.step(param_shard)# 只更新本分片# 收集完整参数forparaminmodel.parameters():dist.all_gather(param,param)8.2 ZeRO-2 通信ZeRO-2 在反向传播时使用 Reduce-Scatter 分发梯度defzero2_communication(model,world_size,rank):ZeRO-2 通信模式# 前向传播无需通信lossmodel.forward(batch)# 反向传播Reduce-Scatter 梯度forparaminmodel.parameters():# 计算梯度后 Reduce-Scattershard_sizeparam.numel()//world_size chunksparam.grad.view(world_size,shard_size)reduce_scatter_outputtorch.zeros(shard_size,deviceparam.device)dist.reduce_scatter(reduce_scatter_output,chunks)param.gradreduce_scatter_output# 只保留本分片梯度# 优化器更新只更新本分片optimizer.step()# 收集完整参数forparaminmodel.parameters():shard_sizeparam.numel()//world_size shardparam.data[:shard_size]dist.all_gather(param.data.view(world_size,shard_size),shard)8.3 ZeRO-3 通信ZeRO-3 在前向和反向传播时都需要 All-Gatherdefzero3_communication(layer,input_data,world_size,rank):ZeRO-3 通信模式# 前向传播先收集完整参数shard_sizelayer.weight.numel()//world_size shardlayer.weight.data[:shard_size]full_weighttorch.zeros_like(layer.weight.data)dist.all_gather(full_weight.view(world_size,shard_size),shard)# 使用完整参数计算outputlayer.forward(input_data)# 丢弃非本分片参数layer.weight.datashardreturnoutput9. ZeRO 的卸载策略9.1 优化器卸载优化器卸载将 Adam 动量和方差从 GPU 卸载到 CPU 内存# ZeRO-Offload 优化器卸载配置zero_config{zero_optimization:{stage:3,offload_optimizer:{device:cpu,pin_memory:True,buffer_count:4,fast_init:False}}}卸载策略GPU 显存节省训练速度影响适用场景无卸载0%基准显存充足优化器卸载50%慢 10-20%显存不足优化器参数卸载66%慢 20-30%显存严重不足全卸载80%慢 30-50%超大模型9.2 CPU 优化器计算defcpu_adam_step(parameters,gradients,optimizer_state):CPU 上的 Adam 优化器步骤forparam,gradinzip(parameters,gradients):# 在 CPU 上更新参数param.dataparam.data-lr*grad/(torch.sqrt(optimizer_state[variance][param])1e-8)9.3 卸载的性能权衡GPU 显存GB可训练模型ZeRO-3可训练模型ZeRO-3 Offload16GB7B13B32GB13B30B80GB30B70B160GB70B175B10. ZeRO 的训练实践10.1 训练脚本importdeepspeeddeftrain_with_deepspeed(model,dataloader,config):使用 DeepSpeed ZeRO 训练# 初始化 DeepSpeedmodel_engine,optimizer,_,_deepspeed.initialize(modelmodel,model_parametersmodel.parameters(),config_paramsconfig)forepochinrange(10):forbatchindataloader:lossmodel_engine(batch)model_engine.backward(loss)model_engine.step()returnmodel_engine10.2 配置示例{train_batch_size:32,gradient_accumulation_steps:4,optimizer:{type:AdamW,params:{lr:1e-4,weight_decay:0.01}},zero_optimization:{stage:3,offload_optimizer:{device:cpu}},fp16:{enabled:true}}10.3 性能调优参数推荐值说明reduce_bucket_size5e8梯度通信 bucket 大小allgather_bucket_size5e8参数收集 bucket 大小stage3_prefetch_bucket_size5e8预取 bucket 大小stage3_max_live_parameters1e9最大存活参数数gradient_accumulation_steps4梯度积累步数11. DeepSpeed ZeRO 的进阶功能11.1 梯度裁剪# 启用梯度裁剪zero_config{zero_optimization:{stage:3,gradient_clipping:1.0# 梯度裁剪阈值}}11.2 学习率调度# 学习率调度配置zero_config{scheduler:{type:WarmupLR,params:{warmup_min_lr:0,warmup_max_lr:1e-4,warmup_num_steps:1000}}}11.3 混合精度训练# 混合精度配置zero_config{bf16:{enabled:True# 使用 BF16 替代 FP16},fp16:{enabled:False}}总结DeepSpeed ZeRO 通过分阶段消除数据并行中的冗余存储将显存占用降低到原来的 1/N。ZeRO-1 分片优化器状态节省 4x 显存ZeRO-2 分片梯度节省 8x 显存ZeRO-3 全参数分片节省 Nx 显存。ZeRO-Offload 将计算和存储卸载到 CPU/NVMe进一步减少 GPU 显存占用。ZeRO 阶段的选择取决于模型规模和硬件资源。外部引用ZeRO 原始论文https://arxiv.org/abs/1910.02054DeepSpeed 官方文档https://www.deepspeed.ai/ZeRO-Offload 卸载https://arxiv.org/abs/2101.06840ZeRO-Infinity 超大模型https://arxiv.org/abs/2204.12047DeepSpeed 混合精度https://www.deepspeed.ai/ZeRO 与 FSDP 对比https://www.deepspeed.ai/ZeRO-1 优化器分片https://arxiv.org/abs/1910.02054ZeRO-2 梯度分片https://arxiv.org/abs/1910.02054ZeRO-3 全参数分片https://arxiv.org/abs/1910.02054分布式训练显存优化https://arxiv.org/abs/2303.04226