ICML 2022 · Intra-op

张量并行与 AllReduce

一张卡装不下大模型时，把同一个矩阵乘切开、多卡一起算。关键是：切哪一维，要不要 AllReduce。

配套 code/03_tp_shapes_and_allreduce.py · 不必上 4 机 Docker

1. 为什么要并行（两种切法）

GPT-3 量级权重几百 GB，单卡只有几十 GB。办法：把图切开，多设备一起算，并付通信账。

视角人话常见名通信
Inter-op不同层放到不同卡，像流水线PP点对点；容易空等（bubble）
Intra-op同一层的矩阵切开，多卡各算一块TP集合通信（AllReduce 等）；卡更忙

真实训练常是 DP × TP × PP（3D 并行）。本页把 TP / AllReduce 钉死即可。

Inter vs Intra Inter-op / PP 卡1：第 1～8 层 卡2：第 9～16 层 激活在卡之间传递 有人忙、有人等 Intra-op / TP 同一层 Linear 每卡只持有权重的一块 算完要集合通信 设备更满载
图 1 · PP 切「层」；TP 切「同一个算子里的块」

2. 集合通信先认这四个

原语直觉
AllReduce各卡一份同 shape 张量 → 加总（常 sum）后，每张卡都有完整结果
ReduceScatter先加总，再按块散开，每卡只留一块
AllGather各卡一块 → 拼成完整，每卡都有全量
AllToAll按分区对换（MoE 路由常用）
 
AllReduce ≈ ReduceScatter + AllGather

DDP 训练里，梯度同步就是隐式 AllReduce。

3. 一个 Matmul 怎么切（飞书自测：切 m / n / k）

C[m,n] = A[m,k] @ B[k,n]，设备数 p。

切 m n k 切 m（输出行） 每卡 A 的行块 × 完整 B → C 的行块 前向通常无 AllReduce 切 n（输出列 · Megatron 列并行） 每卡完整 A × B 的列块 → C 的列块 下一层若要完整向量，可能 AllGather；或配「行并行」抵消 切 k（规约维） 每卡 A 列块 × B 行块 → 部分 C 必须 AllReduce(sum) 才能得到完整 C
图 2 · 三种切法；只有切 k 前向必 AllReduce

脚本会打印具体 shape：

python code/03_tp_shapes_and_allreduce.py

4. Megatron MLP：列为行配对 + 一次 AllReduce

x ──► Linear W1（列并行）──► GeLU ──► Linear W2（行并行）──► AllReduce ──► +残差
  1. W1 按输出通道切开 → 每卡算一部分激活；GeLU 本地做
  2. W2 按输入通道切开（和上一层对齐）→ 每卡只有部分输出
  3. 前向末尾 AllReduce 把部分输出加总
  4. 反向在对称位置再来一次 AllReduce

Attention 的 QKV（列）和 Wo（行）也是这套配对，少做多余 AllGather。

通信量直觉：一层 MLP 前向 AllReduce 元素量约 B · T · H（激活）。卡很多时启动延迟变大，朴素 Ring AllReduce 不一定最优。

列并行接行并行 x 完整 W1 列切 每卡 I/p 列 W2 行切 部分 y AllReduce 完整 y 列切的输出维 = 下一层行切的输入维，中间不必 AllGather
图 3 · 为何「列并行 + 行并行」成对出现

5. Pipeline 对照（浅记即可）

6. 和编译器主线

话题关系
选切轴、插 AllReduce图级并行策略
通信与计算重叠调度 / 通算融合
单算子内 tilingTriton（你的主战场）

7. 自测

  1. 切 m / n / k：每卡存什么？要不要 AllReduce？
  2. 8 卡纯 TP，一层里 AllReduce 出现在哪些投影之后？
  3. 为什么列并行 Linear 后面常接行并行 Linear？
python code/03_tp_shapes_and_allreduce.py
# 可选：
torchrun --nproc_per_node=2 code/03_tp_shapes_and_allreduce.py --dist

← MLA · 下一篇：torch.compile →