多卡并行——通信原语

多卡并行（主要指数据并行 / 模型并行）中，通信操作本质是在多进程/多设备之间搬运张量。下面按你图里的典型操作逐一说明，并给出直观语义与常见用途。

1. Reduce

语义：多卡各自有一个张量 → 按某个规约算子合成一个结果 → 只保留在指定一张卡
[
y=op(x0,x1,…,xN−1)y = \text{op}(x_0, x_1, \dots, x_{N-1})y=op(x0​,x1​,…,xN−1​)
]

常见算子：sum / mean / max / min
用途：

• 统计全局 loss（只在 rank0 用）
• 收集指标但不回传给所有卡

2. Broadcast

语义：一张卡上的张量 → 复制到所有卡

x0→x0,x0,…,x0x_0 \rightarrow x_0, x_0, \dots, x_0x0​→x0​,x0​,…,x0​

用途：

• 初始化模型参数（rank0 加载，其余同步）
• 同步配置、超参数

3. Gather

语义：每张卡一个张量 → 按 rank 顺序拼接 → 放在指定一张卡

y=[x0∣x1∣…∣xN−1]y = [x_0 | x_1 | \dots | x_{N-1}]y=[x0​∣x1​∣…∣xN−1​]

特点：

• 输出 shape 会变大
• 非对称（只有 root 有结果）

用途：

• 验证 / 推理阶段收集所有样本结果
• 日志、评测（只在 rank0）

4. Scatter

语义：一张卡有一个“大张量” → 切分 → 分发给各卡

[x0,x1,…,xN−1]→xi[x_0, x_1, \dots, x_{N-1}] \rightarrow x_i[x0​,x1​,…,xN−1​]→xi​

用途：

• 将数据 batch 切分到不同 GPU
• pipeline / tensor 并行中的输入分发

5. All-Reduce

语义：Reduce + Broadcast

y=op(x0,…,xN−1),每张卡都拿到 yy = \text{op}(x_0, \dots, x_{N-1}),\quad \text{每张卡都拿到 } yy=op(x0​,…,xN−1​),每张卡都拿到y

这是训练中最核心的操作

用途：

• 梯度同步（DDP 标配）

  grad_i → AllReduce(sum) → grad_global

关键点：

• 通信量大，但高度优化（ring / tree）
• 通常与 backward 重叠（bucket + async）

6. All-Gather

语义：Gather + Broadcast

y=[x0∣x1∣…∣xN−1],每张卡都有 yy = [x_0 | x_1 | \dots | x_{N-1}],\quad \text{每张卡都有 } yy=[x0​∣x1​∣…∣xN−1​],每张卡都有y

用途：

• 对比学习（CLIP / InfoNCE）：需要“全局 batch”
• MoE / 检索 / 多卡负样本
• 模型并行中拼接中间特征

常见陷阱：

• 显存暴涨（batch × world_size）
• 反向梯度是否需要all_gather_with_grad

总结对照表