parallel_state and GroupCoordinator: TP/PP/DP communication groups

源码版本v0.25.1

Responsibilities

vLLM runs distributed inference on top of PyTorch torch.distributed ProcessGroups, but the raw ProcessGroup API is verbose and inconsistent: CPU and device communication live in separate groups, custom all-reduce has to go through a custom op, and broadcasting objects versus broadcasting tensors use two different APIs. GroupCoordinator (parallel_state.py:358) is vLLM's wrapper layer on top of ProcessGroup. It bundles one process group's CPU group, device group, device_communicator, and mq_broadcaster into a single object and exposes a unified all_reduce / all_gather / broadcast / send_object / recv_tensor_dict / barrier interface.

The parallel_state module maintains several global singletons: _TP (tensor parallel), _PP (pipeline parallel), _DP (data parallel), _EP (expert parallel), _DCP (decode context parallel), _PCP (prefill context parallel), and _EPLB (expert load balancing). Each group is read through get_tp_group / get_pp_group / get_dp_group and friends (parallel_state.py:1368-1424), and the model layers ColumnParallelLinear / RowParallelLinear rely on these getters to obtain their communication groups. Groups are initialized once at startup by initialize_model_parallel (parallel_state.py:1713), which splits the global rank along tensor_model_parallel_size * pipeline_model_parallel_size * ... and calls new_group for each segment.

GroupCoordinator also registers all_reduce / all_gather as torch.ops.vllm.* custom ops, so that during Dynamo compilation the communication ops can be fused into the graph as ordinary ops (parallel_state.py:641-663), rather than forcing self and other Python objects into the compiled graph. When the world size is 1, all communication ops return their input directly, bypassing NCCL entirely.

Design motivation

• Two ProcessGroups per group: cpu_group (gloo backend) and device_group (nccl/gloo/platform backend) are split apart (parallel_state.py:435-452). Object broadcasts and coordinator-to-coordinator coordination happen on CPU, while tensor traffic happens on device.
• Three rank concepts kept distinct: rank (global), local_rank (per machine), and rank_in_group (within the group) (parallel_state.py:369-380). In multi-node setups these differ, and the model layer must use rank_in_group to slice correctly.
• Fixed rank layout: ExternalDP × DP × PP × PCP × TP (parallel_state.py:1779-1794). After torch.arange(world_size).reshape(...), unbinding along a dimension yields the rank list for the corresponding group.
• Single-GPU bypass: all_reduce and friends return early when world_size == 1 (parallel_state.py:657-658), so a single GPU does not need to spin up NCCL.
• Custom ops for Dynamo compatibility: when use_custom_op_call=True, calls go through torch.ops.vllm.all_reduce(input_, group_name=self.unique_name) (parallel_state.py:660-661). The group is referenced by string name so Dynamo never has to trace through a Python object.
• Pluggable DeviceCommunicator: the device_communicator field lets CUDA / XPU / CPU / custom all-reduce each provide their own implementation (parallel_state.py:463-478). is_cuda_aliked binds cuda:N, XPU binds xpu:N.
• Message-queue broadcaster for TP: init_model_parallel_group opens use_message_queue_broadcaster=True on the TP group (parallel_state.py:1805-1811), so workers can share tensor dicts without going through NCCL.
• StatelessGroupCoordinator fallback: in elastic EP scenarios where torch.distributed is not initialized, _init_stateless_group (parallel_state.py:1316-1338) stands up an equivalent group using a TCP store plus StatelessProcessGroup.

Key files

• GroupCoordinator class:358 — process group wrapper holding rank/ranks/world_size/local_rank/rank_in_group/cpu_group/device_group.
• GroupCoordinator.__init__:387-462 — picks split_group vs new_group based on VLLM_DISTRIBUTED_USE_SPLIT_GROUP, binds device.
• GroupCoordinator.all_reduce:641 — routes through torch.ops.vllm.all_reduce or _all_reduce_out_place, bypasses when world_size=1.
• GroupCoordinator.all_gather:670 — same custom-op path, with a dim check and a bypass.
• broadcast:745 — broadcasts a tensor on the device group, with a src argument.
• broadcast_object:760 — uses broadcast_object_list on the cpu_group to sync configs and sampler generator state.
• broadcast_tensor_dict:864 — broadcasts an entire dict, used at worker startup to sync KV cache metadata.
• barrier:1179 — device + cpu double barrier, aligning GPU operations with CPU coordination.
• initialize_model_parallel:1713 — splits ranks along ExternalDP × DP × PP × PCP × TP and initializes all groups.
• TP group construction:1796-1811 — group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0).
• init_model_parallel_group:1298 — factory that forwards to GroupCoordinator(...).
• get_tp_group:1368 — global getter used by the model layer to fetch the TP group.
• get_tensor_model_parallel_world_size:2031 — convenience accessor, internally get_tp_group().world_size.
• _register_group:126 — registers each GroupCoordinator in a global dict so custom ops can look it up by group_name.

Data flow

At startup, Worker calls initialize_model_parallel(tp_size, pp_size, ...) in its own process. This function reshapes the global rank in ExternalDP × DP × PP × PCP × TP order, then unbinds along each dimension to get the rank list for every group. The TP group construction looks like this:

# 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",
)

When the model layer invokes communication, it does not touch ProcessGroup directly. It calls get_tp_group().all_reduce(tensor), and all_reduce internally decides — based on use_custom_op_call — whether to go through the custom op or call device_communicator.all_reduce directly:

# 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's forward computes the local shard and then calls all_reduce to combine the shards; RowParallelLinear does the reverse — all-gather first, compute local, then reduce_scatter. All of these ops are unified at the GroupCoordinator layer.

Boundaries and failure modes

• rank_in_group must be the in-group rank: the model layer must slice with rank_in_group, not the global rank, otherwise multi-node setups will slice the wrong shard (parallel_state.py:380).
• TP double initialization fails the assert: initialize_model_parallel checks assert _TP is None for each group (parallel_state.py:1798). Calling it twice raises AssertionError; you need destroy_model_parallel first.
• DCP cannot exceed TP: dcp_size <= tp_size, because DCP reuses the TP group's GPUs (parallel_state.py:1816-1820).
• Single GPU takes the bypass: all_reduce returns the input directly when world_size == 1 (parallel_state.py:657-658), so no NCCL process is needed. But if device_communicator is None and the call is not in the bypass path, it raises ValueError.
• Missing DeviceCommunicator raises: _all_reduce_out_place checks device_communicator is None and raises (parallel_state.py:665-668), preventing calls before the backend is initialized.
• _replace_active_groups must be called collectively: parallel_state.py:1341-1362 requires every rank to call it simultaneously, otherwise it deadlocks. During elastic EP scale-up/down it is orchestrated by the supervisor.

Summary

parallel_state is the foundation of vLLM's multi-GPU inference: GroupCoordinator wraps PyTorch ProcessGroup and bundles the CPU/device dual groups, custom ops, device_communicator, and message-queue broadcaster into one place. The model layer fetches the relevant group via get_tp_group() / get_pp_group() and then calls all_reduce and other methods. For how TP linear layers use these communications, see /model-loading/tp-layers; for how workers start these groups, see /worker/worker; for the entry point where the engine core ties these workers together, see /engine/engine-core.

See official docs: vLLM docs · README