多卡并行——通信原语
2026/5/9 9:16:19 网站建设 项目流程

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


1. Reduce

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

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

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

2. Broadcast

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

x0→x0,x0,…,x0x_0 \rightarrow x_0, x_0, \dots, x_0x0x0,x0,,x0

用途

  • 初始化模型参数(rank0 加载,其余同步)
  • 同步配置、超参数

3. Gather

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

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

特点

  • 输出 shape 会变大
  • 非对称(只有 root 有结果)

用途

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

4. Scatter

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

[x0,x1,…,xN−1]→xi[x_0, x_1, \dots, x_{N-1}] \rightarrow x_i[x0,x1,,xN1]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,,xN1),每张卡都拿到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=[x0x1xN1],每张卡都有y

用途

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

常见陷阱

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

总结对照表

操作输入输出位置是否对称典型用途
ReduceN → 1单卡指标、统计
Broadcast1 → N所有卡参数同步
GatherN → 1单卡验证收集
Scatter1 → N所有卡数据分发
All-ReduceN → N所有卡梯度同步
All-GatherN → N所有卡全局特征

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询