大模型规模化训练 · 第 3 篇

集合通信与分片矩阵乘法

从 broadcast 到 all-to-all，再到 $A[I_X, J] \cdot B[J, K_Y]$

按 → 或点击屏幕翻页 · 带交互的页面可以直接拖动

动机

一张卡放不下，或者太慢

70B 参数的 bf16 权重就要 140 GB，单张 80 GB 的 H100 连权重都放不下，更别说梯度、优化器状态和激活。
就算放得下，想把一步训练的时间缩短，也只能把计算摊到更多卡上。
于是要把张量切开（shard）放到很多卡上。Transformer 的计算几乎全是矩阵乘法，所以问题归结为：矩阵被切开以后，怎么乘？要搬多少数据？要多久？
回答这个问题需要两样东西：一套集合通信原语以及它们的代价公式；一套描述"谁拿着哪一块"的记号。

三个部分

N 个进程，每个叫一个 rank

一组独立的节点（CPU 核、GPU 或机器），每个算完自己那份，然后要和别人交换结果，才能进入下一步。
它们组成一个 process group，每个成员有编号 rank ∈ {0,…,N−1}，world_size = N。
有的操作有一个特殊成员 root：数据的唯一来源或唯一目的地。
集合（collective）操作：组内所有 rank 都要调用，一起完成。和点对点的 send / recv 相对。

import torch, torch.distributed as dist

def init_process():
    # 选后端，建立连接，分配 rank
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())

# 启动 3 个进程，每个占一张卡：
#   torchrun --nproc_per_node=3 dist_op.py

I · 原语

Broadcast：一份数据发给所有人

root（这里是 rank 0）有一份张量，其余 rank 拿到完整拷贝。典型用途：训练开始时把初始权重从 rank 0 同步给所有副本。

if rank == 0:
    t = torch.tensor([1., 2, 3, 4, 5]).cuda()
else:
    t = torch.zeros(5).cuda()
dist.broadcast(t, src=0)
# 之后每个 rank 都是 [1,2,3,4,5]

I · 原语

Reduce 与 AllReduce：把大家的数合起来

每个 rank 有一个同形状的张量，用一个函数 $f$（求和、平均、最大…）逐元素合并。Reduce 只把结果交给 root；AllReduce 让每个 rank 都拿到结果，等于 Reduce + Broadcast。数据并行的梯度同步就是一个 AllReduce。

t = torch.full((5,), rank + 1.).cuda()  # 1,2,3
dist.reduce(t, dst=0, op=dist.ReduceOp.SUM)
# rank0: [6,6,6,6,6]; rank1/2 不变
dist.all_reduce(t, op=dist.ReduceOp.SUM)
# 每个 rank: [6,6,6,6,6]

I · 原语

Gather 与 AllGather：把碎片拼成整体

每个 rank 有不同的一块，数据本身不做运算，只是搬运拼接。Gather 拼到 root 上；AllGather 让每个 rank 都拿到拼好的完整张量。ZeRO-3 / FSDP 在用到某层之前把这层被切开的权重 AllGather 回来。

out = [torch.zeros(5).cuda()
       for _ in range(world)]
dist.all_gather(out, t)
# 每个 rank: [t0, t1, t2]
# gather 只在 dst 上准备 out 列表：
dist.gather(t, out if rank == 0 else None,
            dst=0)

I · 原语

Scatter 与 ReduceScatter：切开再分发

Scatter 是 Gather 的反向：root 把张量切成 N 块，第 r 块发给 rank r（不同于 Broadcast 发整份）。ReduceScatter：每个 rank 都有一个分成 N 块的张量，先逐块求和，再把第 r 块的和只交给 rank r。图里先演示 Scatter，再演示 ReduceScatter。

# 3 个 rank，rank r 的第 j 块是
#   [r+1, 2(r+1)] 的 j+1 次方
inp = [torch.tensor([rank+1., 2*(rank+1)]).cuda()
       ** (j + 1) for j in range(3)]
out = torch.zeros(2).cuda()
dist.reduce_scatter(out, inp, op=dist.ReduceOp.SUM)
# rank 0: [1+2+3, 2+4+6] = [6, 12]
# rank 1: [14, 56]    rank 2: [36, 288]

I · 原语

AllToAll：分布式的"转置"

每个 rank 有 N 块，第 j 块要送给 rank j；收完以后 rank j 拿着"所有人的第 j 块"。把 rank 当行、块当列，这就是一次矩阵转置。不做运算，也不复制数据。

它不会自然出现在分片矩阵乘里，但在 MoE 里随处可见：每个 token 要送到它选中的 expert 所在的卡上（见专家并行）。分片记号里它表示"把下标从一个维度挪到另一个维度"。

I · 原语

Barrier：等所有人到齐

Barrier 不搬数据，只做同步：所有 rank 都到达之前，谁也不能继续往下走。

time.sleep(rank)          # rank r 睡 r 秒
dist.barrier()
# rank 0/1/2 都在约 2.0 秒后才通过

最快的 rank 一点没睡，也要等最慢的那个。所以 barrier 要慎用：很多时候让快的 rank 先往下做更好，下一轮它可能变成慢的那个，延迟会平均掉。

I · 原语 · 小结

九个操作，一张表

