chore: import upstream snapshot with attribution

tests/distributed/test_split_group.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for split_group in GroupCoordinator.
+
+These tests verify that:
+1. split_group is used for both device and CPU group creation.
+2. Multiple subgroups work correctly with split_group.
+3. Both GPU and CPU all-reduce work on split groups.
+"""
+
+import os
+from typing import Any
+
+import multiprocess as mp
+import pytest
+import torch
+import torch.distributed
+
+import vllm.envs as envs
+from vllm.distributed.parallel_state import (
+    GroupCoordinator,
+    init_distributed_environment,
+)
+from vllm.utils.system_utils import update_environment_variables
+
+# The whole module exercises the split_group code path, which is opt-in
+# behind VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1.
+pytestmark = pytest.mark.skipif(
+    not envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP,
+    reason=("VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1 not set; split_group path is opt-in."),
+)
+
+mp.set_start_method("spawn", force=True)
+
+
+def distributed_run(fn, world_size):
+    number_of_processes = world_size
+    processes: list[mp.Process] = []
+    for i in range(number_of_processes):
+        env: dict[str, str] = {}
+        env["RANK"] = str(i)
+        env["LOCAL_RANK"] = str(i)
+        env["WORLD_SIZE"] = str(number_of_processes)
+        env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
+        env["MASTER_ADDR"] = "localhost"
+        env["MASTER_PORT"] = "12346"
+        # propagate the opt-in flag to the spawned child workers
+        env["VLLM_DISTRIBUTED_USE_SPLIT_GROUP"] = "1"
+        p = mp.Process(target=fn, args=(env,))
+        processes.append(p)
+        p.start()
+
+    for p in processes:
+        p.join()
+
+    for p in processes:
+        assert p.exitcode == 0
+
+
+def worker_fn_wrapper(fn):
+    def wrapped_fn(env):
+        update_environment_variables(env)
+        local_rank = os.environ["LOCAL_RANK"]
+        device = torch.device(f"cuda:{local_rank}")
+        torch.accelerator.set_device_index(device)
+        init_distributed_environment()
+        fn()
+
+    return wrapped_fn
+
+
+def _verify_device_group(coordinator: GroupCoordinator):
+    """Verify device group works via all-reduce."""
+    local_rank = torch.distributed.get_rank()
+    device = torch.device(f"cuda:{local_rank}")
+    tensor = torch.ones(16, 16, dtype=torch.float32, device=device)
+    torch.distributed.all_reduce(tensor, group=coordinator.device_group)
+    torch.accelerator.synchronize()
+    expected = coordinator.world_size
+    assert torch.all(tensor == expected).cpu().item(), (
+        f"Device group all-reduce failed: expected {expected}, "
+        f"got {tensor.flatten()[0].item()}"
+    )
+
+
+def _verify_cpu_group(coordinator: GroupCoordinator):
+    """Verify CPU group works via all-reduce."""
+    tensor = torch.ones(16, dtype=torch.float32)
+    torch.distributed.all_reduce(tensor, group=coordinator.cpu_group)
+    expected = coordinator.world_size
+    assert torch.all(tensor == expected).cpu().item(), (
+        f"CPU group all-reduce failed: expected {expected}, "
+        f"got {tensor.flatten()[0].item()}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Test 1: Basic split_group path with 2 GPUs
+# ---------------------------------------------------------------------------
+@worker_fn_wrapper
+def split_group_basic_worker():
+    rank = torch.distributed.get_rank()
+    world_size = torch.distributed.get_world_size()
+    group_ranks = [list(range(world_size))]
+
+    coordinator = GroupCoordinator(
+        group_ranks=group_ranks,
+        local_rank=rank,
+        torch_distributed_backend="nccl",
+        use_device_communicator=False,
+        group_name="test_split_basic",
+    )
+
+    _verify_device_group(coordinator)
+    _verify_cpu_group(coordinator)
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 2,
+    reason="Need at least 2 GPUs to run the test.",
+)
+def test_split_group_basic():
+    """Test basic GroupCoordinator creation with split_group."""
+    distributed_run(split_group_basic_worker, 2)
+
+
+# ---------------------------------------------------------------------------
+# Test 2: Multiple subgroups with split_group (4 GPUs)
+# ---------------------------------------------------------------------------
+@worker_fn_wrapper
+def split_group_multiple_subgroups_worker():
+    rank = torch.distributed.get_rank()
+    group_ranks = [[0, 1], [2, 3]]
+
+    coordinator = GroupCoordinator(
+        group_ranks=group_ranks,
+        local_rank=rank,
+        torch_distributed_backend="nccl",
+        use_device_communicator=False,
+        group_name="test_split_multi",
+    )
+
+    assert coordinator.world_size == 2
+
+    _verify_device_group(coordinator)
+    _verify_cpu_group(coordinator)
+
+    if rank in [0, 1]:
+        assert coordinator.ranks == [0, 1]
+    else:
+        assert coordinator.ranks == [2, 3]
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 4,
+    reason="Need at least 4 GPUs to run the test.",
+)
+def test_split_group_multiple_subgroups():
+    """Test GroupCoordinator with multiple independent subgroups."""
+    distributed_run(split_group_multiple_subgroups_worker, 4)
+
+
+# ---------------------------------------------------------------------------
+# Test 3: split_group contract — every parent rank must enter with the same
+# ``split_ranks``. NCCL happens to produce
+# correct subgroups for disjoint partitions because the wrapper hashes
+# ``my_group`` to derive the comm-split color, but the contract violation is
+# real and would break under non-partition / non-NCCL backends. This test
+# captures the actual ``split_ranks`` argument passed on every rank and
+# asserts they match.
+# ---------------------------------------------------------------------------
+@worker_fn_wrapper
+def split_group_contract_worker():
+    rank = torch.distributed.get_rank()
+    group_ranks = [[0, 1], [2, 3]]
+
+    captured: list[list[list[int]]] = []
+    original_split_group = torch.distributed.split_group
+
+    def capturing_split_group(*args, split_ranks=None, **kwargs):
+        captured.append([list(g) for g in split_ranks])
+        return original_split_group(*args, split_ranks=split_ranks, **kwargs)
+
+    torch.distributed.split_group = capturing_split_group
+    try:
+        GroupCoordinator(
+            group_ranks=group_ranks,
+            local_rank=rank,
+            torch_distributed_backend="nccl",
+            use_device_communicator=False,
+            group_name="test_split_contract",
+        )
+    finally:
+        torch.distributed.split_group = original_split_group
+
+    # GroupCoordinator builds two subgroups (device + cpu) per coordinator,
+    # so every rank must have made exactly two split_group calls.
+    if len(captured) != 2:
+        raise AssertionError(
+            f"rank {rank} expected 2 split_group calls (device + cpu), "
+            f"got {len(captured)}: {captured}"
+        )
+
+    world_size = torch.distributed.get_world_size()
+    for call_idx in range(2):
+        gathered: list[Any] = [None] * world_size
+        torch.distributed.all_gather_object(gathered, captured[call_idx])
+        # Normalize for stable comparison (sort each subgroup and the outer list).
+        norm = [
+            sorted([sorted(sg) for sg in per_rank_args]) for per_rank_args in gathered
+        ]
+        reference = norm[0]
+        for r, args in enumerate(norm):
+            if args != reference:
+                raise AssertionError(
+                    f"split_group contract violation on call #{call_idx}: "
+                    f"rank {r} passed split_ranks={gathered[r]}, but rank 0 "
+                    f"passed split_ranks={gathered[0]}. PyTorch requires every "
+                    "parent rank to enter split_group with the same split_ranks."
+                )
+
+
+@pytest.mark.skipif(
+    torch.accelerator.device_count() < 4,
+    reason="Need at least 4 GPUs to run the test.",
+)
+def test_split_group_contract_same_split_ranks_on_all_ranks():
+    """All parent ranks must call torch.distributed.split_group with the same
+    ``split_ranks`` argument. This catches the bug where each rank passed
+    only its own subgroup (``split_ranks=[ranks]``), which NCCL forgives for
+    disjoint partitions but is a documented contract violation.
+    """
+    distributed_run(split_group_contract_worker, 4)