DeepSpeed 通信优化:ZeRO 梯度 AllReduce、参数 All-Gather 与序列并行 AlltoAll 的三大优化实践

本文基于 DeepSpeed 官方技术博客 blogs/comm-opt/README.md 展开,系统讲解 DeepSpeed 面向大规模训练的三类通信集合操作优化:ZeRO 1/2 的梯度 AllReduce 优化(多 rank 桶合并与 MoE 专家布局 D+E)、ZeRO-2 参数 All-Gather 的 all_gather_into_tensor 单调用优化、以及 DeepSpeed-Ulysses 序列并行的 AlltoAll 融合。读完后你将理解每项优化解决的通信瓶颈、对应的配置开关(如 use_multi_rank_bucket_allreduce)、源码中的实现位置与验证方式,并能在自己的训练配置中正确启用这些优化。

1. 背景:通信集合操作为什么成为大规模训练的瓶颈

大模型训练的成本极高,除了合理组合硬件资源外,还需要一个可规模化的训练库保证在时间约束内完成训练。DeepSpeed 的可扩展性核心之一就是通信优化。集合通信操作(all-reduce、all-gather、all-to-all 等)是 ZeRO、MoE、AutoTP 等众多 DeepSpeed 关键技术的关键组成部分,这些优化在 DeepSpeed 版本 0.x.x 及以上可用。

具体来看三类优化各自针对的场景:

• 梯度 AllReduce 变慢:随着 GPU 数量增加,每卡分摊的参数量(参数分区大小 = #参数 / #GPU)越来越小,导致 AllReduce 的消息粒度变小、通信开销占比上升;
• 参数 All-Gather 变慢:ZeRO-2 在每步结束需要 all-gather 更新后的参数,分区越多窄化(narrow)操作越多;
• AlltoAll 开销:序列并行(DeepSpeed-Ulysses)中 q/k/v 的 all-to-all 分发随 SP 度增大而成为瓶颈。

2. ZeRO stage 1/2 的梯度 AllReduce 优化

2.1 问题:GPU 越多,消息越小

AllReduce 是训练过程中的重要环节。ZeRO 以桶(bucket)为单位处理梯度归约,可通过 reduce_bucket_size 配置以获得良好的通信吞吐。但当 GPU 数量增长时,会遇到"更小的分区 AllReduce"——此时现有的桶切分方案无法抵消通信开销。

官方博客给出的案例研究表明,这一问题主要在大规模 GPU 上训练中小规模模型(如 Llama-7B)时显著:

• 在 ZeRO stage 1/2 下训练 dense-7B 架构时,从 256 张 A100 扩展到 512、1024 张时,AllReduce 时间分别增加约 1 秒和 2 秒;
• MoE 架构问题更严重(增加 3-12 秒),因为专家并行 + 数据并行的默认布局下,需要跨 rank 通信时专家参数之间的距离更远。

根因在于:梯度平均发生在每卡 rank 的更小分区上(分区大小 = #参数 / #GPU),消息粒度不足以打满网络带宽。

2.2 优化一:同一进程组内的多 rank 桶合并(multi-rank bucketing)

该优化的做法是:将同一进程组中来自不同 rank、都需要归约的数据打包进一个大的一维展平张量,调用 AllReduce 而不是多次 reduce 操作;归约后再把各 rank 应得的那部分数据 scatter 回对应的 rank。

从源码结构看,这一特性由 ZeRO 配置项 use_multi_rank_bucket_allreduce 控制,定义在 zero_optimization 配置模型,默认值为 True,其文档注释写道:

Combine the reduce buckets of the different ranks and do an All-Reduce instead of multiple Reduce ops. This feature is useful when the model is small and we want to scale it on too many GPUs which therefore reduces the message sizes of each packet.

这与博客中"模型小、GPU 多、消息粒度小"的场景描述完全对应。

在 ZeRO 1/2 优化器实现 中可以看到具体的分派逻辑:遍历桶内每个参数分片,计算目标分区(partition_id)与偏移量后,按"归约方式 + 进程组"生成桶键(bucket key):

# 源码位置:deep speed/runtime/zero/stage_1_and_2.py(逻辑示意)
if self.use_multi_rank_bucket_allreduce or len(copy_ranks) > 1:
    bucket_key = ("allreduce", process_group)
else:
    bucket_key = ("reduce", dst, process_group)
...
if bucket_key[0] == "allreduce":
    self.allreduce_and_scatter(buckets[bucket_key], ...)   # 多 rank 桶合并:AllReduce 后再 scatter
else:
    self.allreduce_no_retain(buckets[bucket_key], rank=dst, ...)  # 传统路径:定向 reduce

即当开关打开时,多个 rank 的梯度分片被合并进同一个桶,统一走 allreduce_and_scatter(AllReduce + scatter);关闭时则退化为针对单一目标 rank 的 reduce(博客中称 reduce operation)。MoE 参数还使用独立的 expert_dp_process_group 进程组参与归约(stage_1_and_2.py),这与下文专家布局优化直接相关;此外源码对 MoE 有一条断言要求必须启用 reduce_scatter(stage_1_and_2.py),说明 MoE 场景下这条优化路径是实际被使用的代码路径。

2.3 优化二:MoE 专家-数据并行布局从 E+D 改为 D+E

MoE 架构的默认并行布局(E + D)是:专家先排在 E 个专家并行 GPU 上,再沿数据并行维度复制 D 份。这种布局下,数据并行的 rank 距离较远,尤其是存在跨 rank 通信时 AllReduce 更慢。

DeepSpeed 引入的新布局 D + E 是:先将每个专家沿数据并行维度复制 D 次,再沿专家并行维度拼接。官方给出的两种布局对比如下:

布局说明引自原文:左)E + D,先沿 EP 维度放置 GPU 再添加 DP;右)D + E,先按 DP 大小复制每个专家再构造 EP。第二种布局下 AllReduce 更快,代价是 AlltoAll 时间增加。由于 AllReduce 的通信量(全部参数量)通常远大于 AlltoAll(MLP 激活内存),端到端训练时间通常更快。