操作	调用前 rank r 有	调用后 rank r 有	PyTorch
Broadcast	root：$x$	$x$	`broadcast(t, src)`
Reduce	$x_r$	root：$\sum_r x_r$	`reduce(t, dst, op)`
AllReduce	$x_r$	$\sum_r x_r$	`all_reduce(t, op)`
Gather	$x_r$	root：$[x_0,\dots,x_{N-1}]$	`gather(t, list, dst)`
AllGather	$x_r$	$[x_0,\dots,x_{N-1}]$	`all_gather(list, t)`
Scatter	root：$[x^0,\dots,x^{N-1}]$	$x^r$	`scatter(t, list, src)`
ReduceScatter	$[x_r^0,\dots,x_r^{N-1}]$	$\sum_s x_s^r$	`reduce_scatter(out, list)`
AllToAll	$[x_r^0,\dots,x_r^{N-1}]$	$[x_0^r,\dots,x_{N-1}^r]$	`all_to_all(outs, ins)`
Barrier	—	—（同步点）	`barrier()`

下标 = 来自哪个 rank，上标 = 第几块。注意 AllReduce = ReduceScatter + AllGather，这也正是最省带宽的实现方式。

I · 原语 · 实现

NCCL：GPU 上用的那个库

PyTorch 支持三种后端：经典的 MPI、Meta 的 Gloo、NVIDIA 的 NCCL（读作 "nickel"）。接口相同，针对的硬件不同。
NCCL 专门优化 GPU 与 GPU 之间的通信（NVLink、InfiniBand、RoCE），MPI 和 Gloo 面向 CPU-CPU 或 CPU-GPU。经验法则：GPU 训练用 NCCL，CPU 训练用 Gloo。
NCCL 内部会按消息大小和拓扑选择算法（ring、tree 等）和协议。实测带宽会比"理论最优"模型低一些。
集合通信在 GPU 上是异步的 kernel，跑在单独的 CUDA stream 上，所以能和计算重叠。重叠是所有并行策略的核心手段。

II · 代价

α-β 模型：延迟项 + 带宽项

把一次消息传递的时间拆成两部分：

$$T_\text{msg}(n) = \alpha + \frac{n}{W}$$

$\alpha$：每一跳的固定开销（启动、同步、链路延迟），与数据量无关。TPU ICI 和 H100 NVLink 上都在 1 µs 左右。
$W$：链路带宽（字节/秒），$n$：这一跳发的字节数。
一个集合操作 = 若干步，每步所有链路并行发送。总时间 ≈ 步数 × α + 每条链路上最多的字节数 / W。

单向 vs 双向

一条全双工链路两个方向互不干扰，各有 $W_\text{uni}$。所谓"双向带宽" $W_\text{bi} = 2W_\text{uni}$。

H100 每张卡 NVLink：18 条 × 25 GB/s = 450 GB/s 每个方向。

TPU 的公式用双向带宽 $W_\text{bi}$（环面上两个方向同时用），GPU 的公式用单向出口带宽 $W$，两者差 2 倍。

II · 代价

笨办法：全部发给 root

AllReduce 一个 $V$ 字节的张量：其余 N−1 个 rank 都把整份发给 root，root 求和，再把整份广播回去。

$$T_\text{naive} \approx \frac{(N-1)V}{W} \cdot 2$$

root 的入口和出口各要扛 $(N-1)V$ 字节，其他链路大部分时间闲着。时间随 N 线性增长，卡越多越慢。

想要的效果：让每条链路同时都在忙，而且每个 rank 只收发 $O(V)$ 字节，与 N 无关。环形算法做到了这一点。

II · 代价 · 环形算法

环形 AllGather

把 N 个 rank 排成环，每个只向右邻居发、从左邻居收。张量总共 $V$ 字节，每个 rank 开始时有一块 $V/N$。

第 1 步：每个 rank 把自己那块发给右邻居。
第 k 步：把上一步收到的那块再转发给右邻居。
N−1 步后，每块都走遍了整个环，每个 rank 都有全部 N 块。

每步每条链路发 $V/N$ 字节，共 $N-1$ 步：

$$T_\text{AG} = (N-1)\,\alpha + \frac{N-1}{N}\cdot\frac{V}{W}$$

带宽项 $\to V/W$，与 N 无关。环上所有链路每一步都在工作，这是关键。

GPU 节点内是交换机而不是物理环，但可以按任意顺序"假装"成环，每跳都用满单卡出口带宽，结论相同。

II · 代价 · 环形算法

环形 ReduceScatter

每个 rank 把自己的张量切成 N 块。第 $k$ 步（$k$ 从 0 数起），rank $i$ 把第 $(i-k) \bmod N$ 块（已经累加了前面几个 rank 的贡献）发给右邻居，右邻居把自己的同一块加上去。

N−1 步以后，rank $i$ 手上第 $(i+1) \bmod N$ 块已经加了全部 N 个 rank 的贡献，就是这一块的最终和。

通信模式和 AllGather 完全一样，只是"收到后相加"而不是"收到后保存"：

$$T_\text{RS} = (N-1)\,\alpha + \frac{N-1}{N}\cdot\frac{V}{W}$$

另外有 $(N-1)V/N$ 次加法，通常远小于通信时间。

注意 $V$ 是整个张量的字节数（NCCL 里 = 每 rank 输出大小 × N）。

II · 代价 · 交互

环形 AllReduce = ReduceScatter + AllGather，一步一步看

格子里是"已累加的 GPU 数 / N"。实线框：本步正在发出的块；虚线框：本步接收方。橙色 = 已是最终和。发送量以 $V$（整个张量）为单位。

II · 代价

2(N−1)/N 是怎么来的

ReduceScatter 之后，rank $i$ 拥有第 $(i+1)\bmod N$ 块的完整和，每块恰好在一个 rank 上。再做一次 AllGather，把这些块发给所有人，每个 rank 就都有了完整的和，这正是 AllReduce 的定义。

$$\text{AllReduce} = \text{AllGather} \circ \text{ReduceScatter}$$

