AllToAll 全交换通信
AllToAll 是最通用的集合通信原语:每个节点向每个其他节点发送不同的数据,同时接收来自每个其他节点的不同数据。它是 MoE 专家路由(Expert Routing)的核心通信操作。
语义定义
输入:节点 $i$ 持有 $N$ 块数据 $[D_i^{(0)}, D_i^{(1)}, ..., D_i^{(N-1)}]$,其中 $D_i^{(j)}$ 将发送给节点 $j$
输出:节点 $j$ 持有 $[D_0^{(j)}, D_1^{(j)}, ..., D_{N-1}^{(j)}]$
$$\begin{equation} \forall j \in [0, N): \quad B_j = [D_0^{(j)}, D_1^{(j)}, ..., D_{N-1}^{(j)}] \label{eq:a2a-semantic-def} \end{equation}$$直观理解:输入是 $N \times N$ 数据矩阵,AllToAll 执行矩阵转置——行分布变为列分布。
与其他原语的区别
| 原语 | 每节点发送 | 每节点接收 |
|---|---|---|
| AllGather | 相同数据给所有人 | 所有人的不同数据 |
| AllReduce | 不同数据(被规约) | 相同结果 |
| AllToAll | 不同数据给不同人 | 不同人的不同数据 |
AllToAll 是唯一每对节点之间交换独特数据的原语,是最通用也是通信量最大的原语。
理论下界
延迟下界:AllToAll 中每个节点需要和其他 $N-1$ 个节点各交换一次独特数据。在单端口模型下(每步只能与一个节点通信),至少需要 $N-1$ 步:
$$\begin{equation} T_{\alpha}^{\min} = (N-1) \cdot \alpha \label{eq:a2a-delay-lb} \end{equation}$$注意与 AllReduce 的区别:AllReduce 的延迟下界是 $\lceil\log_2 N\rceil \cdot \alpha$(信息可以通过中间节点级联传播),而 AllToAll 每对节点之间传输的数据是独特的,无法通过中间节点合并,因此下界更高。
带宽下界:设 $m$ 为每对节点间的数据块大小(总消息大小 $M = N \cdot m$)。节点 $i$ 需要向其他 $N-1$ 个节点各发送 $m$ 字节,总发送量 $(N-1)m = \frac{(N-1)M}{N}$:
$$\begin{equation} T_{\beta}^{\min} = \frac{(N-1)M}{N \cdot d\beta} \label{eq:a2a-bw-lb} \end{equation}$$Bisection 下界:将 $N$ 个节点对半分为 $A$、$B$ 两组各 $N/2$ 个。$A$ 中每个节点需要向 $B$ 中 $N/2$ 个节点各发送 $m$ 字节,从 $A$ 侧跨 bisection 的总数据量为 $(N/2) \cdot (N/2) \cdot m = N^2 m / 4 = NM/4$:
$$\begin{equation} T_{\text{bisection}} = \frac{NM}{4 \cdot B_{\text{bisection}}} \label{eq:a2a-bisection-lb} \end{equation}$$对 Ring 拓扑($B_{\text{bisection}} = 2\beta$):$T_{\text{bisection}} = NM / (8\beta)$。当 $N$ 较大时,bisection 下界可能比节点级带宽下界 $\eqref{eq:a2a-bw-lb}$ 更紧。
带宽下界推导
带宽下界的通用推导框架见 01-理论下界。
语义:每个节点持有 $M$ 字节数据,需要将其中 $M/N$ 字节发给其他每个节点(总共向 $N-1$ 个节点各发 $M/N$)。
推导:节点 $i$ 需要向其他 $N-1$ 个节点各发送 $M/N$ 字节,总发送量:
$$\begin{equation} S_i = (N-1) \cdot \frac{M}{N} = \frac{(N-1)}{N} \cdot M \label{eq:alltoall-send} \end{equation}$$同理,节点 $i$ 从其他 $N-1$ 个节点各接收 $M/N$ 字节,总接收量:
$$\begin{equation} R_i = \frac{(N-1)}{N} \cdot M \label{eq:alltoall-recv} \end{equation}$$ $$\begin{equation} T_{\beta}^{\min}(\text{AllToAll}) = \frac{(N-1) \cdot M}{N \cdot d \cdot \beta} \label{eq:alltoall-bw-lower-bound} \end{equation}$$MoE Dispatch/Combine:非均匀 AllToAll
上述推导假设均匀通信(每对节点间 $M/N$)。MoE 场景中,AllToAll 承担 Expert 路由,分为两次:
- Dispatch:每个节点将本地 token 按路由结果发送到持有对应 Expert 的节点
- Combine:Expert 计算完成后,结果返回原始节点
MoE 路由导致通信量不均匀——热门 Expert 所在节点接收大量 token,冷门 Expert 几乎空闲。设节点 $j$ 接收的 token 总量为 $R_j$,则:
$$\begin{equation} T_{\beta}^{\min}(\text{Dispatch}) \ge \max_j \frac{R_j}{d \cdot \beta} \label{eq:dispatch-bw-lower-bound} \end{equation}$$带宽下界由负载最重的节点决定,而非平均值。均匀路由下 $R_j = \frac{(N-1)}{N} M$ 退化为标准 AllToAll 下界。
不均匀程度的量化:设 Expert Capacity Factor 为 $C$($C = 1.0$ 无冗余,$C = 1.25$ 有 25% 余量),每个 Expert 最多接收:
$$\begin{equation} \text{capacity}_j = C \cdot \frac{B \cdot K}{E} \label{eq:expert-capacity-bound} \end{equation}$$其中 $B$ 为 batch token 数,$K$ 为 Top-K,$E$ 为 Expert 总数。当路由严重不均匀时,最大负载节点的 $R_j$ 趋近 $C \cdot \frac{B \cdot K}{E} \cdot h \cdot \text{dtype}$,可能远大于均匀情况下的 $R_j$。
Dispatch vs Combine 的对称性:Dispatch 的发送量分布取决于 token 路由分布(不均匀),Combine 的发送量分布取决于 Expert 输出分布(同样不均匀且与 Dispatch 模式相同)。因此两次 AllToAll 面临相同的不均匀性约束,下界相同。
算法一:Pairwise Exchange
最直观的算法:$N-1$ 轮,每轮每个节点与一个不同的伙伴交换数据。
通信调度
使用循环配对(round-robin pairing),第 $r$ 轮节点 $i$ 与 $(i+r) \bmod N$ 交换数据:
| 轮次($N=4$) | Node 0 伙伴 | Node 1 伙伴 | Node 2 伙伴 | Node 3 伙伴 |
|---|---|---|---|---|
| 1 | 1 | 0 | 3 | 2 |
| 2 | 2 | 3 | 0 | 1 |
| 3 | 3 | 2 | 1 | 0 |
代价推导
$N-1$ 轮,每轮每个节点与一个伙伴互相发送 $m = M/N$ 字节。每轮代价为 $\alpha + \frac{M/N}{\beta}$,$N-1$ 轮累加:
$$\begin{equation} T_{\text{pairwise}} = (N-1) \cdot \left(\alpha + \frac{M/N}{\beta}\right) = (N-1)\alpha + \frac{(N-1)M}{N\beta} \label{eq:pairwise-a2a} \end{equation}$$- 延迟项:$(N-1)\alpha$ —— 达到延迟下界 $\eqref{eq:a2a-delay-lb}$,单端口模型下最优
- 带宽项:$\frac{(N-1)M}{N\beta}$ —— 达到带宽下界 $\eqref{eq:a2a-bw-lb}$
在完全图上,每轮所有通信对之间无链路共享,无冲突。在 Ring 上远距离配对需多跳,可能产生带宽竞争。
算法二:Ring AllToAll
利用 Ring 拓扑特性:$N-1$ 步,每步所有节点同时向右邻发送、从左邻接收,等价于在 Ring 上执行 $N$ 个并行的 Scatter。
步骤($N=4$)
每个节点 $D_i = [D_i^{(0)}, D_i^{(1)}, D_i^{(2)}, D_i^{(3)}]$
| 步骤 | Node 0 发送 | Node 0 接收 |
|---|---|---|
| Step 1 | $D_0^{(1)}$ → Node 1 | $D_3^{(0)}$ from Node 3 |
| Step 2 | $D_0^{(2)}$ → Node 1(转发) | $D_2^{(0)}$(Node 3 转发) |
| Step 3 | $D_0^{(3)}$ → Node 1(转发) | $D_1^{(0)}$(两跳转发) |
代价推导
与 Pairwise 相同:$N-1$ 步,每步传输 $M/N$ 字节:
$$\begin{equation} T_{\text{Ring-A2A}} = (N-1)\alpha + \frac{(N-1)M}{N\beta} \label{eq:ring-a2a} \end{equation}$$代价公式与 Pairwise $\eqref{eq:pairwise-a2a}$ 相同,但在 Ring 上的优势是每步只使用相邻链路,无链路冲突。
SCCL 实现:AllToAll 基于 Scatter 实现——每个节点作为 root 执行一次 Scatter,$N$ 个 Scatter 并行执行。
算法三:Bruck 算法
由 Bruck 等人提出(TPDS 1997),步数为 $O(\log N)$,适合小消息延迟优先场景。
算法思想
利用 Recursive Doubling:$\lceil \log_2 N \rceil$ 步,第 $j$ 步($j = 0, 1, ..., \lceil\log_2 N\rceil - 1$):
- 节点 $i$ 向节点 $(i + 2^j) \bmod N$ 发送所有尚未到达目标的 chunk
- 从节点 $(i - 2^j) \bmod N$ 接收数据
- 对接收到的数据进行本地重排
代价推导
共 $\lceil\log_2 N\rceil$ 步。第 $j$ 步中,每个节点将约 $\lceil N/2 \rceil$ 个目标节点的数据打包发送给距离 $2^j$ 的节点。每步的平均传输量约为 $(N-1)m/2$(因为每步发送约一半目标的数据),加上启动延迟 $\alpha$。操作完成后还需要 $(N-1)$ 次本地数据重排(内存拷贝,非网络开销):
$$\begin{equation} T_{\text{Bruck}} = \lceil \log_2 N \rceil \cdot \alpha + \lceil\log_2 N\rceil \cdot \frac{(N-1)m}{2\beta} + (N-1) \cdot t_{\text{copy}} \label{eq:bruck-a2a} \end{equation}$$- 延迟项:$\lceil\log_2 N\rceil \cdot \alpha$ —— 延迟最优($O(\log N)$ 步,远优于 Pairwise 的 $O(N)$)
- 带宽项:$\lceil\log_2 N\rceil \cdot \frac{(N-1)m}{2\beta}$ —— 非最优(每步传输冗余数据,总带宽开销是 Pairwise 的 $\frac{\lceil\log_2 N\rceil}{2}$ 倍)
切换点
当延迟项主导($m$ 小)时用 Bruck,否则用 Pairwise:
$$\begin{equation} m < m^* \approx \frac{(N-1)\alpha\beta}{\lceil\log_2 N\rceil - 1} \label{eq:a2a-bruck-crossover} \end{equation}$$算法对比
| 算法 | 步数 | 延迟项 | 带宽项 | 适用场景 |
|---|---|---|---|---|
| Pairwise | $N-1$ | $(N-1)\alpha$ | $\frac{(N-1)M}{N\beta}$ | 大消息,完全图 |
| Ring | $N-1$ | $(N-1)\alpha$ | $\frac{(N-1)M}{N\beta}$ | 大消息,Ring 拓扑 |
| Bruck | $\lceil\log_2 N\rceil$ | $\lceil\log_2 N\rceil\alpha$ | $\lceil\log_2 N\rceil \frac{(N-1)m}{2\beta}$ | 小消息 |
MoE 场景下的 AllToAll
MoE 中的角色
Mixture of Experts 模型中,每个 token 被路由到 Top-K 个 Expert:
- Dispatch AllToAll:将 token 从各节点按路由结果分发到持有对应 Expert 的节点
- Expert 计算
- Combine AllToAll:将 Expert 处理结果返回给原始节点
通信量估算
设 batch 中有 $B$ 个 token,每个 token 路由到 $K$ 个 Expert,共 $N$ 个节点,hidden dimension 为 $h$,均匀路由下:
$$\begin{equation} m = \frac{B \cdot K \cdot h \cdot \text{dtype}}{N} \label{eq:a2a-moe-msg-size} \end{equation}$$以 DeepSeek-V3($B=4096$, $K=8$, $h=7168$, $N=8$, BF16)为例:
$$\begin{equation} m = \frac{4096 \times 8 \times 7168 \times 2}{8} = 58.7 \text{ MB per pair} \label{eq:a2a-dsv3-msg-size} \end{equation}$$总 AllToAll 通信量:$N \times m = 469$ MB,属于大消息场景,应选 Pairwise 或 Ring 算法。
MoE AllToAll 的特殊性
不均匀通信量:路由结果决定每对节点间的通信量,热门 Expert 接收大量 token,冷门 Expert 几乎无 token。传统均匀分块算法效率降低,链路利用率不均。
Expert Capacity Factor:为防止不均匀路由导致某些 Expert 过载,设置容量因子 $C$:
$$\begin{equation} \text{capacity} = C \cdot \frac{B \cdot K}{E} \label{eq:a2a-expert-capacity} \end{equation}$$$C = 1.0$ 为无冗余,$C = 1.25$ 为 25% 冗余。超出 capacity 的 token 被丢弃(token drop),引入精度损失。
分层 AllToAll(2D Torus 上)
16 芯 2D Torus 上可分解为两步:
- 行内 AllToAll:每行 8 芯做 Ring AllToAll
- 本地数据重组:每节点对收到的数据做 local transpose(内存操作,非网络开销)
- 列内 AllToAll:每列 2 芯之间做 AllToAll
比直接在 16 芯上做 AllToAll 更高效,利用了 2D Torus 的结构化连接。
参考资料
- Bruck et al., "Efficient Algorithms for All-to-All Communications"(TPDS, 1997)- Bruck 算法($O(\log N)$ 步)
- Thakur et al., "Optimization of Collective Communication Operations in MPICH"(IJHPCA, 2005)
- Fedus et al., "Switch Transformers"(JMLR, 2022)- MoE 中 AllToAll 的使用
- Lepikhin et al., "GShard"(ICLR, 2021)- Expert Capacity 和 Token Drop
- DeepSeek-V3 Technical Report(DeepSeek, 2024)- MoE AllToAll 通信分析
- SCCL Documentation v1.0(SOPHGO)- AllToAll 基于 Scatter 实现