在 A100-DGX 集群(每节点 8 卡)上,从 E + D 切换到 D + E 后,参数更新过程的跨节点 InfiniBand 通信量减少约 8 倍——因为这部分通信改由节点内 NVLink 完成。虽然该布局增加了 MoE 部分的 AlltoAll 成本,但官方博客的实测表明 AllReduce 的收益足以覆盖这一开销。

官方实测结果(Table 1,1024 张 GPU)总结如下:

架构	GPUs	AllReduce 时间	单迭代时间
baseline (dense)	1024	1.2 s	5.4 s
optimized (dense)	1024	0.36 s	4.5 s
baseline (MoE)	1024	11.5 s	16.1 s
optimized (MoE)	1024	0.45 s	5.1 s

应用多 rank 桶合并后,dense 架构 AllReduce 时间降低 4 倍,MoE 架构降低 5-8 倍;MoE 架构再叠加 D+E 布局可额外获得约 3 倍节省。因此 GPU 数量越大、MoE 架构的性能收益越明显。例如训练 7B-base MoE 架构时,512 卡上单迭代时间从 13 秒降到 9.5 秒(37%),1024 卡上从 16.1 秒降到 5.1 秒(约 3.2 倍)。

2.4 配置示例

在训练配置文件(zero_optimization 段)中,与上述优化直接相关的开关如下,均可参考 zero_optimization 配置定义:

{
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "use_multi_rank_bucket_allreduce": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 500000000,
    "allgather_bucket_size": 500000000,
    "overlap_comm": true
  }
}

要点:

• use_multi_rank_bucket_allreduce 默认为 true,小模型上大集群训练时建议保持开启;
• reduce_bucket_size / allgather_bucket_size 控制每次集合通信的元素数,用于在通信吞吐与内存占用之间权衡,默认均为 5e8;
• 若训练 MoE 模型,reduce_scatter 必须为 true(源码中有断言约束,见上文)。

3. ZeRO-2 参数 All-Gather 优化:一次调用完成全参数聚合

与 AllReduce 同理,分区越多,all-gather 耗时越长。ZeRO stage-2 中参数存储在展平(flattened)缓冲区中,因此可以一次 all-gather 调用就把参数聚合进这个张量。

原始实现中,桶方案使用多个窄化(narrow)操作,从每个分区创建一系列大小为桶尺寸的张量列表——这是为了与 PyTorch 的 all_gather 操作对齐。而较新版本 PyTorch 提供了 all_gather_into_tensor 操作,支持直接对单个连续张量执行全量 all-gather,只需一次内核调用即可完成全参数聚合。官方实测该优化使大规模训练的步时间(step time)降低约 2 倍。

