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