AllGather
AllGather 是 ReduceScatter 的对偶原语:每个节点贡献一个分片,操作完成后所有节点都持有完整的拼接结果。它是 Ring AllReduce 的后半段,也是 Sequence Parallelism 和 ZeRO 的核心通信操作。
语义定义
输入:节点 $i$ 持有 $A_i$,大小 $M/N$
输出:所有节点都持有完整拼接结果:
$$\begin{equation} \forall i \in [0, N): \quad B_i = [A_0, A_1, ..., A_{N-1}], \quad |B_i| = M \label{eq:ag-semantic-output} \end{equation}$$与 ReduceScatter 的对偶关系
| 维度 | AllGather | ReduceScatter |
|---|---|---|
| 输入大小/节点 | $M/N$ | $M$ |
| 输出大小/节点 | $M$ | $M/N$ |
| 数据流 | 分散 → 聚合 | 聚合 → 分散 |
| 操作 | 拼接 | 规约 + 分片 |
将 AllGather 的数据流反转并加入规约操作,即得 ReduceScatter。
与 AllReduce 的关系
$$\begin{equation} \text{AllReduce} = \text{ReduceScatter} + \text{AllGather} \label{eq:ag-allreduce-decomposition} \end{equation}$$Ring AllReduce 正是按此分解实现的。当只需要 AllGather(而非完整 AllReduce)时,通信量只有 AllReduce 的一半。
理论下界
延迟下界:每个节点需要从 $N-1$ 个其他节点接收数据。$k$ 步后最多 $2^k$ 个节点获得信息,因此至少需要 $\lceil \log_2 N \rceil$ 步:
$$\begin{equation} T_{\alpha}^{\min} = \lceil \log_2 N \rceil \cdot \alpha \label{eq:ag-delay-lb} \end{equation}$$带宽下界:AllGather 是 ReduceScatter 的对偶操作。节点 $i$ 初始只有 $M/N$ 字节,操作后需要持有 $M$ 字节。缺失的 $\frac{N-1}{N} \cdot M$ 字节必须从其他节点传入。每个节点的最小接收量为:
$$\begin{equation} R_i = \frac{N-1}{N} \cdot M \label{eq:allgather-recv} \end{equation}$$即使所有链路并行全速工作,完成时间也不低于:
$$\begin{equation} T_{\beta}^{\min}(\text{AllGather}) = \frac{(N-1) \cdot M}{N \cdot d \cdot \beta} \label{eq:allgather-bw-lower-bound} \end{equation}$$带宽下界的通用推导框架见 01-理论下界。
Ring AllGather
$N$ 个节点在 Ring 上传播数据,每步每个节点将上一步收到的 chunk 转发给下游。
步骤($N=4$)
初始:Node $i$ 持有 chunk $i$(大小 $M/4$)
| 步骤 | Node 0 | Node 1 | Node 2 | Node 3 |
|---|---|---|---|---|
| 初始 | $[A]$ | $[B]$ | $[C]$ | $[D]$ |
| Step 1 | $[A,D]$ | $[B,A]$ | $[C,B]$ | $[D,C]$ |
| Step 2 | $[A,D,C]$ | $[B,A,D]$ | $[C,B,A]$ | $[D,C,B]$ |
| Step 3 | $[A,D,C,B]$ | $[B,A,D,C]$ | $[C,B,A,D]$ | $[D,C,B,A]$ |
最终数据顺序不一定是 $[A,B,C,D]$,但包含所有 chunk。
代价推导
共 $N-1$ 步,每步每个节点将上一步收到的 chunk($M/N$ 字节)转发给右邻,同时从左邻接收一个新 chunk。每步代价为 $\alpha + \frac{M/N}{\beta}$,$N-1$ 步累加:
$$\begin{equation} T_{\text{Ring-AG}} = (N-1) \cdot \left(\alpha + \frac{M/N}{\beta}\right) = (N-1)\alpha + \frac{(N-1)M}{N\beta} \label{eq:ring-ag} \end{equation}$$- 延迟项:$(N-1)\alpha$ —— 非最优(下界 $\eqref{eq:ag-delay-lb}$ 为 $\lceil\log_2 N\rceil \cdot \alpha$)
- 带宽项:$\frac{(N-1)M}{N\beta}$ —— 达到下界 $\eqref{eq:allgather-bw-lower-bound}$,带宽最优
双向 Ring AllGather:同时向两个方向传播,步数减半:
$$\begin{equation} T_{\text{bidir-Ring-AG}} = \lceil\frac{N-1}{2}\rceil \cdot \left(\alpha + \frac{M/N}{\beta}\right) \label{eq:ag-bidir-ring} \end{equation}$$SCCL 的 AllGather 实现
SCCL(SOPHGO)在 4 芯 Ring 上将数据进一步切分为多个 bulk,通过 2 次迭代完成 AllGather:"将数据按 0 维拆分成多个 bulk,再通过流水并行实现效率最大化"。bulk 分块 + 流水线使每次迭代中链路利用率更高。
Recursive Doubling AllGather
适合 $N = 2^k$ 节点,每步数据量翻倍,能同时达到延迟和带宽下界。
步骤($N=8$)
| 步骤 | 通信对 | 交换量 | 操作后每节点数据量 |
|---|---|---|---|
| 1 | (0,1)(2,3)(4,5)(6,7) | $M/8$ | $M/4$ |
| 2 | (0,2)(1,3)(4,6)(5,7) | $M/4$ | $M/2$ |
| 3 | (0,4)(1,5)(2,6)(3,7) | $M/2$ | $M$ |
代价推导
共 $\log_2 N$ 步,第 $j$ 步($j=1,\ldots,\log_2 N$)每个节点与距离 $2^{j-1}$ 的对手互换当前持有的数据。第 $j$ 步交换量为 $M \cdot 2^{j-1} / N$(数据量逐步翻倍:$M/N, 2M/N, \ldots, M/2$)。
每步启动延迟 $\alpha$,带宽项为 $\frac{M \cdot 2^{j-1}}{N\beta}$。总代价:
$$\begin{equation} T_{\text{RD-AG}} = \sum_{j=1}^{\log_2 N} \left(\alpha + \frac{M \cdot 2^{j-1}}{N\beta}\right) = \log_2 N \cdot \alpha + \frac{M}{N\beta} \sum_{j=1}^{\log_2 N} 2^{j-1} \label{eq:ag-recursive-doubling-full} \end{equation}$$等比级数求和:$\sum_{j=1}^{p} 2^{j-1} = 2^p - 1 = N - 1$,因此:
$$\begin{equation} T_{\text{RD-AG}} = \log_2 N \cdot \alpha + \frac{(N-1)M}{N\beta} \label{eq:rd-ag} \end{equation}$$- 延迟项:$\log_2 N \cdot \alpha$ —— 达到延迟下界 $\eqref{eq:ag-delay-lb}$
- 带宽项:$\frac{(N-1)M}{N\beta}$ —— 达到带宽下界 $\eqref{eq:allgather-bw-lower-bound}$
- 同时最优! 但要求 $N = 2^k$ 且拓扑支持(Hypercube 理想)
算法对比
| 算法 | 延迟项 | 带宽项(墙钟时间) | 延迟最优? | 带宽最优? | 拓扑要求 |
|---|---|---|---|---|---|
| Ring | $(N-1)\alpha$ | $\frac{(N-1)M}{N\beta}$ | 否 | 是(1 链路 100%) | Ring |
| 双向 Ring | $\lceil\frac{N-1}{2}\rceil\alpha$ | $\lceil\frac{N-1}{2}\rceil \cdot \frac{M}{N\beta}$(约为单向的 $\frac{1}{2}$) | 否 | 是(2 链路各 100%) | 双向 Ring |
| Recursive Doubling | $\log_2 N \cdot \alpha$ | $\frac{(N-1)M}{N\beta}$ | 是 | 是(1 链路 100%) | Hypercube |
在 Ring 拓扑上,Recursive Doubling 的非直连通信对需要多跳路由,实际性能可能退化,使 Ring AllGather 反而更优。
在 LLM 中的应用
Sequence Parallelism(SP)
SP 将序列维度切分到多个节点,在 Transformer 层边界需要 AllGather + ReduceScatter:
- LayerNorm 前:AllGather 收集完整序列
- 执行 LayerNorm / Attention / MLP
- LayerNorm 后:ReduceScatter 重新分片
SP 与 TP 的总通信量相当(均为 $b \times s \times h$ 每层),但 AllGather/ReduceScatter 分开执行允许更好的计算-通信重叠。
ZeRO Stage 3
权重按节点分片存储,前向/反向计算时:
- AllGather:收集完整权重用于当前层计算(前向 1 次 + 反向 1 次)
- 计算完成后丢弃非本节点的权重分片
- ReduceScatter:反向传播中规约梯度并分片存储
推理场景(无反向传播):每层只需 1 次 AllGather,通信量为 ZeRO-1 的 $1/3$。
通信量对比
| 并行策略 | 使用的原语 | 每层通信量 |
|---|---|---|
| TP | AllReduce | $b \times s \times h$ |
| SP | AllGather + ReduceScatter | $b \times s \times h$(各一半) |
| ZeRO-3 | AllGather + ReduceScatter | layer_params |
参考资料
- Korthikanti et al., "Reducing Activation Recomputation in Large Transformer Models"(MLSys, 2023)- Sequence Parallelism
- Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"(SC, 2020)- ZeRO 中 AllGather 的使用
- Chan et al., "Collective Communication: Theory, Practice, and Experience"(CCPE, 2007)
- SCCL Documentation v1.0(SOPHGO)- AllGather 流水线 bulk 实现