Collective

Collective Communication（集合通信）描述一组 Rank 共同参与的数据交换。它不关心模型结构，而是规定数据从哪些 Rank 出发、经过怎样的聚合或切分、最终落到哪些 Rank；Tensor Parallelism（TP）、Sequence Parallelism（SP）和 Expert Parallelism（EP）都建立在这些原语之上。

图中左侧是操作前的数据布局，右侧是操作后的数据布局。R0★ 表示 Root Rank，颜色表示分片位置，Σ 表示使用 $oplus$ 得到的逐元素归约结果。Reduce-Scatter 中的 $A_i$ 表示 Rank $i$ 持有的 A 段；All-to-All 中的 $x_{ij}$ 表示从源 Rank $i$ 发往目标 Rank $j$ 的数据。该图描述的是数据语义，不限定底层采用 Ring、Tree 或其他通信拓扑。

通信模型

设通信组包含 $p$ 个 Rank，编号为 $0,\dots ,p-1$。Rank $r$ 的本地输入记为 $X_r$，Reduction 使用的逐元素结合操作记为 $\oplus$，例如 Sum、Max 或 Min。

Collective 只定义输入与输出的数据关系，不规定底层必须采用 Ring、Tree 还是其他算法。同一种 Collective 可以根据消息大小、设备拓扑和通信库选择不同实现。

操作	数据关系	结果所在位置
Broadcast	Root 的完整数据复制到所有 Rank	所有 Rank
Scatter	Root 的完整数据切成 $p$ 份	每个 Rank 一份
Gather	所有 Rank 的分片按 Rank 顺序汇集	仅 Root
All-Gather	所有 Rank 的分片按 Rank 顺序汇集	所有 Rank
Reduce	所有 Rank 的同形数据执行 Reduction	仅 Root
All-Reduce	所有 Rank 的同形数据执行 Reduction	所有 Rank
Reduce-Scatter	先执行 Reduction，再把结果切成 $p$ 份	每个 Rank 一份
All-to-All	每个 Rank 分别向所有目标 Rank 发送不同分片	每个 Rank 收到来自所有源 Rank 的分片

Root-based Collectives

Root-based Collective 以一个指定 Rank 为数据源或汇聚点。这里的 Root 是通信组内的 Rank 编号，不等同于物理 GPU 编号。

Broadcast

Broadcast 将 Root Rank 的输入复制到整个通信组：

$$Y_r=X_{\mathrm{root}},r=0,\dots ,p-1$$

所有输出形状都与 Root 输入相同。它常用于同步模型配置、控制信息或需要在各 Rank 保持一致的小型张量。

Scatter 与 Gather

Scatter 将 Root 上的完整输入沿指定维度切成 $p$ 个分片，并把第 $r$ 个分片发送给 Rank $r$：

$$X_{\mathrm{root}}=[X^{(0)};X^{(1)};\dots ;X^{(p-1)}],Y_r=X^{(r)}$$

Gather 执行相反过程。每个 Rank 提供一个分片，Root 按 Rank 顺序接收并拼接：

$$Y_{\mathrm{root}}=[X_0;X_1;\dots ;X_{p-1}]$$

两者适合由单个协调 Rank 分发输入或收集最终输出，但 Root 需要承载完整张量和集中流量。

Reduce

Reduce 对所有 Rank 的同形输入执行逐元素 Reduction，只在 Root 上保留结果：

$$Y_{\mathrm{root}}=X_0\oplus X_1\oplus \cdots \oplus X_{p-1}$$

当 Reduction 为 Sum 时，对应元素会在所有 Rank 间求和。其他 Rank 参与计算，但不获得最终完整结果。

All-rank Collectives

All-rank Collective 不把完整结果限制在 Root。它们通常用于保持分布式张量布局，或者让每个 Rank 获得后续计算所需的数据。

All-Gather

All-Gather 收集每个 Rank 的本地分片，并把完整拼接结果发送给所有 Rank：

$$Y_r=[X_0;X_1;\dots ;X_{p-1}],r=0,\dots ,p-1$$

若每个 $X_r$ 含 $n$ 个元素，则每个 Rank 的输出包含 $pn$ 个元素。它完成的是 Shard → Replicate：通信后显存占用随通信组大小增加。

All-Reduce

All-Reduce 对所有输入执行 Reduction，并让每个 Rank 都获得相同的完整结果：

$$Y_r=X_0\oplus X_1\oplus \cdots \oplus X_{p-1},r=0,\dots ,p-1$$

从语义上看，Reduce + Broadcast 等价于 All-Reduce。对于大张量，Ring All-Reduce 常拆成两个阶段：

Partial → Reduce-Scatter → Reduced Shard → All-Gather → Replicated Result

设每个 Rank 的输入大小为 $N$ Byte，Ring 算法包含 $2(p-1)$ 个通信步骤；每个 Rank 的理论发送量约为下式，接收量与之相同：

$$V_{\mathrm{ring}\mathrm{allreduce}}=2\frac{p-1}{p}N$$

它能较充分利用链路带宽，但步骤数随 $p$ 增长，小消息更容易受启动延迟影响。

Reduce-Scatter

Reduce-Scatter 是 TP、SP 和分片训练中的核心原语。它先对所有 Rank 的输入执行 Reduction，再让 Rank $r$ 只保留结果的第 $r$ 个分片：