从源码看,这一行为由 partition_parameters.py 中的运行时探测决定:

self.use_all_gather_into_tensor = dist.has_all_gather_into_tensor()
if not self.use_all_gather_into_tensor:
    logger.info(f"all_gather_into_tensor API is not available in torch {torch.__version__}")

在后续的参数聚合路径中(如 partition_parameters.py),当该 API 可用时直接调用 dist.all_gather_into_tensor(flat_tensor, ...),否则回退到基于 all_gather 的多张量路径。这意味着该优化无需用户配置——只要底层 PyTorch 版本提供 has_all_gather_into_tensor 能力即自动生效;若日志中出现 "all_gather_into_tensor API is not available" 提示,说明当前 PyTorch 版本过旧,正在使用旧路径。

4. 序列并行训练的 AlltoAll 优化

该部分优化面向 DeepSpeed-Ulysses,目标是让序列并行度(Sp)从 2 扩展到 8 时保持可扩展性。研究基于 A100-DGX 硬件(每节点 8 卡),当并行度超过 8 时会因跨节点通信而性能受损,因此优化重点放在节点内 SP-2 到 SP-8 的扩展上。

融合在两个层面进行:

1. 融合 q、k、v 的序列 AlltoAll:不再预先拆分,而是使用混合(mixed)张量直接 Scatter 注意力头。实现上需要从模型侧获取额外信息(如 q 头数与 kv 头数),在调用 AlltoAll 之前完成头的切分。官方在 Megatron-DeepSpeed 仓库中加入了配套的序列并行改动来整合这些变化。
2. 融合 AlltoAll 张量并调用 PyTorch 的 AlltoAll-single API:对 scatter 维度 reshape 后,用单个张量执行 AlltoAll,避免了使用张量列表时为每个元素调用 contiguous 的额外开销。

官方报告这些优化带来了约 10%-15% 的提速,并在不同 SP 度与上下文长度下获得良好可扩展性。Table 2 展示了"GPU 数量翻倍 + SP 度提升"时的可扩展性(每 4K 样本/秒为样本吞吐指标):

GPUs	bsz	seq	Tokens (M)	SP	Sample (4K)-per-second	加速比 (x)
256	256	8192	2	1	60.71	1
512	256	8192	2	2	111.18	1.83
512	128	16384	2	4	108.81	1.79
512	64	32768	2	8	106.54	1.75
512	64	65536	4	8	110.05	1.81

从表中可以看到:使用 SP-2 从 256 卡扩展到 512 卡时获得超过 80% 的效率;在保持处理 token 数相近的前提下增加序列长度和 SP,2 倍资源仍保持 75% 以上效率;而最后一行(序列长度加到 65536、SP-4)token 吞吐量翻倍时,整体性能提升到 1.81 倍。

5. 小结:如何落地这些通信优化

优化	解决的瓶颈	启用方式	源码/文档位置
多 rank 桶合并 AllReduce	GPU 多、参数分区小导致 AllReduce 消息粒度不足	use_multi_rank_bucket_allreduce: true(默认开启)	config 定义、实现
MoE 专家布局 D+E	专家-数据并行布局下跨节点 AllReduce 慢	采用新的 EP/DP 布局组合(配合 MoE 训练)	官方博客 Fig 1,expert_dp_process_group 归约
all_gather_into_tensor 全参数聚合	ZeRO-2 每步参数 all-gather 的窄化开销	PyTorch 版本支持即自动启用	partition_parameters.py
AlltoAll 融合(Ulysses)	SP 度增大时 q/k/v AlltoAll 与张量列表开销	使用融合后的 Ulysses 通信路径	官方博客 Section 4、DeepSpeed-Ulysses 博客

整体来看,这些优化遵循同一思路:把"多次小消息集合通信"合并为"少量大消息集合通信",让通信量匹配硬件(节点内 NVLink / 跨节点 InfiniBand)的实际吞吐能力。对于 dense 模型,GPU 越多、模型越小,多 rank 桶合并的收益越明显;对于 MoE 模型,叠加 D+E 布局可获得最大的端到端收益;对于长序列训练,序列并行的 AlltoAll 融合保证了 SP 度扩展的效率。相关背景可进一步参阅 ZeRO 配置文档 与 zero_optimization 配置定义。