Infra: Collective Operations

2026-09-05
每行一张 GPU · root = R0颜色跟随数据块，下标表示分片逐元素求和的图例Σ₀ = A₀ + B₀ + C₀ + D₀，表示同一分片位置的逐元素求和。A₀+B₀+C₀+D₀=Σ₀
Broadcast：整份复制R0 的完整张量 A、B、C、D 被复制到所有 rank，每个 rank 都得到完整的一份。输入输出R0ABCDR1—R2—R3—ABCDABCDABCDABCD
Scatter：切开分发R0 持有 A、B、C、D 四块数据。分发后 R0 得到 A，R1 得到 B，R2 得到 C，R3 得到 D。输入输出R0ABCDR1—R2—R3—ABCD
Gather：收集拼接四个 rank 分别提供 A、B、C、D，在 R0 按 rank 顺序拼接为四块数据。只在 R0 指定输出。输入输出R0AR1BR2CR3DABCD———
Reduce：逐项相加四个 rank 的同形状数据 A、B、C、D 逐元素相加，R0 得到同形状的结果 Σ = A + B + C + D。只在 R0 指定输出。输入输出+R0AR1BR2CR3DΣ———
All-gather：每处都拼接四个 rank 分别提供 A、B、C、D，每个 rank 都得到按 rank 顺序拼接的完整数据 A、B、C、D。输入输出R0AR1BR2CR3DABCDABCDABCDABCD
Reduce-scatter：相加后分片各 rank 先按相同分片位置逐元素相加。Σᵢ = Aᵢ + Bᵢ + Cᵢ + Dᵢ。R0 只得到 Σ₀，R1 只得到 Σ₁，R2 只得到 Σ₂，R3 只得到 Σ₃。输入输出+R0A₀A₁A₂A₃R1B₀B₁B₂B₃R2C₀C₁C₂C₃R3D₀D₁D₂D₃Σ₀Σ₁Σ₂Σ₃
All-reduce：每处都得总和各 rank 按相同分片位置逐元素相加，每个 rank 都得到完整结果 Σ₀、Σ₁、Σ₂、Σ₃。输入输出+R0A₀A₁A₂A₃R1B₀B₁B₂B₃R2C₀C₁C₂C₃R3D₀D₁D₂D₃Σ₀Σ₁Σ₂Σ₃Σ₀Σ₁Σ₂Σ₃Σ₀Σ₁Σ₂Σ₃Σ₀Σ₁Σ₂Σ₃
All-to-all：按列交换每个 rank 将第 i 个分片发给 Ri。Ri 按来源 rank 顺序收到 Aᵢ、Bᵢ、Cᵢ、Dᵢ。输入的每一列成为输出的一行，没有求和或复制。可用图下方的单选按钮高亮一个目标 rank。输入输出R0A₀A₁A₂A₃R1B₀B₁B₂B₃R2C₀C₁C₂C₃R3D₀D₁D₂D₃A₀B₀C₀D₀A₁B₁C₁D₁A₂B₂C₂D₂A₃B₃C₃D₃
追踪去向
All-reduce = Reduce-scatter + All-gather
All-reduce 的两步分解左侧每个 rank 有四个输入分片。Reduce-scatter 后，中间的每个 rank 持有一个求和分片 Σᵢ。再经 All-gather，右侧每个 rank 都持有完整结果 Σ₀、Σ₁、Σ₂、Σ₃，与 All-reduce 的输出相同。这是结果等价的分解，不限定实际通信算法。输入每处一片每处一整份R0A₀A₁A₂A₃R1B₀B₁B₂B₃R2C₀C₁C₂C₃R3D₀D₁D₂D₃R0Σ₀R1Σ₁R2Σ₂R3Σ₃R0Σ₀Σ₁Σ₂Σ₃R1Σ₀Σ₁Σ₂Σ₃R2Σ₀Σ₁Σ₂Σ₃R3Σ₀Σ₁Σ₂Σ₃Reduce-scatterAll-gather
更详细的代码可以看 CS336 · Lecture 7。归约以 SUM 为例；箭头表示数据关系。虚线框表示无指定输入或输出，并不表示清空原数据。

Cite this post

@misc{pu2026mlrevisitinfracollectiveop,
  author = {Pu, Fanyi},
  title  = {Infra: Collective Operations},
  year   = {2026},
  month  = {9},
  url    = {https://pufanyi.com/blog/ml-revisit-infra-collective-op}
}