两阶段各 N−1 步、每步 $V/N$：

$$T_\text{AR} = 2(N-1)\,\alpha + \frac{2(N-1)}{N}\cdot\frac{V}{W}$$

和笨办法比

每个 rank 收发约 $2V$，不随 N 增长；笨办法的 root 要收发 $2(N-1)V$。

代价：步数从 2 变成 $2(N-1)$，延迟项随 N 线性增长。消息大时带宽项占主导，环形算法好；消息很小、N 很大时，延迟项会反过来占主导。

按元素数 $K$ 记：AllReduce 每卡收发约 $2K$ 个元素，ReduceScatter 和 AllGather 各约 $K$。

II · 代价

为什么不可能更快（只考虑点对点算法）

AllReduce 的每个输出元素都是 $N$ 个输入之和，而且最后要出现在 $N$ 个 rank 上。
每个元素至少需要 $N-1$ 次"加法所需的搬运"和 $N-1$ 次"把结果送到别处"的搬运，合计 $2(N-1)$ 次跨 rank 传输。
总共 $V$ 字节的元素 × $2(N-1)$ 次传输，由 $N$ 个 rank 的出口（每个带宽 $W$）分担：

$$t \;\ge\; \frac{V \cdot 2(N-1)}{N \cdot W}$$

环形算法正好达到这个下界，所以它是带宽最优的。ReduceScatter / AllGather 只需要其中一半：$V(N-1)/(NW)$。
唯一的例外是让交换机自己做加法（NVIDIA SHARP，in-network reduction），理论上可以把 AllReduce 降到 $V/W$；在 8×H100 节点内实测只提升约 30%（总线带宽约 370 → 480 GB/s）。

II · 代价 · 延迟

小消息：延迟项说了算

每一跳至少要花一个固定的时间 $T_\text{min}$，数据再少也省不掉：

$$T_\text{hop} = \max\!\left[T_\text{min},\ \frac{\text{每跳字节}}{W}\right]$$

当每跳的数据量小到传输时间不足 $T_\text{min}$ 时，带宽再高也没用，总时间 ≈ 跳数 × $T_\text{min}$，这时与 N 有关。

例：TPU v5e 单向 ICI 带宽 4.5e10 B/s，每跳 ~1 µs ⇒ 小于 $4.5\text{e}10 \times 10^{-6} = 45$ kB 的消息是延迟受限的。

GPU 上：8×H100 节点内 AllReduce 一个 128 B 的张量，实测约 16 µs（环形 14 跳，约 1 µs/跳）；跨 2 个、4 个节点约 25、35 µs。

树形算法

两两 AllReduce，再 4 个一组、8 个一组……只要 $\log_2 N$ 轮而不是 $N-1$ 轮，带宽总量不变，所以小消息、大 N 时更好。NCCL 会根据消息大小在 ring 和 tree 之间切换。

交叉点

环形 AllReduce 中延迟项 = 带宽项时：$2(N-1)\alpha = 2\frac{N-1}{N}\frac{V}{W}$，即 $V^\ast = N\alpha W$。消息比 $V^\ast$ 小就是延迟受限。这也是 DDP 要把梯度攒成"桶"再通信的原因（见数据并行）。

II · 代价 · TPU 环面

双向环、环面与多个轴

TPU 的芯片之间是 2D / 3D 环面（torus），每个芯片只和邻居相连，链路是双向的。AllGather 可以把每块同时往左右两个方向送：

只需 $\lfloor N/2\rfloor$ 跳，每跳每个方向发 $V/N$；
每跳耗时 $\frac{2V}{N \cdot W_\text{bi}}$，乘以 $N/2$ 跳：

$$T_\text{AG} = \max\!\left[\frac{T_\text{min}\cdot N}{2},\ \frac{V}{W_\text{bi}}\right]$$

在 $k$ 个 mesh 轴上同时做（例如 $\text{AllGather}_{XY}$），可用链路翻 $k$ 倍：

$$T = \max\!\left[\frac{T_\text{min}\sum_i |X_i|}{2},\ \frac{V}{W_\text{bi}\, N_\text{axes}}\right]$$

AllReduce 仍是 2 倍，ReduceScatter 与 AllGather 相同。只要是带宽受限，耗时只取决于数组大小和链路带宽，与分片数无关。

注意：v5e 只有某个轴长为 16 时才有回绕链路；没有回绕就只能单向走满 N−1 跳。

II · 代价 · AllToAll

AllToAll 为什么便宜

每个 rank 的 $V/N$ 字节再分成 N 小块（每块 $V/N^2$），第 j 小块只需要到 rank j，不需要复制给所有人。

单向环：rank 0 的小块分别要走 1, 2, …, N−1 跳，一条链路上合计 $(1+\dots+(N-1))\frac{V}{N^2} = \frac{N-1}{2N}V$，是 AllGather 的 1/2。

双向环：小块按近的方向走，最远 $N/2$ 跳；一条链路一个方向合计约 $\frac{N^2}{8}\cdot\frac{V}{N^2} = \frac{V}{8}$。AllGather 在双向化时只快 2 倍，AllToAll 快 4 倍：

$$T_\text{A2A} \approx \frac{V}{4\,W_\text{bi}} = \frac{T_\text{AG}}{4}$$

GPU 节点内（全互连交换）

每个 GPU 直接把 $N-1$ 小块发给目标：

$$T = \frac{V(N-1)}{N^2 W} \approx \frac{V}{N\,W}$$

跨节点

AllToAll 不是层级化的，跨 $M$ 个节点时 $T \approx V/(M\,W_\text{node})$。从单节点 $V/(8\times 450\text{e}9)$ 到两节点 $V/(2\times 400\text{e}9)$，慢了 4 倍多。MoE 专家并行最好放在节点内，原因就在这里。

