跳到主要内容

ReduceScatter

与 AllGather 对偶的两类规约算法及在 LLM 并行中的应用

核心要点：

  • 输入：每节点持 $M$；输出：每节点持 $M / N$ 的完全规约结果
  • AllReduce = ReduceScatter + AllGather，独立用 RS 通信量减半
  • 带宽下界 $(N-1) M / (N d \beta)$，延迟下界 $\lceil \log_2 N \rceil \cdot \alpha$
  • Ring RS 带宽最优 / 延迟非最优 ($O(N) \alpha$)
  • Recursive Halving 同时最优，但要求 $N = 2^k$ + Hypercube
  • LLM 中 SP 前向 / ZeRO-3 反向 / TP 是核心用户

语义​

核心问题：ReduceScatter 的输入输出和数学语义是什么？与 AllGather 为什么是对偶关系？

节点 $i$ 持 $A_i = [A_i^{(0)}, \ldots, A_i^{(N-1)}]$ 大小 $M$, ReduceScatter 后只持第 $i$ 分片的完全规约结果:

$$\begin{equation} \forall i \in [0, N): \quad B_i = \bigoplus_{j=0}^{N-1} A_j^{(i)}, \quad |B_i| = M / N \label{eq:cc-rs-semantic} \end{equation}$$

AllGather ↔ ReduceScatter 对偶​

维度AllGatherReduceScatter
输入 / 节点$M / N$$M$
输出 / 节点$M$$M / N$
数据流分散 → 聚合聚合 → 分散
操作拼接规约 + 分片

@tbl-cc-rs-ag-duality AllGather ↔ ReduceScatter 对偶

反转 AllGather 数据流并加规约即得 ReduceScatter。

AllReduce 分解​

$$\begin{equation} \text{AllReduce} = \text{ReduceScatter} + \text{AllGather} \label{eq:cc-rs-ar-decomp} \end{equation}$$

独立用 ReduceScatter 通信量为 AllReduce 一半 — SP 与 ZeRO 用此性质换内存 / 通信。

理论下界​

核心问题：ReduceScatter 的下界与 AllGather 是否一致？

延迟下界 （与对偶版本同）：

$$\begin{equation} T_{\alpha}^{\min} = \lceil \log_2 N \rceil \cdot \alpha \label{eq:cc-rs-delay-lb} \end{equation}$$

带宽下界：节点 $i$ 本地保留自己分片 ($M / N$)，其余 $(N-1) M / N$ 必须发出：

$$\begin{equation} S_i = \frac{(N-1) M}{N} \label{eq:cc-rs-send-budget} \end{equation}$$ $$\begin{equation} T_{\beta}^{\min} = \frac{(N-1) M}{N \cdot d \beta} \label{eq:cc-rs-bw-lb} \end{equation}$$

带宽下界通用推导见 4.2 理论下界。

Ring ReduceScatter​

核心问题：Ring 上如何用 $N-1$ 步达带宽最优 （但延迟非最优）？

AllGather 的对偶，每步发一片给下游，同时接收上游一片并规约，见 。

Ring ReduceScatter 步骤流程 (N=4)@fig-cc-rs-ring

步骤 ($N = 4$)​

初始：节点 $i$ 持 $A_i = [A_i^{(0)}, A_i^{(1)}, A_i^{(2)}, A_i^{(3)}]$。

步骤操作
1节点 $i$ 发 chunk $(i+1) \% N$ 给右邻；接收方将收到的与本地对应 chunk 规约
2节点 $i$ 把上步规约后的 chunk 转发，继续规约
3最后一轮；每节点持一个 chunk 的完全规约结果

@tbl-cc-rs-ring-trace Ring RS 步骤 (N=4)

代价​

$N-1$ 步，每步每节点发 / 收 $M / N$，规约吞吐 $\gamma$:

$$\begin{equation} T_{\text{Ring-RS}} = (N-1) \cdot \left( \alpha + \frac{M / N}{\beta} + \frac{M / N}{\gamma} \right) \label{eq:cc-rs-ring-full} \end{equation}$$

$\gamma \gg \beta$ (SG2260 SUM 512 vs CDMA 64 GB/s) 时计算被通信掩盖：

$$\begin{equation} T_{\text{Ring-RS}} \approx (N-1) \alpha + \frac{(N-1) M}{N \beta} \label{eq:cc-rs-ring} \end{equation}$$

带宽项达下界 （带宽最优），延迟项 $O(N) \alpha$ 远超下界。

Recursive Halving ReduceScatter​

核心问题：能否同时达延迟与带宽下界？

Recursive Doubling AllGather 的对偶，每步交换量减半。

步骤 ($N = 8$)​

步骤通信对交换量操作后数据量 / 节点
1(0,4) (1,5) (2,6) (3,7)$M / 2$$M / 2$
2(0,2) (1,3) (4,6) (5,7)$M / 4$$M / 4$
3(0,1) (2,3) (4,5) (6,7)$M / 8$$M / N$

