跳到主要内容

计算通信 Overlap 分析

在大模型训练和推理中，通信操作可以与计算操作重叠执行，从而隐藏通信延迟。有效的 overlap 能显著降低端到端延迟，是大规模并行系统性能优化的核心手段之一。

基本原理​

为什么 Overlap 可行​

现代加速器（GPU/NPU）的硬件架构天然支持计算与通信并发：

  • 独立执行单元：计算核（Tensor Core / Cube Core）与 DMA 引擎（Copy Engine / GDMA / CDMA）是完全独立的硬件模块，互不抢占执行资源
  • 异步发起通信：软件可以通过异步 API 发出通信请求后立刻返回，DMA 引擎在后台执行数据搬运，计算核继续运行
  • 多 Stream 并行：CUDA Stream（或类似机制）允许不同 Stream 上的操作在不同硬件单元上同时推进，只有存在数据依赖时才需要显式同步

实现 overlap 的三个必要条件：

  1. 数据依赖解耦：通信所需的输入数据已就绪，且通信结果不是当前正在执行的计算的即时输入
  2. 硬件并发能力：芯片有独立的 DMA 引擎，能与计算核同时工作
  3. 软件异步调度：通信操作提前异步发出，而非阻塞等待完成

Overlap 的理论收益​

设某操作的计算时间为 $T_{\text{compute}}$，通信时间为 $T_{\text{comm}}$，overlap 率为 $r \in [0, 1]$：

$T_{\text{eff}} = T_{\text{compute}} + (1 - r) \cdot T_{\text{comm}}$

三种极端情况：

场景公式说明
无 overlap（$r = 0$）$T_{\text{eff}} = T_{\text{compute}} + T_{\text{comm}}$串行执行，通信完全暴露
完全 overlap（$r = 1$）$T_{\text{eff}} = T_{\text{compute}}$通信时间完全被隐藏
部分 overlap（$0 < r < 1$）$T_{\text{eff}} = T_{\text{compute}} + (1-r) T_{\text{comm}}$实际最常见的情况

完全 overlap 还有另一种等价表达：

$T_{\text{eff}} = \max(T_{\text{compute}}, T_{\text{comm}}) \quad \text{（仅当 overlap 窗口足够大时成立）}$

Overlap 的前提：通信时间 ≤ 计算时间​

只有当 $T_{\text{comm}} \leq T_{\text{compute}}$ 时，才能实现完全 overlap。若通信时间超过计算时间，即使硬件支持并发，通信延迟仍会暴露在关键路径上：

$T_{\text{eff}} = T_{\text{comm}} + \max(0,\, T_{\text{comm}} - T_{\text{compute}})$

因此，计算量越大、通信量越小，overlap 的收益越高。

张量并行（TP）中的 Overlap​

标准 TP AllReduce 难以 Overlap 的原因​

TP 的数据流（以 MLP 层为例）：

输入 X
  -> [列并行 Linear W1]   # 各设备独立计算列分片
  -> GeLU
  -> [行并行 Linear W2]   # 各设备计算部分积
  -> AllReduce             # 合并部分积，得到完整输出
  -> 输出（作为下一层输入）

AllReduce 处于当前层与下一层之间的关键路径上：

  • 当前层的计算必须等 AllReduce 完成才能产生有效输出
  • 下一层的计算必须等 AllReduce 完成后才能开始

这一数据依赖关系使得标准 TP AllReduce 无法与同层或下一层计算直接重叠。

序列并行（SP）打开的 Overlap 窗口​

Megatron-LM 提出的序列并行（SP）方案将 AllReduce 拆分为 All-Gather + Reduce-Scatter，与 Attention/MLP 计算形成可重叠的流水线：

[All-Gather 输入序列分片]
       |
       v（All-Gather 完成后开始 Attention 计算）
[Attention / MLP 计算]  ←→  [Reduce-Scatter 异步通信]
       |
       v（Reduce-Scatter 完成后，输出序列分片已就绪）
[LayerNorm / Dropout，在序列分片上操作]
       |
       v
[All-Gather 下一层输入]  ←→  [下一层 Attention/MLP 计算]

关键机制：将权重矩阵按列分块，每计算完一个块（chunk）后立即发起该块对应的 Reduce-Scatter，同时继续计算下一个块。这样通信和计算在时间轴上交叉推进。