II · 代价 · GPU 网络

H100 集群：两层带宽

层级	每卡/每节点出口	AllGather 时间
节点内 (NVLink 4)	450 GB/s / GPU	$\frac{V(N-1)}{N\cdot 450\text{e}9}$
跨节点 (IB NDR)	400 GB/s / 节点	$\approx V / 400\text{e}9$

跨节点时先在节点内归约，再在节点间做环，节点内带宽更高、几乎不会成为瓶颈：$T=\max[V/W_\text{GPU},\,V/W_\text{node}]$。
胖树保证任意二分都有满带宽，所以可以把任意多节点串成环，跨节点的集合带宽约为常数。
实测：8×H100 节点内 AllReduce 的总线带宽，不用交换机内归约（NVLink SHARP）时约 370 GB/s，用时约 480 GB/s，而且都要 GB 级的大消息才能接近。

II · 代价 · 测量

algbw 与 busbw：看 nccl-tests 的输出

algbw = $S / t$：数组大小除以耗时。直接用来估时间，但会随 N 变化，不好和硬件比较。

busbw = algbw × 修正因子，把 N 的影响去掉，直接和链路峰值（NVLink、PCIe、网卡）比较。

修正因子正是前面推出来的每种操作的"每条链路要承载的字节数 / S"。

操作	busbw / algbw	理想时间
AllReduce	$2(n-1)/n$	$\frac{S}{B}\cdot\frac{2(n-1)}{n}$
ReduceScatter	$(n-1)/n$	$\frac{S}{B}\cdot\frac{n-1}{n}$
AllGather	$(n-1)/n$	$\frac{S}{B}\cdot\frac{n-1}{n}$
AllToAll	$(n-1)/n$	$\frac{S}{B}\cdot\frac{n-1}{n}$
Broadcast	1	$S/B$（root 出口是瓶颈）
Reduce	1	$S/B$

nccl-tests 的 AllToAll 里 $S$ 是每个 rank 发出的总量，前面用的 $V$ 是整个数组，两者差 N 倍，换算后一致。

II · 代价 · 交互

集合通信耗时计算器

数组大小 V rank 数 N 单向链路带宽每跳延迟 α

估计耗时

—

带宽项延迟项（跳） algbw = V / T busbw 受限于

带宽项按环形算法的公式计算，延迟项 = 环形算法的跳数 × α。α 默认 1 µs：TPU ICI 每跳约 1 µs；8×H100 节点内小消息 AllReduce 实测约 16 µs，摊到 14 跳也是约 1 µs。

II · 代价 · 算例（GPU）

代入真实带宽

① 8×H100 节点内 AllGather

$\text{bf16}[4096, 65536]$：$V = 4096 \times 65536 \times 2 \approx 537$ MB。

$$T = \frac{537\text{e}6 \times 7}{8 \times 450\text{e}9} \approx 1.04\ \text{ms}$$

实测约 1.5 ms：带宽跑不到标称值，还有延迟。

② 8B 模型的 bf16 梯度 AllReduce

$V = 8\text{e}9 \times 2 = 16$ GB。

节点内 8 卡：$2 \cdot \frac{7}{8} \cdot \frac{16\text{e}9}{450\text{e}9} \approx 62$ ms

跨节点（层级化，按节点出口 400 GB/s）：$\approx 2 \times \frac{16\text{e}9}{400\text{e}9} = 80$ ms

若只能跑到约 370 GB/s，节点内会变成约 76 ms。它能不能藏在反向计算后面，决定了数据并行的效率。

③ AllToAll 跨节点： 同样 $V$ 字节，节点内 $V/(8 \times 450\text{e}9)$，跨 2 个节点 $V/(2 \times 400\text{e}9)$，慢 4.5 倍。④ 只跨 2 个节点的 AllGather： 节点间是 $V\cdot\frac{1}{2}/400\text{e}9 = V/800\text{e}9$，节点内是 $V\cdot\frac78/450\text{e}9 = V/514\text{e}9$，瓶颈反而在节点内。

II · 代价 · 算例（TPU）

带宽受限与延迟受限

① TPU v5e，$\text{AllGather}_Y([E_Y, F]) \to [E, F]$

mesh {X: 8, Y: 4}，$E=2048,\ F=8192$，bf16。整个数组 $2048\times8192\times2 \approx 34$ MB，每卡 8.4 MB。

理想双向环：$34\text{e}6 / 9\text{e}10 \approx 377$ µs。但长 4 的轴没有回绕链路，只能单向走 3 跳：$3 \times 8.4\text{e}6 / 4.5\text{e}10 \approx 560$ µs。实测约 680 µs。

② 同样的操作，$E = F = 256$

每卡只有 $64 \times 256 \times 2 = 32$ kB，一跳传输 $32\text{e}3/4.5\text{e}10 \approx 0.7$ µs，比 1 µs 的每跳开销还小 ⇒ 延迟受限。3 跳约 3 µs，实测约 8 µs。

教训

大数组只看带宽（与 N 无关）；小数组（比如推理时采样）只看跳数（与 N 有关）。实际系统总比理论慢一些，用公式估量级，用 profiler 定结论。

III · 分片矩阵乘法

Device mesh：给卡排个阵

把所有加速器排成一个 2D 或 3D 的阵列，每个轴起一个名字（比如 X、Y、Z），叫 mesh 轴（mesh axis）。

Mesh(devices=((0, 1), (2, 3)),
     axis_names=('X', 'Y'))
# 4 块卡排成 2×2，|X| = 2, |Y| = 2