@tbl-cc-rs-rh-trace Recursive Halving RS 步骤 (N=8)

第 $j$ 步交换 $M / 2^j$ 字节。

代价​

$$\begin{equation} T_{\text{RH-RS}} = \sum_{j=1}^{\log_2 N} \left( \alpha + \frac{M}{2^j \beta} + \frac{M}{2^j \gamma} \right) \label{eq:cc-rs-rh-full} \end{equation}$$

$\sum 1 / 2^j = (N-1) / N$，忽略计算项：

$$\begin{equation} T_{\text{RH-RS}} \approx \log_2 N \cdot \alpha + \frac{(N-1) M}{N \beta} \label{eq:cc-rs-rh} \end{equation}$$

同时达延迟下界与带宽下界。

适用约束​

  • 要求 $N = 2^k$ (其他 $N$ 需 padding 或 hybrid 算法)
  • 要求拓扑近似 Hypercube （Fat-tree / 全连接可仿真；Ring 不行）

算法对比​

核心问题：Ring、Recursive Halving、多维分层三种 ReduceScatter 算法在延迟、带宽、适用场景上如何对比？

算法延迟项带宽项同时最优拓扑要求
Ring$(N-1) \alpha$$(N-1) M / (N \beta)$否 （延迟差）Ring
Recursive Halving$\log_2 N \cdot \alpha$$(N-1) M / (N \beta)$是Hypercube + $N = 2^k$

@tbl-cc-rs-algo-cmp Ring vs Recursive Halving

多维分层 ReduceScatter​

核心问题：实际集群带宽分层 （芯内 / 节内 / 机内 / 机间），如何减少慢链路通信？

分层框架通用推导见 4.8 AllReduce, ReduceScatter 是其前半段。

分解​

$N$ 节点组织为 $d$ 维网格 $p_0 \times p_1 \times \cdots \times p_{d-1}$ ($N = \prod p_i$)，逐维 RS:

$$\begin{equation} \text{Hier-RS}(N) = \text{RS}_0 \to \text{RS}_1 \to \cdots \to \text{RS}_{d-1} \label{eq:cc-rs-hier-decomp} \end{equation}$$

每级数据量缩 $p_i$ 倍。设 $P_i = \prod_{j=0}^{i} p_j$ ($P_{-1} = 1$):

级别输入 / 节点输出 / 节点
0$M$$M / p_0$
1$M / p_0$$M / (p_0 p_1)$
.........
$d - 1$$M / P_{d-2}$$M / N$

@tbl-cc-rs-hier-flow 多维分层 RS 各级数据量

代价​

$$\begin{equation} T_{\text{multi-dim-RS}} = \sum_{i=0}^{d-1} \left[ (p_i - 1) \alpha_i + \frac{p_i - 1}{p_i} \cdot \frac{M}{P_{i-1} \beta_i} \right] \label{eq:cc-rs-multidim} \end{equation}$$

外层 （低带宽） 只处理内层压缩后的数据，通信量减 $P_{i-1}$ 倍。

LLM 中的应用​

核心问题：ReduceScatter 在 LLM 训练的 DP/SP 等场景下有哪些典型使用、通信量和规模特征？

Sequence Parallelism (SP)​

Transformer 层结束后，将输出沿序列维度重新分片：

  • 输入：每节点持完整序列 ($b \times s \times h$)
  • 输出：每节点持 $1 / N$ 序列分片的规约结果 ($b \times (s / N) \times h$)

比 AllReduce 节省一半：不广播回所有节点。

ZeRO Stage 3 （梯度规约）​

反向传播中：

  1. 每节点持完整梯度 $\nabla W_i$ 大小 $M$
  2. ReduceScatter：每节点得 $1 / N$ 分片完全规约梯度
  3. 每节点更新 / 持久化自己负责的分片

梯度存储从 $M$ 降至 $M / N$，是 ZeRO-3 内存节省核心机制。

通信量对比​

并行策略原语每层通信量备注
TPAllReduce$b s h$RS + AG
SPRS + AG$b s h$各占一半，可重叠
ZeRO-3 反向ReduceScatterlayer_params每层梯度规约 + 分片

@tbl-cc-rs-llm-traffic LLM 中 ReduceScatter 通信量

Takeaway​

知识点核心结论
语义规约后按分片分发，每节点持 $M / N$
AllGather 对偶反转数据流 + 加规约
AllReduce 分解RS + AG，独立用 RS 通信量减半
带宽下界$(N-1) M / (N d \beta)$
延迟下界$\lceil \log_2 N \rceil \cdot \alpha$
Ring RS 性质带宽最优，延迟 $O(N) \alpha$ 非最优
Recursive Halving同时最优，需 $N = 2^k$ + Hypercube
LLM 用户SP 前向 / ZeRO-3 反向 / TP

@tbl-cc-rs-takeaway ReduceScatter 要点

参考资料​