de7d70676d
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: hjjq <50634613+hjjq@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Tests for cpu_all_reduce_ (gloo, CPU only)."""
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.multiprocessing as mp
|
|
|
|
from vllm.distributed.parallel_state import GroupCoordinator
|
|
from vllm.distributed.utils import (
|
|
cpu_all_reduce_,
|
|
stateless_init_torch_distributed_process_group,
|
|
)
|
|
from vllm.utils.network_utils import get_open_port
|
|
|
|
pytestmark = pytest.mark.cpu_test
|
|
|
|
|
|
def _make_input(rank: int, world_size: int, step: int) -> torch.Tensor:
|
|
# Same shape as the V2 DP metadata tensor; values differ per rank/step.
|
|
t = torch.zeros(6, world_size, dtype=torch.int32)
|
|
t[:, rank] = torch.arange(6, dtype=torch.int32) * (rank + 1) + step
|
|
return t
|
|
|
|
|
|
def _worker(rank: int, world_size: int, port: int, stateless_port: int):
|
|
dist.init_process_group(
|
|
"gloo",
|
|
init_method=f"tcp://127.0.0.1:{port}",
|
|
rank=rank,
|
|
world_size=world_size,
|
|
)
|
|
coord = GroupCoordinator(
|
|
group_ranks=[list(range(world_size))],
|
|
local_rank=0,
|
|
torch_distributed_backend="gloo",
|
|
use_device_communicator=False,
|
|
group_name="cpu_all_reduce_test",
|
|
)
|
|
|
|
# Several steps to exercise tag reuse across consecutive calls.
|
|
for step in range(5):
|
|
expected = _make_input(rank, world_size, step)
|
|
dist.all_reduce(expected, group=coord.cpu_group)
|
|
t = _make_input(rank, world_size, step)
|
|
out = coord.cpu_all_reduce_(t)
|
|
assert out is t
|
|
torch.testing.assert_close(t, expected, rtol=0, atol=0)
|
|
|
|
# Float sums must be bitwise identical on every rank.
|
|
t = torch.randn(6, world_size, generator=torch.Generator().manual_seed(rank))
|
|
expected = t.clone()
|
|
dist.all_reduce(expected, group=coord.cpu_group)
|
|
coord.cpu_all_reduce_(t)
|
|
torch.testing.assert_close(t, expected)
|
|
gathered = [torch.empty_like(t) for _ in range(world_size)]
|
|
dist.all_gather(gathered, t, group=coord.cpu_group)
|
|
for other in gathered:
|
|
assert torch.equal(other, t)
|
|
|
|
# Unregistered (stateless) process group, as used by the engine-core DP
|
|
# group and elastic EP.
|
|
stateless_pg = stateless_init_torch_distributed_process_group(
|
|
"127.0.0.1", stateless_port, rank, world_size, backend="gloo"
|
|
)
|
|
expected = _make_input(rank, world_size, 0)
|
|
dist.all_reduce(expected)
|
|
t = _make_input(rank, world_size, 0)
|
|
cpu_all_reduce_(t, stateless_pg)
|
|
torch.testing.assert_close(t, expected, rtol=0, atol=0)
|
|
|
|
dist.destroy_process_group()
|
|
|
|
|
|
# Non-power-of-two sizes exercise the cross-block exchange.
|
|
@pytest.mark.parametrize("world_size", [2, 3, 4, 6, 8, 9, 12])
|
|
def test_cpu_all_reduce(world_size: int):
|
|
mp.spawn(
|
|
_worker,
|
|
args=(world_size, get_open_port(), get_open_port()),
|
|
nprocs=world_size,
|
|
join=True,
|
|
)
|