parallel_state 与 GroupCoordinator:TP/PP/DP 通信组

源码版本v0.25.1

职责

vLLM 跑分布式推理靠的是 PyTorch torch.distributed 的 ProcessGroup,但 raw ProcessGroup 用起来又啰嗦又难统一:CPU 通信和 device 通信要分两个 group,自定义 all-reduce 要走 custom op,广播对象和广播张量是两套 API。GroupCoordinator(parallel_state.py:358)是 vLLM 在 ProcessGroup 之上加的一层封装,把一组进程的 CPU group、device group、device_communicator、mq_broadcaster 全装到一个对象里,提供 all_reduce / all_gather / broadcast / send_object / recv_tensor_dict / barrier 这一套统一接口。

parallel_state 模块维护几组全局单例:_TP(张量并行)、_PP(流水线并行)、_DP(数据并行)、_EP(专家并行)、_DCP(decode context parallel)、_PCP(prefill context parallel)、_EPLB(专家负载均衡)。每个 group 用 get_tp_group / get_pp_group / get_dp_group 等(parallel_state.py:1368-1424)读取,模型层 ColumnParallelLinear / RowParallelLinear 就靠这些 getter 拿通信组。组的初始化由 initialize_model_parallel(parallel_state.py:1713)在启动时调一次,按 tensor_model_parallel_size * pipeline_model_parallel_size * ... 把全局 rank 切成多段并 new_group。

GroupCoordinator 还做了一件事:把 all_reduce / all_gather 注册成 torch.ops.vllm.* 自定义 op,这样 Dynamo 编译时能把通信 op 当作普通 op 融合到 graph 里(parallel_state.py:641-663),而不是把 self 这种 Python 对象塞进编译图。世界大小 (world size) 为 1 时所有通信 op 直接 return input,绕过 NCCL。

设计动机

• 一组双 PG:cpu_group(gloo 后端)和 device_group(nccl/gloo/平台后端)分开(parallel_state.py:435-452),CPU 上做对象广播、coordinator 之间协调,device 上做张量通信。
• rank 三种概念分清:rank(全局)、local_rank(本机内)、rank_in_group(组内)(parallel_state.py:369-380)——多节点场景这三个不一样,模型层要用 rank_in_group 才能正确切片。
• rank layout 固定顺序:ExternalDP × DP × PP × PCP × TP(parallel_state.py:1779-1794),torch.arange(world_size).reshape(...) 后沿某维 unbind 就拿到对应组的 rank list。
• bypass 单卡:all_reduce 等方法在 world_size == 1 时直接 return(parallel_state.py:657-658),避免单 GPU 也要启动 NCCL。
• custom op 适配 Dynamo:use_custom_op_call=True 时走 torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)(parallel_state.py:660-661),把 group 当字符串查表,避免 Dynamo 看到 Python 对象就报错。
• DeviceCommunicator 可插拔:device_communicator 字段允许 CUDA / XPU / CPU / 自定义 all-reduce 各自实现(parallel_state.py:463-478),is_cuda_aliked 时绑 cuda:N、XPU 绑 xpu:N。
• TP 用消息队列广播:init_model_parallel_group 在 TP 组上打开 use_message_queue_broadcaster=True(parallel_state.py:1805-1811),让 worker 之间能共享张量字典而不走 NCCL。
• StatelessGroupCoordinator 兜底:在没初始化 torch.distributed 的弹性 EP 场景,_init_stateless_group(parallel_state.py:1316-1338)用 TCP store + StatelessProcessGroup 起一套等价 group。

关键文件

• GroupCoordinator 类:358 — 进程组封装,持 rank/ranks/world_size/local_rank/rank_in_group/cpu_group/device_group。
• GroupCoordinator.__init__:387-462 — 根据 VLLM_DISTRIBUTED_USE_SPLIT_GROUP 选 split_group 或 new_group 两条路,绑 device。
• GroupCoordinator.all_reduce:641 — 走 torch.ops.vllm.all_reduce 或 _all_reduce_out_place,world_size=1 时 bypass。
• GroupCoordinator.all_gather:670 — 同样走 custom op,dim 校验,bypass。
• broadcast:745 — device group 上广播张量,带 src 参数。
• broadcast_object:760 — 在 cpu_group 上用 broadcast_object_list,用来同步配置、采样 generator state。
• broadcast_tensor_dict:864 — 把 dict 整体广播,worker 启动时同步 KV cache 元数据。
• barrier:1179 — device + cpu 双 barrier,确保 GPU 操作和 CPU 协调都对齐。
• initialize_model_parallel:1713 — 按 ExternalDP × DP × PP × PCP × TP 切 rank,初始化所有 group。
• TP group 构建:1796-1811 — group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)。
• init_model_parallel_group:1298 — 工厂方法,转调 GroupCoordinator(...)。
• get_tp_group:1368 — 全局 getter,模型层调它拿 TP 组。
• get_tensor_model_parallel_world_size:2031 — 便捷查询,内部就是 get_tp_group().world_size。
• _register_group:126 — 把每个 GroupCoordinator 注册到全局 dict,供 custom op 用 group_name 反查。