mesh 轴不一定对应物理链路：{X: 16, Y: 4} 也可以描述一个 4×4×4 的 TPU 立方体，其中 X 占两个物理轴。GPU 上常见的是"节点内一个轴、节点间一个轴"。

III · 分片记号

$A[I_X, J_Y]$：下标写"沿哪个 mesh 轴切"

数组的每个维度有名字（I、J、…）。维度下标是 mesh 轴名，表示这个维度沿该 mesh 轴均匀切开。
$A[I, J]$：完全复制，每卡一份完整的。
$A[I_X, J]$：I 沿 X 切，沿 Y 部分复制（同一行的两张卡拿相同的块）。
$A[I_X, J_Y]$：两维都切，每卡 $1/(|X||Y|)$。
$A[I_{XY}, J]$：把 X、Y 合成一个大轴来切 I；下标顺序决定遍历顺序。
$A[I_X, J_X]$ 非法：一个 mesh 轴只能用一次，用过就"花掉了"。

全局形状 vs 本地形状

分片后的数组逻辑上形状不变（全局形状），每张卡实际存的是本地形状：

$$A[I_X, J_Y]:\ \left(\tfrac{|I|}{|X|},\ \tfrac{|J|}{|Y|}\right)$$

本地形状 × 每元素字节 = 每卡占用的显存。没用到的 mesh 轴上数据是复制的，所以总占用 = 单份大小 × 未使用轴长度之积。

JAX 里写作 PartitionSpec('X', 'Y')，None 表示不切。

III · 分片记号 · 交互

在 2×2 mesh 上摆一个 4×4 分块的 A

每个小方格是 A 的一个块（行 = I，列 = J），紫色 = 这张卡持有的块。设备 $d = 2x + y$，$x$ 是 X 轴坐标、$y$ 是 Y 轴坐标。

III · 分片记号 · 小测

算一算每张卡存多少

题 1

$\text{fp32}[1024, 4096]$，分片 $A[I_{XY}, J]$，mesh {X: 8, Y: 2}。每卡多少字节？H100 从 HBM（3.4e12 B/s）读一遍要多久？

本地形状 $[1024/16, 4096] = [64, 4096]$，$4 \times 64 \times 4096 = 1$ MiB；$1\text{e}6/3.4\text{e}12 \approx 294$ ns（这么小的数组实际会被各种开销淹没）。

题 2

$\text{int8}[128, 2048]$，分片 $A[I_{XY}, J]$，mesh {X: 2, Y: 8, Z: 2}（32 卡）。每卡多少？总共多少？

每卡 $[128/16, 2048] = [8, 2048]$ = 16,384 字节。Z 轴上复制 2 份，所以总量 = $128\times2048\times2 = 512$ KiB = 32 × 16 KiB。

题 3：$A[I_X, J, K, \ldots]$，mesh {X: 4, Y: 8, Z: 2}。所有卡上的总字节 / 单份大小 = ？答：每卡 1/4 份，共 64 卡 ⇒ 16 倍，正好是没用到的 $|Y|\cdot|Z|$。

III · 矩阵乘法

分块矩阵乘法与"收缩维度"

$C[I, K] = A[I, J] \cdot B[J, K]$，即 $C_{ik} = \sum_j A_{ij} B_{jk}$。被求和掉的 $J$ 叫收缩维度（contracting dimension），$I, K$ 是非收缩维度。

矩阵按块写时，乘法规则对块同样成立：

$$\begin{pmatrix} A_{00} & A_{01} \\ A_{10} & A_{11} \end{pmatrix}\!\begin{pmatrix} B_{00} & B_{01} \\ B_{10} & B_{11} \end{pmatrix} = \begin{pmatrix} A_{00}B_{00}+A_{01}B_{10} & \cdots \\ A_{10}B_{00}+A_{11}B_{10} & \cdots \end{pmatrix}$$

所以分布式矩阵乘法 = 把块搬到该去的地方 + 本地做块乘 + 把部分和加起来。问题只剩：要加什么通信，多贵。

逐元素运算（加法、激活函数）在分片数组上没有任何额外开销。只有沿被切开的维度做收缩时才需要通信。所有情况可以归成 4 种。

III · 情形 1

收缩维度都没切：不需要通信

$$A[I_X, J] \cdot B[J, K_Y] \;\to\; C[I_X, K_Y]$$

坐标为 $(x, y)$ 的卡有 A 的第 $x$ 个行块（完整的 J）和 B 的第 $y$ 个列块（完整的 J），正好能算出 C 的 $(x, y)$ 块，不缺任何数据。
输出的分片直接继承输入的非收缩维度分片。

$$\begin{aligned} A[I, J]\cdot B[J, K] &\to C[I, K] \\ A[I_X, J]\cdot B[J, K] &\to C[I_X, K] \\ A[I, J]\cdot B[J, K_Y] &\to C[I, K_Y] \\ A[I_X, J]\cdot B[J, K_Y] &\to C[I_X, K_Y] \end{aligned}$$

前提：两个输入的非收缩维度没有用同一个 mesh 轴（那是情形 4）。

III · 情形 2

只有一个输入的收缩维度被切：先 AllGather

$$A[I, J_X] \cdot B[J, K] \;\to\; C[I, K]$$

每张卡只有 A 的一部分列，不能直接和 B 乘出完整的和。常规做法：

$$\textbf{AllGather}_X\, A[I, J_X] \to A[I, J]$$

然后每卡做完整的本地乘法。AllGather 在记号上就是删掉一个下标；也可以只删一部分：$A[I_{XY}, J] \to A[I_Y, J]$ 是只沿 X 的 AllGather。

