跳到主要内容

ReduceScatter

ReduceScatter 对所有节点的数据执行规约,并将规约结果按分片分布到各节点——每个节点只获得全局规约结果的 $1/N$ 分片。它是 AllGather 的对偶原语,也是 Ring AllReduce 的前半段。

语义定义

输入:节点 $i$ 持有 $A_i = [A_i^{(0)}, A_i^{(1)}, ..., A_i^{(N-1)}]$,大小 $M$

输出:节点 $i$ 持有第 $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:rs-semantic-output} \end{equation}$$

与 AllGather 的对偶关系

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

将 AllGather 的数据流反转并加入规约操作,即得 ReduceScatter。

与 AllReduce 的关系

$$\begin{equation} \text{AllReduce} = \text{ReduceScatter} + \text{AllGather} \label{eq:rs-allreduce-decomposition} \end{equation}$$

独立执行 ReduceScatter 时,通信量只有 AllReduce 的一半。这在 SP 和 ZeRO 中至关重要。

理论下界

延迟下界:由对偶性,与 AllGather 相同:

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

带宽下界:节点 $i$ 持有 $M$ 字节数据,其中 $M/N$ 字节对应自己负责的分片(可本地保留),剩余 $\frac{N-1}{N} \cdot M$ 字节必须发送给负责对应分片的其他节点。因此每个节点的最小发送量为:

$$\begin{equation} S_i = \frac{N-1}{N} \cdot M \label{eq:reducescatter-send} \end{equation}$$

即使所有链路并行全速工作,完成时间也不低于:

$$\begin{equation} T_{\beta}^{\min}(\text{ReduceScatter}) = \frac{(N-1) \cdot M}{N \cdot d \cdot \beta} \label{eq:reducescatter-bw-lower-bound} \end{equation}$$

带宽下界的通用推导框架见 01-理论下界

Ring ReduceScatter

AllGather 的对偶。每步,每个节点发送一个 chunk 给下游,从上游接收并规约。

步骤($N=4$

初始:每个节点 $i$ 持有完整数据 $A_i = [A_i^{(0)}, A_i^{(1)}, A_i^{(2)}, A_i^{(3)}]$

步骤操作
Step 1Node $i$ 发送 chunk $(i+1)\%N$ 给 Node $(i+1)\%N$;接收方将收到的 chunk 与本地对应 chunk 规约
Step 2Node $i$ 将上步规约后的 chunk 发给下游;继续规约
Step 3最后一轮传递和规约;每节点得到一个 chunk 的完全规约结果

$N-1 = 3$ 步后:

  • Node 0 持有 chunk 1 的完全规约结果
  • Node 1 持有 chunk 2 的完全规约结果
  • ...

代价推导

$N-1$ 步,每步每个节点发送 1 个 chunk($M/N$ 字节)给右邻,从左邻接收 1 个 chunk 并与本地对应 chunk 执行规约。每步的代价为启动延迟 $\alpha$、传输时间 $\frac{M/N}{\beta}$、以及规约计算时间 $\frac{M/N}{\gamma}$$\gamma$ 为规约吞吐量),$N-1$ 步累加:

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

当规约吞吐远高于通信带宽($\gamma \gg \beta$,如 SG2260 的 SUM 单元 512 GB/s vs CDMA 64 GB/s)时,计算时间被通信时间覆盖,简化为:

$$\begin{equation} T_{\text{Ring-RS}} \approx (N-1)\alpha + \frac{(N-1)M}{N\beta} \label{eq:ring-rs} \end{equation}$$
  • 延迟项:$(N-1)\alpha$ —— 非最优(下界 $\eqref{eq:rs-delay-lb}$$\lceil\log_2 N\rceil \cdot \alpha$
  • 带宽项:$\frac{(N-1)M}{N\beta}$ —— 达到下界 $\eqref{eq:reducescatter-bw-lower-bound}$,带宽最优

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/8 = M/N$

每步交换量减半:第 $j$ 步交换 $M/2^j$

代价推导

$\log_2 N$ 步,第 $j$ 步($j=1,\ldots,\log_2 N$)交换 $M/2^j$ 字节。每步代价为 $\alpha + \frac{M/2^j}{\beta} + \frac{M/2^j}{\gamma}$,总代价:

$$\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) = \log_2 N \cdot \alpha + \frac{M}{\beta}\sum_{j=1}^{\log_2 N}\frac{1}{2^j} + \frac{M}{\gamma}\sum_{j=1}^{\log_2 N}\frac{1}{2^j} \label{eq:rs-recursive-halving-full} \end{equation}$$

等比级数求和:$\sum_{j=1}^{p}\frac{1}{2^j} = 1 - \frac{1}{2^p} = \frac{N-1}{N}$,忽略计算项后:

$$\begin{equation} T_{\text{RH-RS}} \approx \log_2 N \cdot \alpha + \frac{(N-1)M}{N\beta} \label{eq:rh-rs} \end{equation}$$
  • 延迟项:$\log_2 N \cdot \alpha$ —— 达到延迟下界 $\eqref{eq:rs-delay-lb}$
  • 带宽项:$\frac{(N-1)M}{N\beta}$ —— 达到带宽下界 $\eqref{eq:reducescatter-bw-lower-bound}$
  • 同时最优

算法对比

算法延迟项带宽项同时最优?拓扑要求
Ring ReduceScatter$(N-1)\alpha$$\frac{(N-1)M}{N\beta}$否(延迟差)Ring
Recursive Halving$\log_2 N \cdot \alpha$$\frac{(N-1)M}{N\beta}$Hypercube($N=2^k$

在 LLM 中的应用

Sequence Parallelism(SP)

SP 中 Transformer 层计算完成后,通过 ReduceScatter 将输出重新分片到各节点:

  • 输入:每节点持有完整序列的计算结果(Attention/MLP 输出),大小 $b \times s \times h$
  • 输出:每节点持有 $1/N$ 序列分片的规约结果,大小 $b \times (s/N) \times h$

与 AllReduce 相比,ReduceScatter 只完成 AllReduce 的一半工作——规约结果不需要广播到所有节点,节省了 AllGather 部分的通信量和时间。

ZeRO Stage 3(梯度规约)

反向传播中,每层的梯度需要 ReduceScatter:

  1. 每个节点持有完整梯度 $\nabla W_i$(大小 $M$,来自本节点的数据)
  2. ReduceScatter:规约所有节点的梯度,每节点得到 $1/N$ 分片的完全规约梯度
  3. 每节点更新并持久化自己负责的梯度分片

这使得梯度存储内存从 $M$ 降至 $M/N$,是 ZeRO Stage 3 内存节省的核心机制。

MoE Expert Parallelism

MoE 层中,Expert 的输出经过 Combine AllToAll 后,需要 ReduceScatter 将各 Expert 的结果按 token 分片收集:

  • 每个 Expert 节点持有对部分 token 的处理结果
  • ReduceScatter 将输出按 token 分布回各节点(每节点负责的 token 子集)

通信量对比

并行策略原语每层通信量说明
TPAllReduce$b \times s \times h$ReduceScatter + AllGather
SPReduceScatter + AllGather$b \times s \times h$各占一半,可重叠
ZeRO-3 反向ReduceScatterlayer_params每层梯度规约+分片

参考资料