parallel_state と GroupCoordinator:TP/PP/DP 通信グループ

源码版本v0.25.1

役割

vLLM が分散推論を行うには PyTorch torch.distributed の ProcessGroup に依存しますが、生の ProcessGroup は扱いにくく統一も難しいです:CPU 通信と device 通信でグループを分け、カスタム all-reduce は custom op 経由、オブジェクトのブロードキャストとテンサーのブロードキャストは別 API です。GroupCoordinator(parallel_state.py:358)は vLLM が ProcessGroup の上に被せたラッパーで、あるプロセスグループの cpu group、device group、device_communicator、mq_broadcaster を 1 つのオブジェクトにまとめ、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(エキスパート負荷分散)。各グループは 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.* の custom op として登録し、Dynamo コンパイル時に通信 op を通常の op としてグラフに融合できるようにします(parallel_state.py:641-663)。self のような Python オブジェクトをコンパイルグラフに押し込むのを避けるためです。ワールドサイズ (world size) が 1 のときはすべての通信 op が入力をそのまま return し、NCCL をバイパスします。

設計動機

• 1 グループでデュアル PG:cpu_group(gloo バックエンド)と device_group(nccl/gloo/プラットフォームバックエンド)を分け(parallel_state.py:435-452)、CPU でオブジェクトブロードキャストや coordinator 間の調整、device でテンサー通信を行います。
• rank の 3 つの概念を区別:rank(グローバル)、local_rank(ホスト内)、rank_in_group(グループ内)(parallel_state.py:369-380)。マルチノードではこの 3 つは異なり、モデル層は rank_in_group を使わないと正しくスライスできません。
• rank layout は固定順序:ExternalDP × DP × PP × PCP × TP(parallel_state.py:1779-1794)で、torch.arange(world_size).reshape(...) のあとある次元で unbind すると対応グループの rank list が得られます。
• 単カードはバイパス: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)、グループを文字列で検索し、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 で等価なグループを立ち上げます。

主要ファイル

• 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 の 2 経路を選び、device をバインドします。
• GroupCoordinator.all_reduce:641 — torch.ops.vllm.all_reduce または _all_reduce_out_place、world_size=1 でバイパス。
• GroupCoordinator.all_gather:670 — 同様に custom op、dim 検証、バイパス。
• 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 を切り、全グループを初期化します。
• TP グループ構築: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 グループの構築です:

# 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 ではなく rank_in_group を使い、さもなければマルチノードで誤ったスライスを切ります(parallel_state.py:380)。
• TP の重複初期化は assertion 失敗:initialize_model_parallel は各グループで 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 はバイパス:all_reduce は world_size == 1 のとき直接 input を return し(parallel_state.py:657-658)、NCCL プロセスは不要です。ただし device_communicator is None かつ非バイパスのシナリオは 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 ブロードキャスタを 1 つにまとめ、モデル層は get_tp_group() / get_pp_group() で対応グループを取得して all_reduce などを呼びます。TP 線形層がこれらの通信をどう使うかは /model-loading/tp-layers、Worker がこれらの group をどう起動するかは /worker/worker、エンジンコアがこれらの worker をどう繋ぐかは /engine/engine-core を参照してください。

公式資料: vLLM 文档 · README