另一种选择：如果 B 在 X 上是复制的，可以把 B 本地切成 $B[J_X, K]$（免费，丢掉一部分即可），本地乘出部分和，再 AllReduce 输出（变成情形 3）。这样计算也被分摊了，但通信的是输出 $C$ 而不是输入 $A$。哪个更好取决于 $|A|$ 和 $|C|$ 谁大（见"怎么选"）。实际中 B 通常也已经被切了，所以这种情况不常见。

III · 情形 3

两个输入的收缩维度沿同一轴切：本地乘，再 AllReduce

$$A[I, J_X] \cdot_\text{LOCAL} B[J_X, K] \;\to\; C[I, K]\{U_X\} \quad\xrightarrow{\ \textbf{AllReduce}_X\ }\quad C[I, K]$$

每张卡拿着同一段 J，可以本地乘，但得到的只是部分和。记号 $\{U_X\}$ 读作"沿 X 未归约"（unreduced）。从外积的角度看：

$$A \cdot B = \sum_{j} A_{:,j} \otimes B_{j,:}$$

每张卡算了这个求和里属于自己那段 $j$ 的项，剩下的求和要跨卡完成。

AllReduce 删掉 $\{U_X\}$：$\textbf{AllReduce}_Y\, A[I_X, J]\{U_Y\} \to A[I_X, J]$，其他分片不变。它比 AllGather 贵：被归约的每个"分片"都有完整输出那么大。AllReduce = ReduceScatter + AllGather，所以是 AllGather 的 2 倍：$T_\text{AR} = \frac{2V}{W_\text{bi}}$ (TPU) 或 $\frac{2(N-1)V}{N\,W}$ (GPU)。

III · 情形 4

两个非收缩维度用了同一个轴：非法，先 AllGather 一个

$$A[I_X, J] \cdot B[J, K_X] \;\to\; C[I_X, K_X]\ \ \text{✗}$$

X 坐标为 $i$ 的卡只有 A 的第 $i$ 行块和 B 的第 $i$ 列块，只能算出 C 的 $(i, i)$ 块，也就是对角块。非对角块的信息根本不存在于任何一张卡上，而且输出写成 $C[I_X, K_X]$ 违反了"一个轴只用一次"。

两种解法，看后续算子需要哪种输出分片；只看通信量，就 gather 小的那个：

$$\begin{aligned} &\textbf{AllGather}_X A[I_X, J] \to A[I, J];\ \ A[I,J]\cdot B[J,K_X] \to C[I, K_X] \\ &\textbf{AllGather}_X B[J, K_X] \to B[J, K];\ \ A[I_X,J]\cdot B[J,K] \to C[I_X, K] \end{aligned}$$

III · 原语的记号

四个核心原语 = 四种改下标的方式

操作	记号	作用	带宽受限时的耗时
AllGather	$[A_X, B] \to [A, B]$	删掉一个下标	$V / (W_\text{bi} \cdot N_\text{axes})$；GPU：$\frac{N-1}{N}V/W$
ReduceScatter	$[A, B]\{U_X\} \to [A_X, B]$	删掉 $\{U_X\}$，同时加一个下标	同 AllGather
AllReduce	$[A_X, B]\{U_Y\} \to [A_X, B]$	删掉 $\{U_Y\}$	2 × AllGather
AllToAll	$[A, B_X] \to [A_X, B]$	把下标挪到另一维	双向环 AllGather / 4；GPU 节点内约 $V/(NW)$

$V$ 一律指整个数组（所有分片合起来、但只算一份）的字节数。如果数组还沿另一个轴 Y 切着，每组只需处理 $V/|Y|$（TPU 上直接减少；GPU 上要看 Y 是否跨节点）。

严格说 AllReduce 不是原语，它就是 ReduceScatter 后接 AllGather。

III · 一个漂亮的事实

AllGather 与 ReduceScatter 互为转置

设 $u = (1,\dots,1) \in \mathbb{R}^p$（$p$ 个设备）。广播和归约都是线性算子：

$$\text{broadcast} = u \otimes I_n,\qquad \text{reduce} = u^\top \otimes I_n$$

所以 $\text{reduce} = \text{broadcast}^\top$。AllGather、ReduceScatter 分别是它们与 $I_p$ 的 Kronecker 积：

$$\text{AG} = \text{broadcast} \otimes I_p,\quad \text{RS} = \text{reduce} \otimes I_p \;\Rightarrow\; \text{AG}^\top = \text{RS}$$

这对反向传播意味着什么

对线性算子 $y = Mx$，有 $\partial L/\partial x = M^\top\, \partial L/\partial y$。因此：

前向 $\text{AllGather}_X A[I_X] \to A[I]$ ⇒ 反向 $\text{ReduceScatter}_X A'[I]\{U_X\} \to A'[I_X]$
前向 ReduceScatter ⇒ 反向 AllGather

FSDP 前向 AllGather 权重，反向就 ReduceScatter 权重梯度，不是巧合，是求导的必然结果。

III · 情形 3 的变体

输出不必完整：用 ReduceScatter 代替 AllReduce

$$A[I, J_X] \cdot_\text{LOCAL} B[J_X, K] \to C[I, K]\{U_X\} \xrightarrow{\ \textbf{ReduceScatter}_{X,K}\ } C[I, K_X]$$

AllReduce = RS + AG，最后的 AllGather 往往可以推迟甚至省掉：下一个算子要的本来就是切开的输入。
RS 必须选一个维度来"加下标"（这里选 K，也可以选 I），所以写成 $\text{ReduceScatter}_{X,K}$。通常由上下文决定。
省一半通信：这正是 ZeRO-2 把数据并行里的梯度 AllReduce 换成 ReduceScatter，以及序列并行把张量并行里的 AllReduce 拆开的基础。

