DeepSpeed ZeRO 通信优化深度解析用 reduce-scatter 分区感知规约替代 all-reduce【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed本文基于仓库内发布于 2020 年 3 月的公告 docs/_posts/2020-03-17-reduce-scatter.md 展开。该公告预告了 DeepSpeed ZeRO 第一阶段引入分区感知partition-aware的梯度规约方式用 reduce-scatter 取代初版实现中一次性全局 all-reduce的通信模式宣称可将总通信量从数据并行基线的 1.5 倍降至 1.0 倍、通信耗时最多降低 2 倍。时至今日这一设计已成为仓库中 ZeRO Stage 1/2 的默认梯度平均实现。阅读本文后你将理解为什么梯度平均不需要 all-reduce、reduce-scatter 如何与 ZeRO 的参数分区天然契合以及当前仓库源码中该路径的实现与相关配置参数。一、这篇预览公告的背景与三大结论ZeRO stage 1 with reduced communication是 DeepSpeed 官方在早期发布的一篇预览sneak preview性质公告正文以三个要点概括了 ZeRO 训练通信优化的核心思路分区感知partition-aware方案取代了初版实现采用的全局集合通信all-reduce总通信量由数据并行data parallelism的1.5 倍降至 1.0 倍相比 all-reduce通信时间最多可降低 2 倍。需要说明的是上述 1.5x、2x 等量化结论是官方发布该公告时的口径与测量结果具体数值会随模型规模、GPU 数量、梯度桶bucket切分策略等条件而变化本文后续将聚焦为什么能省通信的机制这部分在当前仓库源码中有完整、可直接核验的实现证据。二、机制背景为什么梯度平均用 all-reduce 是浪费的在标准数据并行训练中每个 rank 在各自数据分片上反向传播后需要把各副本上相同参数的梯度取平均再用于更新参数。最直接的做法是对整份梯度做一次all-reduceRing All-Reduce。从通信原语上看一次 all-reduce 在功能上等价于reduce-scatter all-gather先把各 rank 的数据规约并打散到所有 rank每 rank 得到全量结果的一个分片再通过 all-gather 把完整结果广播到每个 rank。ZeRO 的省显存思路是分区优化器状态Stage 1、乃至 16-bit 梯度Stage 2都被切分成数据并行度等份每个 rank 只负责持有并更新自己那份分区。因此梯度平均的真正产物并不需要是全量平均梯度——每个 rank 只需要与自己持有的参数分区相对应的那一段平均梯度切片即可完成本地优化器更新。于是问题就显现了all-reduce 多做了最后一步 all-gather把每个 rank 都不需要的其它分区数据也全员广播了一遍。这正是公告所说的初始实现使用全局集体通信all-reduce存在的通信冗余而 reduce-scatter 只做规约 打散恰好把平均梯度的最终落点对准各自的分区从而省去 all-gather 阶段。三、当前仓库中的实现佐证分区感知规约已写入 ZeRO 优化器该公告预告的优化并非停留在概念上。在当前仓库中ZeRO Stage 1 与 Stage 2 由同一个优化器类统一实现 deepspeed/runtime/zero/stage_1_and_2.py。从源码可以清晰看到分区感知规约的落地1.reduce_scatter是默认开启的开关构造参数中reduce_scatterTrue为默认值stage_1_and_2.py#L164并在初始化时被保存为self.reduce_scatterstage_1_and_2.py#L230。同一文件中用partition_gradients区分两个阶段True时为 ZeRO-2梯度也被分区False时为 ZeRO-1仅优化器状态分区见 stage_1_and_2.py#L224-L226。可见分区感知、只保留本地分区所需梯度切片的逻辑对两个阶段是统一的。2.average_tensor一个方法内的两条规约路径梯度规约的核心函数是average_tensorstage_1_and_2.py#L1360-L1478关闭 reduce_scatter 时走gradient_reduction_w_predividestage_1_and_2.py#L1269-L1298内部对整桶梯度调用dist.all_reduce并通过gradient_predivide_factor、postscale_gradients、gradient_average等参数在 fp16 下做先除后归约以控制数值稳定性开启 reduce_scatter 时代码遍历桶内每个参数的分区元数据param_to_partition_ids与grad_start_offset把梯度张量按目标分区(dst_rank, bucket_offset, numel)切成若干连续切片stage_1_and_2.py#L1386-L1442随后按切片目标分组切片只属于单一目标 rank 时使用(reduce, dst, process_group)为键、对该目标做dist.reduce经allreduce_no_retain→allreduce_bucket(rankdst)其内部即 stage_1_and_2.py#L1866-L1871 的dist.reduce到目标全局 rank存在多 rank 副本需求等特殊场景时才回退到allreduce_and_scatter路径stage_1_and_2.py#L1466-L1470。换言之一个参数可能横跨多个分区、一个规约桶里也往往有多个连续片段指向同一目标 rank代码会将这些片段合并后再一次性规约stage_1_and_2.py#L1435-L1442尽量让每次集合通信都有足够大的数据量。3. 归约结果直接对位本地分区非本地梯度尽早释放在reduce_ipg_gradsstage_1_and_2.py#L1701-L1765中规约完成后若partition_gradients为真代码会对不属于当前 rank 分区的参数直接清空其梯度clear_grad_attribute从而在不持有完整平均梯度的前提下省下显存对属于本地分区的参数调用copy_grads_in_partition把切片写入本地连续分区缓冲供后续更新该分区的优化器状态使用。ZeRO-2只保留自己那份梯度的显存收益正是以这套分区感知的规约落位为前提的。文档 docs/_tutorials/zero.md 对 ZeRO 各阶段的划分给出了同一口径的描述Stage 2 即规约后的 16-bit 梯度也被分区每个进程只保留与其优化器状态分区对应的那部分梯度。四、与反向传播重叠IPG 桶与overlap_commreduce-scatter 不止省通信量还便于与反向传播重叠。ZeRO Stage 1/2 实现了独立分区梯度桶IPGIndependent Partition Gradient机制每个参数的反向梯度一经产生就立即按分区切片拷贝进reduce_bucket_size大小的连续桶stage_1_and_2.py#L1202-L1263当桶满或本轮 backward 结束时触发average_tensor执行上述分区感知规约。开启overlap_comm后规约在独立的 reduction 流stream上进行stage_1_and_2.py#L1360-L1375从而把梯度通信藏在后续反向计算背后进一步摊薄通信开销。五、配置参数与启用方式在 DeepSpeed 配置中只需为zero_optimization启用 ZeRO 即可使用该机制reduce_scatter相关项均位于zero_optimization键下。以下为参考配置融合了文档 docs/_tutorials/zero.md 中 Stage 1 与 Stage 2 的示例写法{ train_batch_size: 32, gradient_accumulation_steps: 1, zero_optimization: { stage: 1, reduce_bucket_size: 5e8, contiguous_gradients: true, reduce_scatter: true, overlap_comm: true } }其中各字段的含义与默认值见 docs/_pages/config-json.md#L514-L536 与 deepspeed/runtime/zero/config.py#L108-L113参数说明默认值reduce_scatter是否用 reduce/reduce-scatter 代替 all-reduce 平均梯度即分区感知规约的开关对应本公告主题truereduce_bucket_size单次规约处理的元素数上限限制一次集合通信占用的内存桶大小5e8约 5 亿元素contiguous_gradients反向过程中把梯度拷贝进连续缓冲避免内存碎片化trueoverlap_comm是否尝试将梯度规约与反向计算重叠falsestage取1即 ZeRO Stage 1仅分区优化器状态取2时配合reduce_scatter: true同时分区梯度无必填要点提示字段描述与文档口径一致Uses reduce or reduce scatter instead of allreduce to average gradients默认即开启docs/_pages/config-json.md#L520-L524。因此在当前版本的 DeepSpeed 中本文讨论的通信优化是默认生效的梯度平均路径reduce_bucket_size是影响通信效率与内存的关键调优项桶越大单次集合通信的效率越高但临时缓冲占用也越大它对上文分片后合并、尽量一次规约的合并效率有直接约束overlap_comm需要与 IPG 桶、连续梯度等机制配合才能发挥效果并非所有后端/加速器都支持流级重叠。六、正确性保护与扩展场景分区感知规约在实现上还包含若干正确性与扩展性设计同样可在源码中核验均值语义开启reduce_scatter时先按dp_world_size / sequence_parallel_size对梯度做除法再发往目标分区stage_1_and_2.py#L1445-L1446保证结果等价于全体数据并行副本的平均fp16 下的数值安全关闭 reduce_scatter 的 all-reduce 路径保留了gradient_predivide_factor与postscale_gradients的前置/后置缩放这是 fp16 大世界规模梯度规约的经典数值保护手段MoE / 专家并行当桶内存在 MoE 参数时会切换到专家数据并行进程组执行规约stage_1_and_2.py#L1399-L1401避免在错误的通信域内混算序列并行规约前的除数使用dp_world_size / sequence_parallel_size说明该通信路径与序列并行维度按设计协作且代码在序列并行大于 1 时会把通信 dtype 提升为 fp32stage_1_and_2.py#L1857-L1858。七、总结从 2020 年 3 月这篇分区感知、降低通信的预览公告到当前仓库中 deepspeed/runtime/zero/stage_1_and_2.py 内以reduce_scatter为默认开关的完整实现可以清晰看到一条一以贯之的设计主线ZeRO 既分区存储也分区通信。既然每个 rank 最终只需要属于自己的那份平均梯度就用 reduce-scatter 精确投递、免去 all-reduce 中多余的 all-gather 阶段再配合reduce_bucket_size约束桶内存、contiguous_gradients对抗碎片化、overlap_comm实现通信与计算重叠最终让梯度平均的通信量与数据并行基线持平并把通信耗时压低到官方公告所称的一半以内。对于希望深入 ZeRO 训练通信原理的读者建议从average_tensor与reduce_ipg_grads这两个函数入手研读源码并结合 docs/_tutorials/zero.md 与 docs/_pages/config-json.md 的配置说明做小规模实验验证。【免费下载链接】DeepSpeedDeepSpeed is a deep learning optimization library that makes distributed training and inference easy, efficient, and effective.项目地址: https://gitcode.com/GitHub_Trending/de/DeepSpeed创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考