TP Overlap 的实际效果​

TP AllReduce 的时间与矩阵乘时间的比值决定能否完全 overlap：

$\text{Overlap 可行性} = \frac{T_{\text{AllReduce}}}{T_{\text{MatMul}}}$

当该比值 < 1 时，可以完全 overlap。典型分析（BF16，$b=1$，$s=4096$，$h=7168$，TP=8）：

量典型值
每次 AllReduce 消息大小~58.7 MB
NVLink 带宽（节点内）400～900 GB/s
AllReduce 时间估算~0.1～0.3 ms
单层 MatMul 计算时间（FP8）~0.5～2 ms
比值0.05～0.6（通常可 overlap）

实际工程中，TP 度越大，单次 AllReduce 数据量不变但每设备分片更小，计算/通信比下降，overlap 收益减小。

流水并行（PP）中的 Overlap​

1F1B 调度下的 P2P 通信​

PP 的 stage 间 P2P 传输（激活值前向、梯度反向）必须等待前一个 stage 计算完成后才能发送，本质上是强数据依赖——P2P 通信本身难以与同 micro-batch 的计算重叠。

PP 的 overlap 机会主要出现在以下两个场景：

Pipeline Bubble 阶段

启动（startup）和排空（drain）阶段存在设备空闲时间（bubble）。Bubble 期间可以插入：

  • DP 梯度的 AllReduce（若与 DP 组合使用，梯度已在反向传播中积累完毕）
  • EP AllToAll 的待处理通信（若与 MoE 结合使用）
  • 下一个 global batch 的数据预取

Interleaved 1F1B（交错流水线）

将每个 PP stage 切分为 $v$ 个 chunk，Bubble ratio 从 $\frac{p-1}{m}$ 降至 $\frac{1}{v} \cdot \frac{p-1}{m}$。chunk 之间的 P2P 通信窗口更小但更频繁，调度更复杂，但增加了细粒度 overlap 的机会。

PP P2P 通信量与 Overlap 潜力​

每次 stage 边界传输的张量大小：

$M_{\text{PP}} = b \times s \times h \times \text{dtype\_size}$

与 TP AllReduce 消息量级相同（10～100 MB），但 P2P 通信不需要跨全组参与，仅涉及两个相邻 stage。相邻 stage 间保持低延迟连接（同板/同机架内）可以最小化 P2P 延迟对 bubble 宽度的影响。

数据并行（DP）中的 Overlap​

反向传播中的梯度 AllReduce Overlap​

DP 是所有并行策略中 overlap 潜力最高的。原因：梯度 AllReduce 不需要阻塞当前 step 的前向计算，且反向传播是逐层进行的，每一层的梯度计算完成后即可立即发起该层的 AllReduce，而无需等待所有层都完成。

Bucket 机制（PyTorch DDP 的实现方式）：

反向传播: Layer N → Layer N-1 → ... → Layer 1
                 ↓ 梯度就绪        ↓ 梯度就绪
AllReduce:  [bucket 填满后触发]  [bucket 填满后触发]

参数按通信顺序（反向传播顺序）分组到若干 bucket，每个 bucket 的所有参数梯度计算完成后立即触发 AllReduce，与后续层的反向计算并行。这使得 DP AllReduce 可以与反向传播充分重叠，是大模型训练中最成熟的 overlap 优化。

ZeRO 的 Overlap 特殊性​

ZeRO Stage 3 将参数、梯度、优化器状态均分片存储，每次前向/反向时动态重建：

  • 前向前 All-Gather（参数重建）可以与前一层的计算 overlap（Prefetch）
  • 反向后 Reduce-Scatter（梯度归约 + 分片）可以与后续层的反向计算 overlap

ZeRO 的 overlap 依赖精细的 Prefetch 调度，通常需要框架（如 DeepSpeed）内置支持。

专家并行（EP）中的 Overlap​

MoE 层的 AllToAll 通信（token dispatch + expert combine）分别在 Expert 计算的前后各发生一次：

AllToAll (dispatch)  ->  Expert 计算  ->  AllToAll (combine)

Dispatch AllToAll 与前一层的 MLP 输出计算有数据依赖，难以直接 overlap。