III · 交互

分片矩阵乘法探索器：$A[I, J] \cdot B[J, K] \to C[I, K]$

A 的 I：

A 的 J：

B 的 J：

B 的 K：

|X| |Y| I J K

bf16；GPU 模型：每卡单向，989 TFLOP/s

通信时间每卡计算若完全重叠

III · 怎么选

先 gather 输入，还是先乘再归约输出？

$X[B, D] \cdot Y[D_X, F] \to Z[B, F]$，记每卡算力 $C$，双向带宽 $W$：

策略 1：AllGather 权重 Y

$$T_1 = \max\!\left(\frac{2BDF}{C},\ \frac{2DF}{W}\right)$$

计算没有分摊（每卡都做全部 FLOPs），通信量是权重大小。$B \gt C/W$ 时计算受限。

策略 2：本地乘 + AllReduce 输出

$$T_2 = \max\!\left(\frac{2BDF}{X\,C},\ \frac{4BF}{W}\right)$$

FLOPs 分摊了 X 倍，但要 AllReduce $2BF$ 字节的输出。只有 $X \lt D/(2\cdot C/W)$ 才计算受限，D≈8k 时几乎总是通信受限。

TPU v5p 上 $C/W \approx 2550$。小 batch（$B \lt 2550$）时两者都通信受限，比较 $4BF$ 和 $2DF$：$D \gt 2B$ 时策略 2 更好。大 batch 时比较 $4BF/W$ 和 $2BDF/C$：$D \gt 5100$ 时策略 2 更好。结论：输出比输入小时，归约输出更划算。

III · 重叠

把通信藏在矩阵乘法后面

上面所有估计都默认：只要通信时间不超过计算时间，就能完全重叠。实现上用的是 collective matmul：

把矩阵乘按块拆开；
算第 k 块的同时，第 k−1 块的部分结果已经在环上传递（或下一块输入正在被 gather）；
最后只剩一小块通信露在外面。

实现并不简单：需要两个 stream、正确的依赖和足够大的块。

III · 并行策略

并行策略 = 给几个矩阵选分片

把 Transformer 简化成一串 MLP：$\text{In}[B, D] \cdot W_\text{in}[D, F] \cdot W_\text{out}[F, D]$，$B$ 是 token 数。

策略	分片	出现的情形 / 通信
数据并行	$\text{In}[B_X, D] \cdot W_\text{in}[D, F] \cdot W_\text{out}[F, D]$	前向情形 1；反向 $dW = \text{In}^\top \cdot d\text{Tmp}$ 沿 $B_X$ 收缩 ⇒ 情形 3，AllReduce 梯度
FSDP / ZeRO-3	$\text{In}[B_X, D] \cdot W_\text{in}[D_X, F] \cdot W_\text{out}[F, D_X]$	前向 AllGather 权重；反向 ReduceScatter 梯度（AG 的转置）
张量并行	$\text{In}[B, D_Y] \cdot W_\text{in}[D, F_Y] \cdot W_\text{out}[F_Y, D]$	进块前 AllGather 激活，出块后 ReduceScatter 激活

数据并行的 AllReduce 什么时候藏得住、ZeRO 怎么把 AllReduce 拆成 RS + AG 来省显存，见数据并行与 ZeRO。

小结

记住这几条

AllReduce = ReduceScatter + AllGather。环形实现带宽最优，每卡收发 $2\frac{N-1}{N}V$。
带宽受限时，AG / RS ≈ $V/W$，AR ≈ $2V/W$，与卡数无关；小消息时延迟项 ∝ 跳数才是瓶颈。
AllToAll 只送不复制：TPU 双向环上是 AllGather 的 1/4，GPU 节点内约 $V/(NW)$，跨节点退化严重。
H100：节点内每卡 450 GB/s，跨节点每节点 400 GB/s；实测往往只到七八成。
分片矩阵乘：收缩维没切 → 不通信；一边切 → AllGather；两边同轴切 → AllReduce / ReduceScatter；非收缩维撞轴 → 先 AllGather 一个。
AllGather 的导数是 ReduceScatter，反之亦然。

收尾

参考与延伸阅读

Ultra-Scale Playbook，附录 A0: Parallel Programming Crash Course：第 I 部分各个原语的 PyTorch 例子（包括 ReduceScatter 的那个算例）、Ring AllReduce 和 NCCL 的介绍。
Austin et al., How to Scale Your Model：Part 3 · Sharded Matrices and How to Multiply Them（分片记号、四种情形、AllToAll、ReduceScatter，以及 TPU 算例）；Part 12 · GPU 上的集合通信（H100 网络、SHARP、8×H100 实测带宽）；Part 10 · Programming TPUs in JAX 有 collective matmul 的实现。
Wang et al., Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models（ASPLOS 2023），collective matmul 的原始论文。
NVIDIA nccl-tests，PERFORMANCE.md：algbw / busbw 的定义和 2(n−1)/n 下界的推导。小消息延迟的数字来自公开的 nccl-tests 结果：单节点 8×H100、2 / 4 节点 H100。

← 回到博客

操作	busbw / algbw	理想时间
AllReduce	\(2(n-1)/n\)	\(\frac{S}{B}\cdot\frac{2(n-1)}{n}\)
ReduceScatter	\((n-1)/n\)	\(\frac{S}{B}\cdot\frac{n-1}{n}\)
AllGather	\((n-1)/n\)	\(\frac{S}{B}\cdot\frac{n-1}{n}\)
AllToAll	\((n-1)/n\)	\(\frac{S}{B}\cdot\frac{n-1}{n}\)
Broadcast	1	\(S/B\)（root 出口是瓶颈）
Reduce	1	\(S/B\)