$$Z=X_0\oplus X_1\oplus \cdots \oplus X_{p-1},Y_r={\mathrm{chunk}}_r(Z)$$

若每个输入包含 $pn$ 个元素，则每个 Rank 最终只接收 $n$ 个元素。它完成的是 Partial → Shard，避免像 All-Reduce 一样在所有设备上复制完整结果。Reduce-Scatter + All-Gather 在数据关系上等价于 All-Reduce。

All-to-All

All-to-All 是个性化交换：每个 Rank 都为每个目标 Rank 准备不同的数据。将 Rank $i$ 的输入写成：

$$X_i=[X_{i\to 0};X_{i\to 1};\dots ;X_{i\to p-1}]$$

Rank $j$ 的输出为：

$$Y_j=[X_{0\to j};X_{1\to j};\dots ;X_{p-1\to j}]$$

All-to-All 不是简单复制同一份数据，而是完成一次全局重分布。MoE Expert Parallelism 会根据 Router 结果把 token Dispatch 到专家所在 Rank，专家计算完成后再执行一次 All-to-All 把结果送回原 Rank。不同专家收到的 token 数量可能不同，因此实际实现还需要支持可变 Split Size、容量限制和负载均衡。

大模型推理中的使用场景

Collective 的选择由张量当前布局和目标布局决定。可结合并行中的模型切分方式理解。

场景	输入布局 → 输出布局	常用 Collective
TP Row-Parallel Linear	各 Rank 持有部分和 → 每个 Rank 获得完整结果	All-Reduce
TP 与 SP 组合	部分和 → 每个 Rank 保留一段序列	Reduce-Scatter
SP 进入 TP 区域	序列分片 → 每个 Rank 获得完整输入	All-Gather
EP Token Dispatch / Combine	按原 Rank 排列 → 按专家 Rank 排列，再恢复	All-to-All
Root 汇总推理输出	每个 Rank 持有输出分片 → Root 获得完整结果	Gather
参数或控制信息同步	Root 持有完整数据 → 所有 Rank 获得副本	Broadcast

Pipeline Parallelism 的相邻 Stage 传递通常使用 P2P Send/Recv，而不是 Collective。Data Parallel 推理的不同模型副本也可以独立处理请求，只有在需要同步状态或统计量时才需要集合通信。

性能与正确性

通信时间可以粗略拆成启动延迟与数据传输两部分：

$$T\approx S\alpha +V\beta$$

$S$ 是通信步骤数，$\alpha$ 是每步启动延迟，$V$ 是传输数据量，$\beta$ 是单位数据传输时间。优化 Collective 时需要同时考虑：

• 消息大小：小消息通常受 $\alpha$ 主导，大消息通常受链路带宽主导。
• 设备拓扑：NVLink、NVSwitch、PCIe、InfiniBand 和 RoCE 的带宽与层级不同，Rank 映射会影响实际路径。
• 计算重叠：使用异步操作只能表示已提交；读取输出前仍需等待操作完成，并正确处理 CUDA Stream 依赖。
• 调用一致性：通信组内所有 Rank 必须以一致的顺序参与同一 Collective，并满足对应 API 的数据类型、元素数量或 Split Size 契约，否则可能发生挂起或数据损坏。
• 张量布局：All-Gather 和 Reduce-Scatter 的输出分片通常按 Rank 排列，错误的 Rank 顺序会直接改变张量语义。

PyTorch 示例

下面的示例假设使用 torchrun 启动，每个进程绑定一张 CUDA GPU，并采用 NCCL Process Group：

import os

import torch
import torch.distributed as dist

rank = int(os.environ["RANK"])
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])

torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
dist.init_process_group(backend="nccl")

# All-Reduce：每个 Rank 最终都得到所有输入之和。
x = torch.full((4,), rank + 1, dtype=torch.float32, device=device)
dist.all_reduce(x, op=dist.ReduceOp.SUM)

# All-Gather：每个 Rank 贡献 2 个元素，输出按 Rank 顺序拼接。
local = torch.arange(2, device=device) + 2 * rank
gathered = torch.empty(world_size * local.numel(), dtype=local.dtype, device=device)
dist.all_gather_into_tensor(gathered, local)

# Reduce-Scatter：输入大小必须是输出大小的 world_size 倍。
rs_input = torch.full((world_size * 2,), rank + 1, dtype=torch.float32, device=device)
rs_output = torch.empty(2, dtype=rs_input.dtype, device=device)
dist.reduce_scatter_tensor(rs_output, rs_input, op=dist.ReduceOp.SUM)

# All-to-All：沿第 0 维等分；第 j 份发送给 Rank j。
a2a_input = torch.arange(world_size, device=device) + rank * world_size
a2a_output = torch.empty_like(a2a_input)
dist.all_to_all_single(a2a_output, a2a_input)

dist.destroy_process_group()

all_gather_into_tensor 要求各 Rank 输入等长；all_to_all_single 可以通过 input_split_sizes 和 output_split_sizes 表达不等长交换，更适合 MoE 的动态 token 分布。若设置 async_op=True，API 会返回 Work Handle，在复用输入缓冲区或读取输出前应调用 wait() 或建立等价同步关系。

参考

1. NVIDIA NCCL：Collective Operations
2. NVIDIA NCCL：Overview
3. PyTorch：Distributed Communication Package
4. PyTorch DTensor：Placement 与 Collective 映射