数据流

启动时 Worker 在自己的进程里调 initialize_model_parallel(tp_size, pp_size, ...),这个函数把全局 rank 按 ExternalDP × DP × PP × PCP × TP 的顺序 reshape,然后沿每一维 unbind 得到每组的 rank list。下面是 TP group 的构造:

# vllm/distributed/parallel_state.py L1788-L1811
all_ranks = torch.arange(world_size).reshape(
    -1,
    data_parallel_size,
    pipeline_model_parallel_size,
    prefill_context_model_parallel_size,
    tensor_model_parallel_size,
)  # noqa

# Build the tensor model-parallel groups.
global _TP
assert _TP is None, "tensor model parallel group is already initialized"
group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
group_ranks = [x.tolist() for x in group_ranks]
if enable_elastic_ep:
    group_ranks = local_all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
    group_ranks = [x.tolist() for x in group_ranks]
# message queue broadcaster is only used in tensor model parallel group
_TP = init_model_parallel_group(
    group_ranks,
    get_world_group().local_rank,
    backend,
    use_message_queue_broadcaster=True,
    group_name="tp",
)

模型层调通信时不直接碰 ProcessGroup,而是 get_tp_group().all_reduce(tensor),all_reduce 内部根据 use_custom_op_call 决定走 custom op 还是直接调 device_communicator.all_reduce:

# vllm/distributed/parallel_state.py L641-L668
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
    """
    User-facing all-reduce function before we actually call the
    all-reduce operation.
    """
    # Bypass the function if we are using only 1 GPU.
    if self.world_size == 1:
        return input_

    if self.use_custom_op_call:
        return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
    else:
        return self._all_reduce_out_place(input_)

def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:
    if self.device_communicator is None:
        raise ValueError("No device communicator found")
    return self.device_communicator.all_reduce(input_)

ColumnParallelLinear 的 forward 在前向算完 local 部分后调 all_reduce 把分片合到一起;RowParallelLinear 反过来,先 all-gather 再算 local,最后 reduce_scatter。这些 op 都在 GroupCoordinator 这层被统一封装。

边界与失败

• rank_in_group 必须用组内 rank:模型层切片要用 rank_in_group 不是全局 rank,否则在多节点场景会切到错误的分片(parallel_state.py:380)。
• TP 重复初始化会断言失败:initialize_model_parallel 对每个 group 检查 assert _TP is None(parallel_state.py:1798),重复调会 AssertionError,需要先 destroy_model_parallel。
• DCP 不能超 TP:dcp_size <= tp_size,因为 DCP 复用 TP 组的 GPU(parallel_state.py:1816-1820)。
• 单 GPU 走 bypass:all_reduce 在 world_size == 1 时直接 return input(parallel_state.py:657-658),不需要 NCCL 进程;但 device_communicator is None 且非 bypass 场景会抛 ValueError。
• DeviceCommunicator 缺失报错:_all_reduce_out_place 检查 device_communicator is None 抛错(parallel_state.py:665-668),防止后端没初始化就被调用。
• _replace_active_groups 必须集体调用:parallel_state.py:1341-1362 要求所有 rank 同时调,否则会死锁,弹性 EP 扩缩容时由 supervisor 统一调度。

小结

parallel_state 是 vLLM 多卡推理的底座:GroupCoordinator 封装 PyTorch ProcessGroup,把 CPU/device 双 group、custom op、device_communicator、message queue 广播器全装一起,模型层用 get_tp_group() / get_pp_group() 拿到对应组再调 all_reduce 等方法。TP 线性层怎么用这些通信,见 /model-loading/tp-layers;Worker 怎么启动这些 group,见 /worker/worker;引擎核心把这些 worker 串起来的入口见 /engine/engine-core。

对照官方资料:vLLM 文档 · README。