操作	调用前 rank r 有	调用后 rank r 有	PyTorch
Broadcast	root：\(x\)	\(x\)	`broadcast(t, src)`
Reduce	\(x_r\)	root：\(\sum_r x_r\)	`reduce(t, dst, op)`
AllReduce	\(x_r\)	\(\sum_r x_r\)	`all_reduce(t, op)`
Gather	\(x_r\)	root：\([x_0,\dots,x_{N-1}]\)	`gather(t, list, dst)`
AllGather	\(x_r\)	\([x_0,\dots,x_{N-1}]\)	`all_gather(list, t)`
Scatter	root：\([x^0,\dots,x^{N-1}]\)	\(x^r\)	`scatter(t, list, src)`
ReduceScatter	\([x_r^0,\dots,x_r^{N-1}]\)	\(\sum_s x_s^r\)	`reduce_scatter(out, list)`
AllToAll	\([x_r^0,\dots,x_r^{N-1}]\)	\([x_0^r,\dots,x_{N-1}^r]\)	`all_to_all(outs, ins)`
Barrier	—	—（同步点）	`barrier()`

层级	每卡/每节点出口	AllGather 时间
节点内 (NVLink 4)	450 GB/s / GPU	\(\frac{V(N-1)}{N\cdot 450\text{e}9}\)
跨节点 (IB NDR)	400 GB/s / 节点	\(\approx V / 400\text{e}9\)

操作	记号	作用	带宽受限时的耗时
AllGather	\([A_X, B] \to [A, B]\)	删掉一个下标	\(V / (W_\text{bi} \cdot N_\text{axes})\)；GPU：\(\frac{N-1}{N}V/W\)
ReduceScatter	\([A, B]\{U_X\} \to [A_X, B]\)	删掉 \(\{U_X\}\)，同时加一个下标	同 AllGather
AllReduce	\([A_X, B]\{U_Y\} \to [A_X, B]\)	删掉 \(\{U_Y\}\)	2 × AllGather
AllToAll	\([A, B_X] \to [A_X, B]\)	把下标挪到另一维	双向环 AllGather / 4；GPU 节点内约 \(V/(NW)\)

策略	分片	出现的情形 / 通信
数据并行	\(\text{In}[B_X, D] \cdot W_\text{in}[D, F] \cdot W_\text{out}[F, D]\)	前向情形 1；反向 \(dW = \text{In}^\top \cdot d\text{Tmp}\) 沿 \(B_X\) 收缩 ⇒ 情形 3，AllReduce 梯度
FSDP / ZeRO-3	\(\text{In}[B_X, D] \cdot W_\text{in}[D_X, F] \cdot W_\text{out}[F, D_X]\)	前向 AllGather 权重；反向 ReduceScatter 梯度（AG 的转置）
张量并行	\(\text{In}[B, D_Y] \cdot W_\text{in}[D, F_Y] \cdot W_\text{out}[F_Y, D]\)	进块前 AllGather 激活，出块后 ReduceScatter 激活

集合通信与分片矩阵乘法

一张卡放不下，或者太慢

目录

I · 原语

II · 代价

III · 分片矩阵乘法

N 个进程，每个叫一个 rank

Broadcast：一份数据发给所有人

Reduce 与 AllReduce：把大家的数合起来

Gather 与 AllGather：把碎片拼成整体

Scatter 与 ReduceScatter：切开再分发

AllToAll：分布式的"转置"

Barrier：等所有人到齐

九个操作，一张表

NCCL：GPU 上用的那个库

α-β 模型：延迟项 + 带宽项

单向 vs 双向

笨办法：全部发给 root

环形 AllGather

环形 ReduceScatter

环形 AllReduce = ReduceScatter + AllGather，一步一步看

2(N−1)/N 是怎么来的

和笨办法比

为什么不可能更快（只考虑点对点算法）

小消息：延迟项说了算

树形算法

交叉点

双向环、环面与多个轴

AllToAll 为什么便宜

GPU 节点内（全互连交换）

跨节点

H100 集群：两层带宽

algbw 与 busbw：看 nccl-tests 的输出

集合通信耗时计算器

代入真实带宽

① 8×H100 节点内 AllGather

② 8B 模型的 bf16 梯度 AllReduce

带宽受限与延迟受限

① TPU v5e，\(\text{AllGather}_Y([E_Y, F]) \to [E, F]\)

② 同样的操作，\(E = F = 256\)

教训

Device mesh：给卡排个阵

\(A[I_X, J_Y]\)：下标写"沿哪个 mesh 轴切"

全局形状 vs 本地形状

在 2×2 mesh 上摆一个 4×4 分块的 A

算一算每张卡存多少

题 1

题 2

分块矩阵乘法与"收缩维度"

收缩维度都没切：不需要通信

只有一个输入的收缩维度被切：先 AllGather

两个输入的收缩维度沿同一轴切：本地乘，再 AllReduce

两个非收缩维度用了同一个轴：非法，先 AllGather 一个

四个核心原语 = 四种改下标的方式

AllGather 与 ReduceScatter 互为转置

这对反向传播意味着什么

输出不必完整：用 ReduceScatter 代替 AllReduce

分片矩阵乘法探索器：\(A[I, J] \cdot B[J, K] \to C[I, K]\)

先 gather 输入，还是先乘再归约输出？

策略 1：AllGather 权重 Y

策略 2：本地乘 + AllReduce 输出

把通信藏在矩阵乘法后面

并行策略 = 给几个矩阵选分片

记住这几条

参考与延伸阅读