Expert 计算与 Combine AllToAll 之间存在 overlap 机会：可以在 Expert 计算完一批 token 后立即发起该批 token 的 combine 通信，与剩余 token 的 Expert 计算并行。这需要将 Expert batch 进一步分块（chunk），实现细粒度流水。

AllToAll 消息大小通常为 0.1～60 MB（远小于 DP AllReduce），通信时间短，overlap 的绝对收益有限，但对于 Expert 计算量较小的轻量 MoE 配置仍然重要。

Overlap 率建模​

统一公式​

$T_{\text{eff}} = T_{\text{compute}} + (1 - r_{\text{overlap}}) \cdot T_{\text{comm}}$

其中 $r_{\text{overlap}} \in [0, 1]$ 是 overlap 率，取决于：

  1. 计算/通信时间比值：$T_{\text{comm}} / T_{\text{compute}}$ 越小，$r_{\text{overlap}}$ 越接近 1
  2. 硬件 DMA 并发能力：DMA 引擎与计算核能否真正并行（芯片架构特性）
  3. 软件调度质量：能否及时发起异步通信、避免不必要的同步点

各并行策略的 Overlap 率参考​

并行策略通信原语Overlap 潜力典型 $r_{\text{overlap}}$限制因素
TP AllReduce（推理）AllReduce低0～0.3在关键路径上，数据依赖强
TP AllReduce（训练）AllReduce中0.3～0.7反向可部分 overlap
SP All-Gather + RSAll-Gather / Reduce-Scatter高0.7～0.9需要分块计算支持
PP P2PP2P Send/Recv中0.3～0.6bubble 窗口有限
DP AllReduceAllReduce高0.7～0.95bucket 大小设置
EP AllToAllAllToAll中0.3～0.6Expert batch 需分块

硬件能力对 Overlap 的影响​

关键硬件指标​

指标对 Overlap 的影响
DMA 引擎数量引擎越多，可同时处理的通信请求越多
DMA 引擎峰值带宽决定通信时间上限，影响 $T_{\text{comm}} / T_{\text{compute}}$ 比值
NVLink / C2C 带宽节点内高带宽通信可减小 $T_{\text{comm}}$，提升 overlap 可行性
RDMA 支持节点间通信通过 NIC DMA 执行，不占用计算核资源
多 Stream 支持软件并发发起多个通信请求，需硬件支持并发 DMA

SG2262 的 Overlap 特性​

SG2262 有三类 DMA 引擎（详见芯片配置）：

DMA 引擎峰值带宽作用域Overlap 场景
GDMA（Global DMA）68 GB/s片内（GMEM ↔ LMEM）与 Cube Core 计算 overlap，隐藏内存访问延迟
SDMA（Special DMA）64 GB/s片内辅助片内数据搬运与计算重叠
CDMA（C2C DMA）× 464 GB/s × 4片间（C2C 互联）与 Cube Core 计算 overlap，隐藏跨芯片通信延迟

CDMA 专用于跨芯片通信（对应 TP AllReduce、PP P2P 等场景），是实现集合通信与计算 overlap 的关键路径。4 个 CDMA 引擎可并发处理不同方向的通信流，总聚合带宽可达 256 GB/s。

compute_dma_overlap_rate: 0.8 表示 SG2262 在典型负载下，CDMA 通信与 Cube Core 计算的时间重叠率可达 80%，即通信时间仅有 20% 暴露在关键路径上。实际值需通过 benchmark 测量标定，可能因负载类型（计算密集 vs 访存密集）和通信消息大小而变化。

各并行策略 Overlap 能力对比汇总​

并行策略Overlap 机制工程实现难度推理收益训练收益
TP（标准）无（关键路径阻塞）—无低
TP + SPAll-Gather / RS 与 Attention/MLP 分块流水高中高
PP（1F1B）Bubble 期间插入其他通信中低中
PP（Interleaved）细粒度 chunk 间通信流水高中高
DPBucket 级梯度 AllReduce 与反向传播流水低（框架内置）无（推理无 DP）高
EPExpert batch 分块，combine 与计算流水高中高
ZeRO Stage 3Prefetch All-Gather + 分块 Reduce-Scatter高（需框架）无高

参考文献​