DeepSpeed ZeRO-2 深度解析：梯度分区、激活优化与内存碎片消除如何将训练规模提升一个数量级

推理引擎大模型【免费下载链接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.项目地址https://gitcode.com/gh_mirrors/fl/FlexGen点击查看免费下载本文以 DeepSpeed 官方博客《An Order-of-Magnitude Larger and Faster Training with ZeRO-2》为核心骨架结合本仓库内 DeepSpeed 子模块的完整源码deepspeed/runtime/zero/、deepspeed/runtime/activation_checkpointing/系统讲解 ZeRO-2 如何在原始 ZeRO优化器状态分区基础上进一步消解梯度、激活与碎片化内存开销从而支撑 170B 参数级别模型的高效训练。读者读完将掌握 ZeRO-2 的四大内存优化维度、zero_optimization与activation_checkpointing配置项的完整语义以及DeepSpeedZeroOptimizer的底层实现机制。一、ZeRO-2 是什么从优化器状态到全谱系内存优化原始 ZeROZero Redundancy Optimizer零冗余优化器的核心洞察是分布式数据并行训练中优化器状态如 Adam 的 momentum 与 variance在每张 GPU 上被冗余保存造成了大量不必要的显存占用。ZeRO 的第一步Stage 1通过将优化器状态按数据并行度切分、每张卡只持有并更新自己那一份分区消除了这份冗余。ZeRO-2 则将这一思想从优化器状态扩展到了训练过程中的完整内存消耗谱系。官方文档明确指出ZeRO-2 在原始 ZeRO 的优化器状态内存优化之外新增了三项核心技术梯度的内存消减——梯度不再每张卡冗余保存而是像优化器状态一样按数据并行度切分激活内存的优化——通过激活检查点Activation Checkpointing等手段压缩前向传播保存的中间激活内存碎片的消除——通过连续内存分配器预分配大块缓冲区并支持碎片整理避免训练中大量小张量的创建与释放把显存打碎。文档宣称的综合效果是将深度学习训练的规模与速度提升一个数量级order of magnitude允许训练高达 1700 亿170B参数的模型速度相比当时的最先进方案state of the art提升最多 10 倍。需要说明的是这里的数量级10 倍是官方文档给出的项目自述结论实际加速比会随模型、硬件与配置不同而变化。二、ZeRO-2 的四大内存优化维度2.1 优化器状态分区继承自 ZeRO Stage 1训练大模型时以 Adam 优化器 混合精度为例每个模型参数在显存中实际占据 16 份空间fp16 权重 2 份 fp32 权重 4 份 momentum 4 份 variance 4 份 梯度 2 份。其中 optimizer 状态fp32 权重 momentum variance占比最大且在数据并行下每张卡都存一份完整副本。ZeRO Stage 1 将这 16 份优化器相关状态按数据并行度 N 切分每张卡只持有 1/N。在DeepSpeedZeroOptimizer的初始化代码中可以看到对应实现参数组被展平flatten后按数据并行度切成近似相等的分区每个进程只从自己的分区克隆出 fp32 master 权重# deepspeed/runtime/zero/stage_1_and_2.py节选 # divide the flat weights into near equal partition equal to the data parallel degree # each process will compute on a different part of the partition data_parallel_partitions self.get_data_parallel_partitions( self.bit16_groups_flat[i], i) self.parallel_partitioned_bit16_groups.append(data_parallel_partitions) # A partition of the fp32 master weights that will be updated by this process. self.single_partition_of_fp32_groups.append( self.parallel_partitioned_bit16_groups[i][partition_id].to( self.device).clone().float().detach())对应配置stage: 1。这一步完成后本地优化器只包含属于自己分区的那部分参数因此本地优化器的状态momentum、variance 等也只与自己的分区相关源码注释原话the local optimizer only saves the states(momentum, variance, etc.) related to its partitions params。2.2 梯度分区ZeRO-2 的核心新增ZeRO-2 的 Stage 2 在 Stage 1 基础上把梯度也按数据并行度切分。反向传播时每张卡并不保存全部梯度而只保留自己分区对应的那一份梯度之后用它更新自己的优化器状态分区。在源码中Stage 1 与 Stage 2 的区分仅由一个布尔量控制# ZeRO stage 1 (False) or 2 (True) self.partition_gradients partition_grads以及配置枚举中的明确注释# deepspeed/runtime/zero/config.py class ZeroStageEnum(int, Enum): Enum class for possible zero stages disabled 0 optimizer_states 1 # 仅优化器状态分区 gradients 2 # 优化器状态 梯度分区 weights 3 # 优化器状态 梯度 权重分区梯度分区通过create_reduce_and_remove_grad_hooks()在参数上注册反向钩子实现每个梯度计算完成后立即按桶bucket进行 reduce-scatter 归约归约后该进程只持有自己分区那部分的平均梯度其余部分被释放。这避免了全量梯度在显存中的滞留。2.3 激活内存优化ZeRO-2 的另一大内存来源是前向传播保存的激活值——对长序列、大 batch 的 Transformer 模型激活值内存常常超过模型权重本身。ZeRO-2 通过激活检查点Activation Checkpointing技术只保留少量检查点激活用于反向重计算从而大幅降低激活内存峰值。在本仓库中激活检查点的配置独立于 ZeRO 配置位于deepspeed/runtime/activation_checkpointing/config.pyactivation_checkpointing: { partitioned_activations: [true|false], number_checkpoints: 100, contiguous_memory_optimization: [true|false], cpu_checkpointing: [true|false], profile: [true|false], synchronize_checkpoint_boundary: [true|false] }各配置项语义均来自DeepSpeedActivationCheckpointingConfig源码partition_activations默认False是否将激活值按数据并行度分区存储进一步降低单卡激活内存number_checkpoints默认None设置激活检查点的数量控制内存与重计算开销的权衡contiguous_memory_optimization默认False是否将激活存储到连续内存缓冲区减少碎片并提升内存利用率cpu_checkpointing默认False是否将激活检查点卸载到 CPU 内存profile默认False是否输出激活检查点相关内存画像synchronize_checkpoint_boundary默认False是否在检查点边界做同步。2.4 内存碎片消除深度学习训练中前向/反向会不断创建、销毁各种大小的临时张量激活、梯度、中间结果长此以往显存被打碎成大量不连续的小空闲块即使总空闲空间足够也无法分配大张量。ZeRO-2 引入了**连续内存分配器Contiguous Memory Allocator**应对这一问题。其实现位于deepspeed/runtime/zero/contiguous_memory_allocator.py的ContiguousMemoryAllocator类初始化时一次性预分配一块大缓冲区self.buffer torch.zeros(size, ...)后续张量分配都从这块缓冲区中切出当剩余空间不足以容纳请求的大小时分配器会自动执行碎片整理defragmentation把散落的小空闲块合并成连续大块并重新绑定参数张量的数据指针_reset_param_data。这一机制与配置中的contiguous_gradients相互配合梯度被复制进连续缓冲区避免反向传播期间的碎片化。三、ZeRO-2 配置实战zero_optimization完整参数详解DeepSpeedZeroConfigdeepspeed/runtime/zero/config.py对 ZeRO-2 相关的全部配置项做了精确定义。以下是在 JSON 配置中启用 ZeRO-2Stage 2的完整形态zero_optimization: { stage: 2, contiguous_gradients: true, reduce_scatter: true, reduce_bucket_size: 500000000, allgather_partitions: true, allgather_bucket_size: 500000000, overlap_comm: false, load_from_fp32_weights: true, offload_optimizer: { device: cpu, pin_memory: true }, round_robin_gradients: false, ignore_unused_parameters: true }各参数语义与源码依据如下参数默认值含义依据DeepSpeedZeroConfigdocstring 与实现stage0禁用0/1/2/3 分别对应禁用、优化器状态分区、优化器状态梯度分区、再加权重分区ZeRO-2 即stage: 2contiguous_gradientstrue梯度产生时即复制到连续缓冲区避免反向传播期间的内存碎片化reduce_scattertrue用 reduce-scatter 而非 allreduce 求平均梯度是梯度分区的基础通信原语reduce_bucket_size5e8每次 reduce/allreduce 的元素数限制大模型下通信缓冲所需内存allgather_partitionstrue每步结束时用 allgather 集合通信而非一串 broadcast收集所有 GPU 更新后的参数allgather_bucket_size5e8每次 allgather 的元素数限制 allgather 所需内存源码要求必须是 2 的倍数allgather_bucket_size % nccl_start_alignment_factor 0overlap_comm动态默认尝试将梯度归约与反向计算重叠Stage 2 下默认关闭源码中仅当 stage3 时动态置为 Trueload_from_fp32_weightstrue从 checkpoint 的 fp32 副本初始化 fp32 master 权重无精度损失而非从模型 fp16 副本转换offload_optimizer无将优化器状态卸载到 CPU/NVMe优化器计算也放到 CPU对 Stage 1/2/3 均有效round_robin_gradientsfalseStage 1/2 的 CPU 卸载优化按细粒度梯度分区在各 rank 间轮转梯度拷贝任务梯度累积步数越多、GPU 越多收益越大ignore_unused_parameterstrue检测到未使用参数时是否终止训练动态网络dynamic networks中未使用参数是常态故默认忽略⚠️ 注意两点限制ZeRO-2 是数据并行内的优化需要至少 2 张 GPU 才有分区收益单卡场景应选择 ZeRO-Offloadoffload_optimizer或 ZeRO-3。启用reduce_scatter时通信数据类型仅支持 fp16/bf16源码断言ZeRO-2 supports only float16 or bfloat16 communication_data_type with reduce scatter enabled且gradient_predivide_factor必须为 1.0。四、源码级剖析DeepSpeedZeroOptimizer的工作机制ZeRO-1/2 的统一实现类是DeepSpeedZeroOptimizerdeepspeed/runtime/zero/stage_1_and_2.py约 2491 行。其核心流程可以概括为四个阶段4.1 初始化展平、分区与 fp32 master 权重初始化时每个参数组先被移到 CPU 以释放显存空间创建 flat buffer随后展平调用flatten_dense_tensors_aligned将组内参数展平为连续张量并按 NCCL 4 字节对齐要求nccl_start_alignment_factor 2即 fp16 元素数的 2 倍对齐分区get_data_parallel_partitions将扁平权重切成与数据并行度相等的 N 份每张卡持有 1/Nfp32 master 权重每张卡从自己那份 fp16 分区克隆出 fp32 master 权重clone().float().detach()替换本地优化器的参数此后本地优化器只对自己分区做状态维护。同时源码用move_to_cpu与see_memory_usage反复记录内存占用为每步内存分配做诊断partition_count、partition_size等信息会在各 rank 打印。4.2 反向传播桶式梯度归约梯度分区通过反向钩子实现create_reduce_and_remove_grad_hooks在每个参数上注册钩子梯度一旦就绪就按reduce_bucket_size聚合到桶IPG bucket中桶满即发起 reduce-scatter归约完成后各卡只保留自己分区那份平均梯度其余立即释放_release_ipg_buffers。overlap_comm开启时归约在独立的reduction_stream上执行与反向计算重叠。4.3 优化器步骤每卡更新自己的分区由于本地优化器已被替换为只含本卡分区的 fp32 参数见param_group[params] [self.single_partition_of_fp32_groups[i]]优化器 step 时每张卡独立更新自己那份状态全程无需其他卡的优化器状态——这正是零冗余的本质。4.4 每步结束allgather 还原完整参数更新完成后通过allgather_partitionsallgather 集合或一串 broadcast 把各卡的分区权重汇聚回完整参数恢复模型权重供下一轮前向使用。allgather_bucket_size控制单次 allgather 的批大小避免为超大模型一次性申请过大通信缓冲。4.5 CPU 卸载扩展ZeRO-2 Offloadoffload_optimizer将 Stage 1/2 的优化器状态与计算卸载到 CPU此时contiguous_gradients被强制为 True源码self.contiguous_gradients contiguous_gradients or cpu_offload配合DeepSpeedCPUAdam在 CPU 上执行优化器更新并使用 pin-memory 缓冲temp_grad_buffer_for_cpu_offload torch.zeros(...).pin_memory()加速梯度传输。round_robin_gradients进一步将梯度按轮转方式分配给各 rank 拷贝避免单一 rank 成为 CPU 卸载的瓶颈。五、ZeRO-2 与 ZeRO-3 的边界何时用 Stage 2从ZeroStageEnum可以清楚看到 ZeRO 三个阶段的递进关系Stage 1只分区优化器状态Stage 2ZeRO-2分区优化器状态 梯度Stage 3ZeRO-3再进一步分区模型权重参数按需 allgather配合stage3_max_live_parameters、stage3_max_reuse_distance、stage3_prefetch_bucket_size等参数控制驻留参数数量。ZeRO-2 不需要像 Stage 3 那样在前向/反向中反复 gather 权重通信开销更低适合显存紧张但未到单卡装不下权重的规模当模型权重本身都无法放入单卡时则需要切换到 Stage 3甚至配合offload_param卸载权重到 CPU/NVMe即 ZeRO-Infinity 路线。此外overlap_comm的动态默认值也印证了这一设计源码 validator 中只有 Stage 3 才会默认开启通信重叠。六、验证与上手路径本仓库的 DeepSpeed 子模块位于 benchmark/third_party/DeepSpeedZeRO-2 核心实现见 stage_1_and_2.py配置解析见 config.py激活检查点配置与解析见 activation_checkpointing/config.py连续内存分配器实现见 contiguous_memory_allocator.pyDeepSpeed 官方在文档中建议参考 Megatron 教程中 GPT 系列模型的训练示例来上手 ZeRO-2对应本仓库 docs/_tutorials/megatron.md 教程更详细的技术原理可查阅官方技术报告原文档所引 arxiv 论文。配置示例可直接作为 DeepSpeed 训练脚本的ds_config.json使用再通过deepspeed.initialize(configds_config)或命令行--deepspeed_config传入即可启用 Stage 2 训练。对于 170B 级别的模型训练通常还需结合张量并行Megatron 式模型并行与流水线并行将 ZeRO-2 作为数据并行维度上的内存优化手段。七、小结ZeRO-2 的本质是把零冗余原则从优化器状态推广到梯度并叠加激活检查点与连续内存分配两项工程优化全面压制训练过程中的显存开销。它让每个数据并行 rank 只维护自己分区的优化器状态与梯度通过 reduce-scatter 归约 allgather 还原的通信模式以少量额外通信换取显存占用的大幅下降——这正是官方文档所称训练规模与速度提升一个数量级的技术根基。理解这四大优化维度及其在DeepSpeedZeroOptimizer中的落地方式是正确配置 ZeRO-2、评估其收益与通信开销的前提。赞分享推理引擎大模型【免费下载链接】FlexGenRunning large language models on a single GPU for throughput-oriented scenarios.项目地址https://gitcode.com/gh_mirrors/fl/FlexGen点击查看免费下载相关推荐DeepSpeed ZeRO-2 技术指南梯度分区、通信重叠与内存碎片治理如何让百亿级参数模型训练成为可能DeepSpeed ZeRO 2 技术指南梯度分区、通信重叠与内存碎片治理如何让百亿级参数模型训练成为可能 本文围绕 DeepSpeed 仓库中的 ZeRO人工智能大模型深度学习分布式训练预训练强化学习模型优化ftgo-application架构解析深入理解10个核心微服务的设计模式ftgo application架构解析深入理解10个核心微服务的设计模式 什么是ftgo application ftgo application是《Mi如何永久保存你的微信对话记忆WeChatMsg聊天记录提取与情感分析全指南如何永久保存你的微信对话记忆WeChatMsg聊天记录提取与情感分析全指南 你是否曾因手机存储空间不足而不得不删除珍贵的微信聊天记录或是担心手机丢失后那些承上一篇NGA论坛优化摸鱼体验插件让你的论坛浏览效率提升300%的终极指南下一篇TanStack Form React 的 UseFormGroup 类型与 useFormGroup Hook 完全解析构建类型安全的分步子表单创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考