| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for split_group in GroupCoordinator. |
| 4 | |
| 5 | These tests verify that: |
| 6 | 1. split_group is used for both device and CPU group creation. |
| 7 | 2. Multiple subgroups work correctly with split_group. |
| 8 | 3. Both GPU and CPU all-reduce work on split groups. |
| 9 | """ |
| 10 | |
| 11 | import os |
| 12 | from typing import Any |
| 13 | |
| 14 | import multiprocess as mp |
| 15 | import pytest |
| 16 | import torch |
| 17 | import torch.distributed |
| 18 | |
| 19 | import vllm.envs as envs |
| 20 | from vllm.distributed.parallel_state import ( |
| 21 | GroupCoordinator, |
| 22 | init_distributed_environment, |
| 23 | ) |
| 24 | from vllm.utils.system_utils import update_environment_variables |
| 25 | |
| 26 | # The whole module exercises the split_group code path, which is opt-in |
| 27 | # behind VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1. |
| 28 | pytestmark = pytest.mark.skipif( |
| 29 | not envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP, |
| 30 | reason=("VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1 not set; split_group path is opt-in."), |
| 31 | ) |
| 32 | |
| 33 | mp.set_start_method("spawn", force=True) |
| 34 | |
| 35 | |
| 36 | def distributed_run(fn, world_size): |
| 37 | number_of_processes = world_size |
| 38 | processes: list[mp.Process] = [] |
| 39 | for i in range(number_of_processes): |
| 40 | env: dict[str, str] = {} |
| 41 | env["RANK"] = str(i) |
| 42 | env["LOCAL_RANK"] = str(i) |
| 43 | env["WORLD_SIZE"] = str(number_of_processes) |
| 44 | env["LOCAL_WORLD_SIZE"] = str(number_of_processes) |
| 45 | env["MASTER_ADDR"] = "localhost" |
| 46 | env["MASTER_PORT"] = "12346" |
| 47 | # propagate the opt-in flag to the spawned child workers |
| 48 | env["VLLM_DISTRIBUTED_USE_SPLIT_GROUP"] = "1" |
| 49 | p = mp.Process(target=fn, args=(env,)) |
| 50 | processes.append(p) |
| 51 | p.start() |
| 52 | |
| 53 | for p in processes: |
| 54 | p.join() |
| 55 | |
| 56 | for p in processes: |
| 57 | assert p.exitcode == 0 |
| 58 | |
| 59 | |
| 60 | def worker_fn_wrapper(fn): |
| 61 | def wrapped_fn(env): |
| 62 | update_environment_variables(env) |
| 63 | local_rank = os.environ["LOCAL_RANK"] |
| 64 | device = torch.device(f"cuda:{local_rank}") |
| 65 | torch.accelerator.set_device_index(device) |
| 66 | init_distributed_environment() |
| 67 | fn() |
| 68 | |
| 69 | return wrapped_fn |
| 70 | |
| 71 | |
| 72 | def _verify_device_group(coordinator: GroupCoordinator): |
| 73 | """Verify device group works via all-reduce.""" |
| 74 | local_rank = torch.distributed.get_rank() |
| 75 | device = torch.device(f"cuda:{local_rank}") |
| 76 | tensor = torch.ones(16, 16, dtype=torch.float32, device=device) |
| 77 | torch.distributed.all_reduce(tensor, group=coordinator.device_group) |
| 78 | torch.accelerator.synchronize() |
| 79 | expected = coordinator.world_size |
| 80 | assert torch.all(tensor == expected).cpu().item(), ( |
| 81 | f"Device group all-reduce failed: expected {expected}, " |
| 82 | f"got {tensor.flatten()[0].item()}" |
| 83 | ) |
| 84 | |
| 85 | |
| 86 | def _verify_cpu_group(coordinator: GroupCoordinator): |
| 87 | """Verify CPU group works via all-reduce.""" |
| 88 | tensor = torch.ones(16, dtype=torch.float32) |
| 89 | torch.distributed.all_reduce(tensor, group=coordinator.cpu_group) |
| 90 | expected = coordinator.world_size |
| 91 | assert torch.all(tensor == expected).cpu().item(), ( |
| 92 | f"CPU group all-reduce failed: expected {expected}, " |
| 93 | f"got {tensor.flatten()[0].item()}" |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | # --------------------------------------------------------------------------- |
| 98 | # Test 1: Basic split_group path with 2 GPUs |
| 99 | # --------------------------------------------------------------------------- |
| 100 | @worker_fn_wrapper |
| 101 | def split_group_basic_worker(): |
| 102 | rank = torch.distributed.get_rank() |
| 103 | world_size = torch.distributed.get_world_size() |
| 104 | group_ranks = [list(range(world_size))] |
| 105 | |
| 106 | coordinator = GroupCoordinator( |
| 107 | group_ranks=group_ranks, |
| 108 | local_rank=rank, |
| 109 | torch_distributed_backend="nccl", |
| 110 | use_device_communicator=False, |
| 111 | group_name="test_split_basic", |
| 112 | ) |
| 113 | |
| 114 | _verify_device_group(coordinator) |
| 115 | _verify_cpu_group(coordinator) |
| 116 | |
| 117 | |
| 118 | @pytest.mark.skipif( |
| 119 | torch.accelerator.device_count() < 2, |
| 120 | reason="Need at least 2 GPUs to run the test.", |
| 121 | ) |
| 122 | def test_split_group_basic(): |
| 123 | """Test basic GroupCoordinator creation with split_group.""" |
| 124 | distributed_run(split_group_basic_worker, 2) |
| 125 | |
| 126 | |
| 127 | # --------------------------------------------------------------------------- |
| 128 | # Test 2: Multiple subgroups with split_group (4 GPUs) |
| 129 | # --------------------------------------------------------------------------- |
| 130 | @worker_fn_wrapper |
| 131 | def split_group_multiple_subgroups_worker(): |
| 132 | rank = torch.distributed.get_rank() |
| 133 | group_ranks = [[0, 1], [2, 3]] |
| 134 | |
| 135 | coordinator = GroupCoordinator( |
| 136 | group_ranks=group_ranks, |
| 137 | local_rank=rank, |
| 138 | torch_distributed_backend="nccl", |
| 139 | use_device_communicator=False, |
| 140 | group_name="test_split_multi", |
| 141 | ) |
| 142 | |
| 143 | assert coordinator.world_size == 2 |
| 144 | |
| 145 | _verify_device_group(coordinator) |
| 146 | _verify_cpu_group(coordinator) |
| 147 | |
| 148 | if rank in [0, 1]: |
| 149 | assert coordinator.ranks == [0, 1] |
| 150 | else: |
| 151 | assert coordinator.ranks == [2, 3] |
| 152 | |
| 153 | |
| 154 | @pytest.mark.skipif( |
| 155 | torch.accelerator.device_count() < 4, |
| 156 | reason="Need at least 4 GPUs to run the test.", |
| 157 | ) |
| 158 | def test_split_group_multiple_subgroups(): |
| 159 | """Test GroupCoordinator with multiple independent subgroups.""" |
| 160 | distributed_run(split_group_multiple_subgroups_worker, 4) |
| 161 | |
| 162 | |
| 163 | # --------------------------------------------------------------------------- |
| 164 | # Test 3: split_group contract — every parent rank must enter with the same |
| 165 | # ``split_ranks``. NCCL happens to produce |
| 166 | # correct subgroups for disjoint partitions because the wrapper hashes |
| 167 | # ``my_group`` to derive the comm-split color, but the contract violation is |
| 168 | # real and would break under non-partition / non-NCCL backends. This test |
| 169 | # captures the actual ``split_ranks`` argument passed on every rank and |
| 170 | # asserts they match. |
| 171 | # --------------------------------------------------------------------------- |
| 172 | @worker_fn_wrapper |
| 173 | def split_group_contract_worker(): |
| 174 | rank = torch.distributed.get_rank() |
| 175 | group_ranks = [[0, 1], [2, 3]] |
| 176 | |
| 177 | captured: list[list[list[int]]] = [] |
| 178 | original_split_group = torch.distributed.split_group |
| 179 | |
| 180 | def capturing_split_group(*args, split_ranks=None, **kwargs): |
| 181 | captured.append([list(g) for g in split_ranks]) |
| 182 | return original_split_group(*args, split_ranks=split_ranks, **kwargs) |
| 183 | |
| 184 | torch.distributed.split_group = capturing_split_group |
| 185 | try: |
| 186 | GroupCoordinator( |
| 187 | group_ranks=group_ranks, |
| 188 | local_rank=rank, |
| 189 | torch_distributed_backend="nccl", |
| 190 | use_device_communicator=False, |
| 191 | group_name="test_split_contract", |
| 192 | ) |
| 193 | finally: |
| 194 | torch.distributed.split_group = original_split_group |
| 195 | |
| 196 | # GroupCoordinator builds two subgroups (device + cpu) per coordinator, |
| 197 | # so every rank must have made exactly two split_group calls. |
| 198 | if len(captured) != 2: |
| 199 | raise AssertionError( |
| 200 | f"rank {rank} expected 2 split_group calls (device + cpu), " |
| 201 | f"got {len(captured)}: {captured}" |
| 202 | ) |
| 203 | |
| 204 | world_size = torch.distributed.get_world_size() |
| 205 | for call_idx in range(2): |
| 206 | gathered: list[Any] = [None] * world_size |
| 207 | torch.distributed.all_gather_object(gathered, captured[call_idx]) |
| 208 | # Normalize for stable comparison (sort each subgroup and the outer list). |
| 209 | norm = [ |
| 210 | sorted([sorted(sg) for sg in per_rank_args]) for per_rank_args in gathered |
| 211 | ] |
| 212 | reference = norm[0] |
| 213 | for r, args in enumerate(norm): |
| 214 | if args != reference: |
| 215 | raise AssertionError( |
| 216 | f"split_group contract violation on call #{call_idx}: " |
| 217 | f"rank {r} passed split_ranks={gathered[r]}, but rank 0 " |
| 218 | f"passed split_ranks={gathered[0]}. PyTorch requires every " |
| 219 | "parent rank to enter split_group with the same split_ranks." |
| 220 | ) |
| 221 | |
| 222 | |
| 223 | @pytest.mark.skipif( |
| 224 | torch.accelerator.device_count() < 4, |
| 225 | reason="Need at least 4 GPUs to run the test.", |
| 226 | ) |
| 227 | def test_split_group_contract_same_split_ranks_on_all_ranks(): |
| 228 | """All parent ranks must call torch.distributed.split_group with the same |
| 229 | ``split_ranks`` argument. This catches the bug where each rank passed |
| 230 | only its own subgroup (``split_ranks=[ranks]``), which NCCL forgives for |
| 231 | disjoint partitions but is a documented contract violation. |
| 232 | """ |
| 233 | distributed_run(split_group_contract_worker, 4) |