chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,204 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import os
|
||||
import random
|
||||
|
||||
import msgspec
|
||||
import msgspec.msgpack
|
||||
import pytest
|
||||
import zmq
|
||||
|
||||
from vllm.config.kv_events import KVEventsConfig
|
||||
from vllm.distributed.kv_events import EventPublisherFactory
|
||||
|
||||
from .test_events import SampleBatch
|
||||
|
||||
DP_RANK = 0
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def random_port():
|
||||
"""Generate a random port number for testing"""
|
||||
return random.randint(10000, 59900)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def publisher_config(random_port, request):
|
||||
"""Create a publisher config with inproc transport"""
|
||||
how = request.param if hasattr(request, "param") else "inproc"
|
||||
|
||||
if how == "inproc":
|
||||
endpoint = f"inproc://test-{random_port}"
|
||||
replay_endpoint = endpoint + "-replay"
|
||||
else:
|
||||
endpoint = f"tcp://*:{random_port}"
|
||||
replay_endpoint = f"tcp://*:{random_port + 100}"
|
||||
|
||||
return KVEventsConfig(
|
||||
enable_kv_cache_events=True,
|
||||
publisher="zmq",
|
||||
endpoint=endpoint,
|
||||
replay_endpoint=replay_endpoint,
|
||||
buffer_steps=100,
|
||||
hwm=1000,
|
||||
topic="test",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def publisher(publisher_config):
|
||||
"""Create and return a publisher instance"""
|
||||
pub = EventPublisherFactory.create(publisher_config, DP_RANK)
|
||||
yield pub
|
||||
pub.shutdown()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def subscriber(publisher_config):
|
||||
"""Create and return a subscriber for testing"""
|
||||
endpoint = publisher_config.endpoint
|
||||
replay_endpoint = publisher_config.replay_endpoint
|
||||
|
||||
if endpoint.startswith("tcp://*"):
|
||||
endpoint = endpoint.replace("*", "127.0.0.1")
|
||||
if replay_endpoint and replay_endpoint.startswith("tcp://*"):
|
||||
replay_endpoint = replay_endpoint.replace("*", "127.0.0.1")
|
||||
|
||||
sub = MockSubscriber(
|
||||
[endpoint],
|
||||
[replay_endpoint] if replay_endpoint else None,
|
||||
publisher_config.topic,
|
||||
)
|
||||
yield sub
|
||||
sub.close()
|
||||
|
||||
|
||||
class MockSubscriber:
|
||||
"""Helper class to receive and verify published events"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pub_endpoints: str | list[str],
|
||||
replay_endpoints: str | list[str] | None = None,
|
||||
topic: str = "",
|
||||
decode_type=SampleBatch,
|
||||
):
|
||||
self.ctx = zmq.Context.instance()
|
||||
|
||||
# Convert single endpoint to list for consistency
|
||||
if isinstance(pub_endpoints, str):
|
||||
pub_endpoints = [pub_endpoints]
|
||||
if isinstance(replay_endpoints, str):
|
||||
replay_endpoints = [replay_endpoints]
|
||||
|
||||
# Set up subscriber socket - connect to all endpoints
|
||||
self.sub = self.ctx.socket(zmq.SUB)
|
||||
self.sub.setsockopt(zmq.SUBSCRIBE, topic.encode("utf-8"))
|
||||
for endpoint in pub_endpoints:
|
||||
self.sub.connect(endpoint)
|
||||
|
||||
# Set up replay sockets if provided.
|
||||
# DEALER allows receiving multiple replies per request.
|
||||
self.replay_sockets = []
|
||||
if replay_endpoints:
|
||||
for replay_endpoint in replay_endpoints:
|
||||
replay = self.ctx.socket(zmq.DEALER)
|
||||
replay.connect(replay_endpoint)
|
||||
self.replay_sockets.append(replay)
|
||||
|
||||
self.topic = topic
|
||||
self.topic_bytes = topic.encode("utf-8")
|
||||
self.received_msgs: list[tuple[int, SampleBatch]] = []
|
||||
self.last_seq = -1
|
||||
self.decoder = msgspec.msgpack.Decoder(type=decode_type)
|
||||
|
||||
def receive_one(self, timeout=1000) -> tuple[int, SampleBatch] | None:
|
||||
"""Receive a single message with timeout"""
|
||||
if not self.sub.poll(timeout):
|
||||
return None
|
||||
|
||||
topic_bytes, seq_bytes, payload = self.sub.recv_multipart()
|
||||
assert topic_bytes == self.topic_bytes
|
||||
|
||||
seq = int.from_bytes(seq_bytes, "big")
|
||||
data = self.decoder.decode(payload)
|
||||
self.last_seq = seq
|
||||
self.received_msgs.append((seq, data))
|
||||
return seq, data
|
||||
|
||||
def request_replay(self, start_seq: int, socket_idx: int = 0) -> None:
|
||||
"""Request replay of messages starting from start_seq"""
|
||||
if not self.replay_sockets:
|
||||
raise ValueError("Replay sockets not initialized")
|
||||
if socket_idx >= len(self.replay_sockets):
|
||||
raise ValueError(f"Invalid socket index {socket_idx}")
|
||||
|
||||
self.replay_sockets[socket_idx].send_multipart(
|
||||
[b"", start_seq.to_bytes(8, "big")]
|
||||
)
|
||||
|
||||
def receive_replay(self, socket_idx: int = 0) -> list[tuple[int, SampleBatch]]:
|
||||
"""Receive replayed messages from a specific replay socket"""
|
||||
if not self.replay_sockets:
|
||||
raise ValueError("Replay sockets not initialized")
|
||||
if socket_idx >= len(self.replay_sockets):
|
||||
raise ValueError(f"Invalid socket index {socket_idx}")
|
||||
|
||||
replay_socket = self.replay_sockets[socket_idx]
|
||||
replayed: list[tuple[int, SampleBatch]] = []
|
||||
while True:
|
||||
try:
|
||||
if not replay_socket.poll(1000):
|
||||
break
|
||||
|
||||
# DEALER receives [empty_delim, topic, seq, payload]
|
||||
frames = replay_socket.recv_multipart()
|
||||
if frames and frames[0] == b"":
|
||||
frames = frames[1:]
|
||||
if len(frames) != 3 or not frames[-1]:
|
||||
# End of replay marker
|
||||
break
|
||||
|
||||
topic, seq_bytes, payload = frames
|
||||
assert topic == self.topic_bytes
|
||||
seq = int.from_bytes(seq_bytes, "big")
|
||||
data = self.decoder.decode(payload)
|
||||
replayed.append((seq, data))
|
||||
except zmq.ZMQError as _:
|
||||
break
|
||||
|
||||
return replayed
|
||||
|
||||
def close(self):
|
||||
"""Clean up resources"""
|
||||
self.sub.close()
|
||||
for replay in self.replay_sockets:
|
||||
replay.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def enable_ray_v2_backend():
|
||||
"""Set env vars for the Ray V2 executor backend and shut down Ray
|
||||
between tests."""
|
||||
import ray
|
||||
|
||||
saved = {
|
||||
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND": os.environ.get(
|
||||
"VLLM_USE_RAY_V2_EXECUTOR_BACKEND"
|
||||
),
|
||||
"VLLM_ENABLE_V1_MULTIPROCESSING": os.environ.get(
|
||||
"VLLM_ENABLE_V1_MULTIPROCESSING"
|
||||
),
|
||||
}
|
||||
os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1"
|
||||
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
|
||||
if ray.is_initialized():
|
||||
ray.shutdown()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if ray.is_initialized():
|
||||
ray.shutdown()
|
||||
os.environ.update({k: v for k, v in saved.items() if v is not None})
|
||||
for key in (k for k, v in saved.items() if v is None):
|
||||
os.environ.pop(key, None)
|
||||
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import atexit
|
||||
import os
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.parallel_state import (
|
||||
init_distributed_environment,
|
||||
)
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
|
||||
def _distributed_worker_wrapper(fn, env, world_size, args, rank, skip_queue):
|
||||
try:
|
||||
fn(env, world_size, *args)
|
||||
except BaseException as exc:
|
||||
if isinstance(exc, pytest.skip.Exception):
|
||||
skip_queue.put((rank, str(exc)))
|
||||
return
|
||||
raise
|
||||
|
||||
|
||||
def distributed_run(fn, world_size, *args):
|
||||
number_of_processes = world_size
|
||||
processes: list[mp.Process] = []
|
||||
skip_queue: mp.SimpleQueue = mp.SimpleQueue()
|
||||
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"] = "12345"
|
||||
p = mp.Process(
|
||||
target=_distributed_worker_wrapper,
|
||||
args=(fn, env, world_size, args, i, skip_queue),
|
||||
)
|
||||
processes.append(p)
|
||||
p.start()
|
||||
|
||||
for p in processes:
|
||||
p.join()
|
||||
|
||||
skipped: list[tuple[int, str]] = []
|
||||
while not skip_queue.empty():
|
||||
rank, reason = skip_queue.get()
|
||||
skipped.append((rank, reason))
|
||||
|
||||
if len(skipped) == number_of_processes:
|
||||
reason = skipped[0][1]
|
||||
pytest.skip(reason)
|
||||
if 0 < len(skipped) < number_of_processes:
|
||||
skipped_ranks = sorted(rank for rank, _ in skipped)
|
||||
raise AssertionError(
|
||||
"Distributed test had partial skips; expected either all ranks "
|
||||
f"to skip or none. Skipped ranks: {skipped_ranks}, "
|
||||
f"total ranks: {number_of_processes}"
|
||||
)
|
||||
|
||||
for p in processes:
|
||||
assert p.exitcode == 0
|
||||
|
||||
|
||||
def set_env_vars_and_device(env: dict[str, str]) -> None:
|
||||
update_environment_variables(env)
|
||||
local_rank = os.environ["LOCAL_RANK"]
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
# Create a minimal vllm config for init_distributed_environment
|
||||
vllm_config = VllmConfig()
|
||||
with set_current_vllm_config(vllm_config):
|
||||
init_distributed_environment()
|
||||
atexit.register(_destroy_process_group_if_initialized)
|
||||
# Ensure each worker process has the same random seed
|
||||
random.seed(42)
|
||||
torch.manual_seed(42)
|
||||
|
||||
|
||||
def _destroy_process_group_if_initialized() -> None:
|
||||
if torch.distributed.is_available() and torch.distributed.is_initialized():
|
||||
torch.distributed.destroy_process_group()
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# can only run on machines with p2p access across GPUs
|
||||
# can only run with torchrun:
|
||||
# torchrun --nproc_per_node=2 tests/distributed/test_ca_buffer_sharing.py
|
||||
|
||||
import ctypes
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
|
||||
from vllm.distributed.device_communicators.custom_all_reduce import ( # noqa
|
||||
CustomAllreduce,
|
||||
)
|
||||
|
||||
# create a cpu process group for communicating metadata (ipc handle)
|
||||
dist.init_process_group(backend="gloo")
|
||||
rank = local_rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
|
||||
# every process sets its own device (differently)
|
||||
lib = CudaRTLibrary()
|
||||
lib.cudaSetDevice(rank)
|
||||
|
||||
buffer_size_in_bytes = 1024
|
||||
byte_value = 2 # the value we write to the buffer for verification
|
||||
|
||||
pointers = CustomAllreduce.create_shared_buffer(buffer_size_in_bytes)
|
||||
|
||||
print(f"Rank {rank} has pointers {pointers}")
|
||||
|
||||
dist.barrier()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
if rank == 0:
|
||||
# the first rank tries to write to all buffers
|
||||
for p in pointers:
|
||||
pointer = ctypes.c_void_p(p)
|
||||
lib.cudaMemset(pointer, byte_value, buffer_size_in_bytes)
|
||||
|
||||
dist.barrier()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
host_data = (ctypes.c_char * buffer_size_in_bytes)()
|
||||
|
||||
# all ranks read from all buffers, and check if the data is correct
|
||||
for p in pointers:
|
||||
pointer = ctypes.c_void_p(p)
|
||||
lib.cudaMemcpy(host_data, pointer, buffer_size_in_bytes)
|
||||
for i in range(buffer_size_in_bytes):
|
||||
assert ord(host_data[i]) == byte_value, (
|
||||
f"Rank {rank} failed"
|
||||
f" to verify buffer {p}. Expected {byte_value}, "
|
||||
f"got {ord(host_data[i])}"
|
||||
)
|
||||
|
||||
print(f"Rank {rank} verified all buffers")
|
||||
|
||||
dist.barrier()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
CustomAllreduce.free_shared_buffer(pointers)
|
||||
@@ -0,0 +1,382 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Test the communication operators.
|
||||
|
||||
Run `pytest tests/distributed/test_comm_ops.py`.
|
||||
"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
import torch
|
||||
|
||||
from vllm.distributed import (
|
||||
broadcast_tensor_dict,
|
||||
get_pp_group,
|
||||
tensor_model_parallel_all_gather,
|
||||
tensor_model_parallel_all_reduce,
|
||||
tensor_model_parallel_reduce_scatter,
|
||||
)
|
||||
from vllm.distributed.parallel_state import GroupCoordinator, TensorMetadata
|
||||
from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
|
||||
|
||||
from ..utils import (
|
||||
init_test_distributed_environment,
|
||||
multi_gpu_test,
|
||||
multi_process_parallel,
|
||||
)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def all_reduce_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
rank: int,
|
||||
distributed_init_port: str,
|
||||
):
|
||||
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
|
||||
# so that each worker can see all the GPUs
|
||||
# they will be able to set the device to the correct GPU
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
num_elements = 8
|
||||
all_tensors = [
|
||||
torch.arange(num_elements, dtype=torch.float32, device="cuda") * (r + 1)
|
||||
for r in range(tp_size)
|
||||
]
|
||||
expected = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
|
||||
t = all_tensors[rank % tp_size]
|
||||
t = tensor_model_parallel_all_reduce(t)
|
||||
torch.testing.assert_close(t, expected)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def reduce_scatter_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
rank: int,
|
||||
distributed_init_port: str,
|
||||
):
|
||||
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
|
||||
# so that each worker can see all the GPUs
|
||||
# they will be able to set the device to the correct GPU
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
|
||||
num_elements = 8
|
||||
all_tensors = [
|
||||
torch.arange(num_elements, dtype=torch.float32, device="cuda") * (r + 1)
|
||||
for r in range(tp_size)
|
||||
]
|
||||
|
||||
index = rank % tp_size
|
||||
partition_size = num_elements // tp_size
|
||||
all_reduce = torch.sum(torch.stack(all_tensors, dim=0), dim=0)
|
||||
expected = all_reduce[index * partition_size : (index + 1) * partition_size]
|
||||
t = all_tensors[index]
|
||||
t = tensor_model_parallel_reduce_scatter(t, 0)
|
||||
torch.testing.assert_close(t, expected)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def all_gather_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
rank: int,
|
||||
distributed_init_port: str,
|
||||
):
|
||||
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
|
||||
# so that each worker can see all the GPUs
|
||||
# they will be able to set the device to the correct GPU
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
num_dimensions = 3
|
||||
tensor_size = list(range(2, num_dimensions + 2))
|
||||
total_size = 1
|
||||
for s in tensor_size:
|
||||
total_size *= s
|
||||
for all_gather_dimension in range(num_dimensions):
|
||||
all_tensors = [
|
||||
torch.arange(total_size, dtype=torch.float32, device="cuda").reshape(
|
||||
tensor_size
|
||||
)
|
||||
* (r + 1)
|
||||
for r in range(tp_size)
|
||||
]
|
||||
expected = torch.cat(all_tensors, dim=all_gather_dimension)
|
||||
t = all_tensors[rank % tp_size]
|
||||
t = tensor_model_parallel_all_gather(t, all_gather_dimension)
|
||||
torch.testing.assert_close(t, expected)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def broadcast_tensor_dict_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
rank: int,
|
||||
distributed_init_port: str,
|
||||
):
|
||||
# it is important to delete the CUDA_VISIBLE_DEVICES environment variable
|
||||
# so that each worker can see all the GPUs
|
||||
# they will be able to set the device to the correct GPU
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
test_dict = {
|
||||
# device tensor
|
||||
"a": torch.arange(8, dtype=torch.float32, device="cuda"),
|
||||
# CPU tensor
|
||||
"b": torch.arange(16, dtype=torch.int8, device="cpu"),
|
||||
"c": "test",
|
||||
"d": [1, 2, 3],
|
||||
"e": {"a": 1, "b": 2},
|
||||
# empty tensor
|
||||
"f": torch.tensor([], dtype=torch.float32, device="cuda"),
|
||||
}
|
||||
|
||||
if (rank % tp_size) == 0:
|
||||
broadcast_tensor_dict(test_dict, src=0)
|
||||
else:
|
||||
recv_dict = broadcast_tensor_dict(src=0)
|
||||
assert len(recv_dict) == len(test_dict)
|
||||
torch.testing.assert_close(recv_dict["a"], test_dict["a"])
|
||||
torch.testing.assert_close(recv_dict["b"], test_dict["b"])
|
||||
assert recv_dict["c"] == test_dict["c"]
|
||||
assert recv_dict["d"] == test_dict["d"]
|
||||
assert recv_dict["e"] == test_dict["e"]
|
||||
torch.testing.assert_close(recv_dict["f"], test_dict["f"])
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def send_recv_tensor_dict_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
rank: int,
|
||||
distributed_init_port: str,
|
||||
):
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
|
||||
test_dict = {
|
||||
# device tensor
|
||||
"a": torch.arange(8, dtype=torch.float32, device="cuda"),
|
||||
# CPU tensor
|
||||
"b": torch.arange(16, dtype=torch.int8, device="cpu"),
|
||||
"c": "test",
|
||||
"d": [1, 2, 3],
|
||||
"e": {"a": 1, "b": 2},
|
||||
# empty tensor
|
||||
"f": torch.tensor([], dtype=torch.float32, device="cuda"),
|
||||
}
|
||||
|
||||
if not get_pp_group().is_first_rank:
|
||||
recv_dict = get_pp_group().recv_tensor_dict()
|
||||
|
||||
if not get_pp_group().is_last_rank:
|
||||
get_pp_group().send_tensor_dict(test_dict)
|
||||
|
||||
if not get_pp_group().is_first_rank:
|
||||
assert len(recv_dict) == len(test_dict)
|
||||
torch.testing.assert_close(recv_dict["a"], test_dict["a"])
|
||||
torch.testing.assert_close(recv_dict["b"], test_dict["b"])
|
||||
assert recv_dict["c"] == test_dict["c"]
|
||||
assert recv_dict["d"] == test_dict["d"]
|
||||
assert recv_dict["e"] == test_dict["e"]
|
||||
torch.testing.assert_close(recv_dict["f"], test_dict["f"])
|
||||
|
||||
|
||||
class _DummyWork:
|
||||
def __init__(self) -> None:
|
||||
self.wait_calls = 0
|
||||
|
||||
def wait(self) -> None:
|
||||
self.wait_calls += 1
|
||||
|
||||
|
||||
class _DummyAllGatherGroup:
|
||||
def __init__(self, world_size: int, rank_in_group: int) -> None:
|
||||
self.world_size = world_size
|
||||
self.rank_in_group = rank_in_group
|
||||
|
||||
def all_gather(self, t: torch.Tensor, dim: int = 0) -> torch.Tensor:
|
||||
# duplicate local slice across ranks.
|
||||
assert dim == 0
|
||||
return torch.cat([t for _ in range(self.world_size)], dim=0)
|
||||
|
||||
|
||||
def _make_group_for_unit_test(
|
||||
rank_in_group: int = 0, world_size: int = 2
|
||||
) -> GroupCoordinator:
|
||||
# avoid running GroupCoordinator.__init__ (it wires up real process groups).
|
||||
g = GroupCoordinator.__new__(GroupCoordinator)
|
||||
g.world_size = world_size
|
||||
g.rank_in_group = rank_in_group
|
||||
g.ranks = list(range(world_size))
|
||||
g.use_cpu_custom_send_recv = False
|
||||
g.device_group = None
|
||||
g.cpu_group = None
|
||||
return g
|
||||
|
||||
|
||||
def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def fake_irecv(t: torch.Tensor, *args: Any, **kwargs: Any) -> _DummyWork:
|
||||
t.fill_(1)
|
||||
return _DummyWork()
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "irecv", fake_irecv)
|
||||
|
||||
g = _make_group_for_unit_test(rank_in_group=0, world_size=2)
|
||||
# 2 tensors so we can catch late-binding bugs in postprocess closures.
|
||||
metadata_list = [
|
||||
("a", TensorMetadata("cpu", torch.int32, torch.Size([4]))),
|
||||
("b", TensorMetadata("cpu", torch.int32, torch.Size([4]))),
|
||||
]
|
||||
g.recv_object = lambda src=None: metadata_list # type: ignore[method-assign]
|
||||
|
||||
ag = _DummyAllGatherGroup(world_size=2, rank_in_group=0)
|
||||
td, handles, postprocess = g.irecv_tensor_dict(all_gather_group=ag)
|
||||
|
||||
assert td is not None
|
||||
assert len(handles) == 2
|
||||
assert len(postprocess) == 2
|
||||
|
||||
# before postprocess, dict holds the TP slice (shape 2).
|
||||
assert td["a"].shape == torch.Size([2])
|
||||
assert td["b"].shape == torch.Size([2])
|
||||
|
||||
# simulate worker-side "defer wait": wait + postprocess later.
|
||||
for handle in handles:
|
||||
handle.wait()
|
||||
for fn in postprocess:
|
||||
fn()
|
||||
|
||||
# after postprocess, dict values are reconstructed to full shape (shape 4),
|
||||
# and each key should be updated independently
|
||||
assert td["a"].shape == torch.Size([4])
|
||||
assert td["b"].shape == torch.Size([4])
|
||||
torch.testing.assert_close(td["a"], torch.ones(4, dtype=torch.int32))
|
||||
torch.testing.assert_close(td["b"], torch.ones(4, dtype=torch.int32))
|
||||
|
||||
|
||||
def test_async_intermediate_tensors_lazy_wait() -> None:
|
||||
work = _DummyWork()
|
||||
post_calls = {"n": 0}
|
||||
|
||||
def post() -> None:
|
||||
post_calls["n"] += 1
|
||||
|
||||
it = AsyncIntermediateTensors(
|
||||
{"x": torch.tensor([1])},
|
||||
comm_handles=[work],
|
||||
comm_postprocess=[post],
|
||||
)
|
||||
|
||||
# accessing non-tensor attributes should not trigger wait.
|
||||
assert it._comm_handles is not None
|
||||
assert work.wait_calls == 0
|
||||
assert post_calls["n"] == 0
|
||||
|
||||
# first access of `.tensors` triggers wait + postprocess.
|
||||
_ = it.tensors
|
||||
assert work.wait_calls == 1
|
||||
assert post_calls["n"] == 1
|
||||
|
||||
# subsequent access should not re-wait.
|
||||
_ = it.tensors
|
||||
assert work.wait_calls == 1
|
||||
assert post_calls["n"] == 1
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def send_recv_test_worker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
rank: int,
|
||||
distributed_init_port: str,
|
||||
):
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
|
||||
size = 64
|
||||
test_tensor = torch.arange(64, dtype=torch.float32, device="cuda")
|
||||
|
||||
if not get_pp_group().is_first_rank:
|
||||
recv_tensor = get_pp_group().recv(size, dtype=torch.float32)
|
||||
|
||||
if not get_pp_group().is_last_rank:
|
||||
get_pp_group().send(test_tensor)
|
||||
|
||||
if not get_pp_group().is_first_rank:
|
||||
torch.testing.assert_close(test_tensor, recv_tensor)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("tp_size", [2])
|
||||
@pytest.mark.parametrize(
|
||||
"test_target",
|
||||
[all_reduce_test_worker, all_gather_test_worker, broadcast_tensor_dict_test_worker],
|
||||
)
|
||||
def test_multi_process_tensor_parallel(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size: int,
|
||||
test_target: Callable[..., Any],
|
||||
):
|
||||
multi_process_parallel(monkeypatch, tp_size, 1, test_target)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("pp_size", [2])
|
||||
@pytest.mark.parametrize(
|
||||
"test_target", [send_recv_test_worker, send_recv_tensor_dict_test_worker]
|
||||
)
|
||||
def test_multi_process_pipeline_parallel(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
pp_size: int,
|
||||
test_target: Callable[..., Any],
|
||||
):
|
||||
multi_process_parallel(monkeypatch, 1, pp_size, test_target)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
@pytest.mark.parametrize("tp_size", [2])
|
||||
@pytest.mark.parametrize("pp_size", [2])
|
||||
@pytest.mark.parametrize(
|
||||
"test_target",
|
||||
[
|
||||
send_recv_test_worker,
|
||||
send_recv_tensor_dict_test_worker,
|
||||
all_reduce_test_worker,
|
||||
all_gather_test_worker,
|
||||
broadcast_tensor_dict_test_worker,
|
||||
],
|
||||
)
|
||||
def test_multi_process_tensor_parallel_pipeline_parallel(
|
||||
tp_size: int,
|
||||
pp_size: int,
|
||||
test_target: Callable[..., Any],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
multi_process_parallel(monkeypatch, tp_size, pp_size, test_target)
|
||||
@@ -0,0 +1,319 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
WARNING: This test runs in both single-node (4 GPUs) and multi-node
|
||||
(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
|
||||
important to set the distributed backend to "mp" to avoid Ray scheduling
|
||||
all workers in a node other than the head node, which can cause the test
|
||||
to fail.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, NamedTuple
|
||||
|
||||
import lm_eval
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.utils import RemoteOpenAIServer, create_new_process_for_each_test
|
||||
from vllm.config.model import RunnerOption
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..models.registry import HF_EXAMPLE_MODELS
|
||||
|
||||
logger = init_logger("test_context_parallel")
|
||||
|
||||
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
|
||||
|
||||
CP_TEST_MODELS = [
|
||||
# TODO support other models
|
||||
# [LANGUAGE GENERATION]
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat",
|
||||
"Qwen/Qwen2.5-1.5B-Instruct",
|
||||
"Qwen/Qwen3.5-0.8B", # hybrid attention model
|
||||
]
|
||||
|
||||
# GSM8K eval configuration
|
||||
NUM_SHOTS = 5 # Few-shot examples
|
||||
TASK = "gsm8k"
|
||||
FILTER = "exact_match,strict-match"
|
||||
NUM_CONCURRENT = 128
|
||||
# tp accuracy with 2% buffer
|
||||
MIN_ACCURACY = {
|
||||
# .buildkite/lm-eval-harness/configs/DeepSeek-V2-Lite-Chat.yaml
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": 0.64,
|
||||
# .buildkite/lm-eval-harness/configs/Qwen2.5-1.5B-Instruct.yaml
|
||||
"Qwen/Qwen2.5-1.5B-Instruct": 0.52,
|
||||
"Qwen/Qwen3.5-0.8B": 0.33,
|
||||
}
|
||||
|
||||
|
||||
class ParallelSetup(NamedTuple):
|
||||
tp_size: int
|
||||
pp_size: int
|
||||
dcp_size: int
|
||||
cp_kv_cache_interleave_size: int
|
||||
eager_mode: bool
|
||||
chunked_prefill: bool
|
||||
|
||||
|
||||
class CPTestOptions(NamedTuple):
|
||||
multi_node_only: bool
|
||||
attn_backend: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CPTestSettings:
|
||||
parallel_setups: list[ParallelSetup]
|
||||
distributed_backends: list[str]
|
||||
runner: RunnerOption
|
||||
test_options: CPTestOptions
|
||||
|
||||
@staticmethod
|
||||
def detailed(
|
||||
*,
|
||||
tp_base: int = 4,
|
||||
pp_base: int = 1,
|
||||
dcp_multipliers: list[float] | None = None,
|
||||
cp_kv_cache_interleave_size: int = 1,
|
||||
multi_node_only: bool = False,
|
||||
runner: RunnerOption = "auto",
|
||||
attn_backend: str | None = None,
|
||||
):
|
||||
parallel_setups = []
|
||||
if dcp_multipliers is None:
|
||||
dcp_multipliers = [
|
||||
0.5,
|
||||
]
|
||||
for eager_mode_val in [False]:
|
||||
for pp_multiplier in [1]:
|
||||
for dcp_multiplier in dcp_multipliers:
|
||||
for chunked_prefill_val in [True]:
|
||||
parallel_setups.append(
|
||||
ParallelSetup(
|
||||
tp_size=tp_base,
|
||||
pp_size=pp_multiplier * pp_base,
|
||||
dcp_size=int(dcp_multiplier * tp_base),
|
||||
cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
|
||||
eager_mode=eager_mode_val,
|
||||
chunked_prefill=chunked_prefill_val,
|
||||
)
|
||||
)
|
||||
return CPTestSettings(
|
||||
parallel_setups=parallel_setups,
|
||||
distributed_backends=["mp"],
|
||||
runner=runner,
|
||||
test_options=CPTestOptions(
|
||||
multi_node_only=multi_node_only,
|
||||
attn_backend=attn_backend,
|
||||
),
|
||||
)
|
||||
|
||||
def iter_params(self, model_id: str):
|
||||
opts = self.test_options
|
||||
|
||||
for parallel_setup in self.parallel_setups:
|
||||
for backend in self.distributed_backends:
|
||||
yield (
|
||||
model_id,
|
||||
parallel_setup,
|
||||
backend,
|
||||
self.runner,
|
||||
opts,
|
||||
)
|
||||
|
||||
|
||||
if current_platform.is_rocm():
|
||||
CP_TEXT_GENERATION_MODELS = {
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": [
|
||||
CPTestSettings.detailed(dcp_multipliers=[1]),
|
||||
],
|
||||
"Qwen/Qwen2.5-1.5B-Instruct": [
|
||||
CPTestSettings.detailed(dcp_multipliers=[1]),
|
||||
],
|
||||
}
|
||||
else:
|
||||
CP_TEXT_GENERATION_MODELS = {
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": [
|
||||
CPTestSettings.detailed(dcp_multipliers=[1]),
|
||||
CPTestSettings.detailed(
|
||||
dcp_multipliers=[0.5],
|
||||
cp_kv_cache_interleave_size=64,
|
||||
attn_backend="FLASHMLA",
|
||||
),
|
||||
],
|
||||
"Qwen/Qwen2.5-1.5B-Instruct": [
|
||||
CPTestSettings.detailed(
|
||||
cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
|
||||
),
|
||||
CPTestSettings.detailed(
|
||||
cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
|
||||
),
|
||||
],
|
||||
"Qwen/Qwen3.5-0.8B": [
|
||||
CPTestSettings.detailed(
|
||||
cp_kv_cache_interleave_size=16,
|
||||
attn_backend="FLASH_ATTN",
|
||||
),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _test_cp_gsm8k(
|
||||
model_id: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: CPTestOptions,
|
||||
num_gpus_available: int,
|
||||
*,
|
||||
method: Literal["generate"],
|
||||
is_multimodal: bool,
|
||||
):
|
||||
(
|
||||
tp_size,
|
||||
pp_size,
|
||||
dcp_size,
|
||||
cp_kv_cache_interleave_size,
|
||||
eager_mode,
|
||||
chunked_prefill,
|
||||
) = parallel_setup
|
||||
|
||||
multi_node_only, attn_backend = test_options
|
||||
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
trust_remote_code = model_info.trust_remote_code
|
||||
tokenizer_mode = model_info.tokenizer_mode
|
||||
hf_overrides = model_info.hf_overrides
|
||||
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
if num_gpus_available < tp_size * pp_size:
|
||||
pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
|
||||
if VLLM_MULTI_NODE and distributed_backend == "mp":
|
||||
pytest.skip(
|
||||
"Skipping multi-node pipeline parallel test for "
|
||||
"multiprocessing distributed backend"
|
||||
)
|
||||
if multi_node_only and not VLLM_MULTI_NODE:
|
||||
pytest.skip("Not in multi-node setting")
|
||||
|
||||
server_args = [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"bfloat16",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
"64",
|
||||
]
|
||||
if chunked_prefill:
|
||||
server_args.append("--enable-chunked-prefill")
|
||||
if eager_mode:
|
||||
server_args.append("--enforce-eager")
|
||||
if runner != "auto":
|
||||
server_args.extend(["--runner", runner])
|
||||
if trust_remote_code:
|
||||
server_args.append("--trust-remote-code")
|
||||
if tokenizer_mode:
|
||||
server_args.extend(["--tokenizer-mode", tokenizer_mode])
|
||||
if hf_overrides:
|
||||
server_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
|
||||
|
||||
server_args.extend(
|
||||
[
|
||||
"--tensor-parallel-size",
|
||||
str(tp_size),
|
||||
"--pipeline-parallel-size",
|
||||
str(pp_size),
|
||||
"--decode-context-parallel-size",
|
||||
str(dcp_size),
|
||||
"--dcp-kv-cache-interleave-size",
|
||||
str(cp_kv_cache_interleave_size),
|
||||
"--distributed-executor-backend",
|
||||
distributed_backend,
|
||||
]
|
||||
)
|
||||
|
||||
if attn_backend:
|
||||
server_args.append(f"--attention-backend={attn_backend}")
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
model_id,
|
||||
server_args,
|
||||
max_wait_seconds=720,
|
||||
) as remote_server:
|
||||
url = f"{remote_server.url_for('v1')}/completions"
|
||||
|
||||
model_args = (
|
||||
f"model={model_id},"
|
||||
f"base_url={url},"
|
||||
f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
|
||||
)
|
||||
|
||||
results = lm_eval.simple_evaluate(
|
||||
model="local-completions",
|
||||
model_args=model_args,
|
||||
tasks=TASK,
|
||||
num_fewshot=NUM_SHOTS,
|
||||
)
|
||||
|
||||
# Validate accuracy is reasonable
|
||||
accuracy = results["results"][TASK][FILTER]
|
||||
min_accuracy = MIN_ACCURACY[model_id]
|
||||
assert accuracy >= min_accuracy, (
|
||||
f"TP+DCP accuracy too low: {accuracy:.3f} < {min_accuracy:.3f}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
"model_id",
|
||||
"parallel_setup",
|
||||
"distributed_backend",
|
||||
"runner",
|
||||
"test_options",
|
||||
),
|
||||
[
|
||||
params
|
||||
for model_id, settings in CP_TEXT_GENERATION_MODELS.items()
|
||||
for setting in settings
|
||||
for params in setting.iter_params(model_id)
|
||||
if model_id in CP_TEST_MODELS
|
||||
],
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_cp_generation(
|
||||
model_id: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: CPTestOptions,
|
||||
num_gpus_available,
|
||||
):
|
||||
if (
|
||||
model_id == "deepseek-ai/DeepSeek-V2-Lite-Chat"
|
||||
and torch.cuda.get_device_capability() < (9, 0)
|
||||
):
|
||||
pytest.skip(reason="MLA+DCP requires compute capability of 9.0 or higher")
|
||||
if (
|
||||
model_id == "Qwen/Qwen2.5-1.5B-Instruct"
|
||||
and torch.cuda.get_device_capability() != (9, 0)
|
||||
):
|
||||
pytest.skip(reason="GQA+DCP currently requires compute capability of 9.0")
|
||||
|
||||
_test_cp_gsm8k(
|
||||
model_id,
|
||||
parallel_setup,
|
||||
distributed_backend,
|
||||
runner,
|
||||
test_options,
|
||||
num_gpus_available,
|
||||
method="generate",
|
||||
is_multimodal=False,
|
||||
)
|
||||
@@ -0,0 +1,132 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
|
||||
from vllm.distributed.parallel_state import get_tp_group, graph_capture
|
||||
|
||||
from ..utils import (
|
||||
ensure_model_parallel_initialized,
|
||||
init_test_distributed_environment,
|
||||
multi_process_parallel,
|
||||
)
|
||||
|
||||
random.seed(42)
|
||||
test_sizes = [random.randint(1024, 2048 * 1024) for _ in range(8)]
|
||||
for i, v in enumerate(test_sizes):
|
||||
test_sizes[i] -= v % 8
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
ensure_model_parallel_initialized(tp_size, pp_size)
|
||||
group = get_tp_group().device_group
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
torch.distributed.all_reduce(data, group=group)
|
||||
torch.accelerator.synchronize()
|
||||
del data
|
||||
|
||||
# we use the first group to communicate once
|
||||
# and the second group to communicate twice
|
||||
# and so on
|
||||
# this is used to demonstrate that each group can
|
||||
# communicate independently
|
||||
num_communication = rank // tp_size + 1
|
||||
|
||||
for sz in test_sizes:
|
||||
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
|
||||
with graph_capture(device=device) as graph_capture_context:
|
||||
# use integers so result matches NCCL exactly
|
||||
device_idx = torch.accelerator.current_device_index()
|
||||
inp1 = torch.randint(1, 16, (sz,), dtype=dtype, device=device_idx)
|
||||
inp2 = torch.randint(1, 16, (sz,), dtype=dtype, device=device_idx)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=graph_capture_context.stream):
|
||||
for i in range(num_communication):
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
# the input buffer is immediately modified to test
|
||||
# synchronization
|
||||
dist.all_reduce(inp1, group=group)
|
||||
out2 = tensor_model_parallel_all_reduce(inp2)
|
||||
dist.all_reduce(inp2, group=group)
|
||||
graph.replay()
|
||||
torch.testing.assert_close(out1, inp1)
|
||||
torch.testing.assert_close(out2, inp2)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
|
||||
# we use the first group to communicate once
|
||||
# and the second group to communicate twice
|
||||
# and so on
|
||||
# this is used to demonstrate that each group can
|
||||
# communicate independently
|
||||
num_communication = rank // tp_size + 1
|
||||
sz = 1024
|
||||
fa = get_tp_group().device_communicator.ca_comm
|
||||
inp = torch.ones(sz, dtype=torch.float32, device=device)
|
||||
out = inp
|
||||
for _ in range(num_communication):
|
||||
out = fa.all_reduce(out, registered=False)
|
||||
torch.testing.assert_close(out, inp * (tp_size**num_communication))
|
||||
|
||||
inp = torch.ones(sz * 4, dtype=torch.bfloat16, device=device)
|
||||
out = inp
|
||||
for _ in range(num_communication):
|
||||
out = fa.all_reduce(out, registered=False)
|
||||
torch.testing.assert_close(out, inp * (tp_size**num_communication))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp_size", [2])
|
||||
@pytest.mark.parametrize("pipeline_parallel_size", [1, 2])
|
||||
@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
|
||||
def test_custom_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pipeline_parallel_size,
|
||||
test_target,
|
||||
):
|
||||
world_size = tp_size * pipeline_parallel_size
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
|
||||
@@ -0,0 +1,502 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for DCP A2A communication backend (no GPU required).
|
||||
|
||||
Tests cover:
|
||||
1. DCP A2A config validation (--dcp-comm-backend)
|
||||
2. KVP group function exists
|
||||
3. LSE-weighted combination correctness
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import multiprocess as mp
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.config.parallel import ParallelConfig
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
mp.set_start_method("spawn", force=True)
|
||||
|
||||
|
||||
class _FakeCPGroup:
|
||||
def __init__(self, world_size: int, device_group: dist.ProcessGroup):
|
||||
self.world_size = world_size
|
||||
self.device_group = device_group
|
||||
|
||||
|
||||
def _dtype_from_name(dtype_name: str) -> torch.dtype:
|
||||
return {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}[dtype_name]
|
||||
|
||||
|
||||
def _packed_a2a_reference(
|
||||
cp_attn_out: torch.Tensor,
|
||||
cp_attn_lse: torch.Tensor,
|
||||
world_size: int,
|
||||
h_per_rank: int,
|
||||
is_lse_base_on_e: bool,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
B, _H, D = cp_attn_out.shape
|
||||
outputs = (
|
||||
cp_attn_out.view(B, world_size, h_per_rank, D)
|
||||
.permute(1, 0, 2, 3)
|
||||
.contiguous()
|
||||
.float()
|
||||
)
|
||||
lses = cp_attn_lse.view(B, world_size, h_per_rank).permute(1, 0, 2).contiguous()
|
||||
return _lse_weighted_combine(
|
||||
outputs,
|
||||
lses,
|
||||
return_lse=True,
|
||||
is_lse_base_on_e=is_lse_base_on_e,
|
||||
)
|
||||
|
||||
|
||||
def _assert_packed_a2a_close(
|
||||
actual: torch.Tensor,
|
||||
expected: torch.Tensor,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
if dtype == torch.float32:
|
||||
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
|
||||
else:
|
||||
torch.testing.assert_close(
|
||||
actual.float(), expected.float(), rtol=3e-2, atol=3e-2
|
||||
)
|
||||
|
||||
|
||||
def _distributed_run(fn, world_size: int, extra_env: dict[str, str]) -> None:
|
||||
port = str(get_open_port())
|
||||
processes: list[mp.Process] = []
|
||||
for rank in range(world_size):
|
||||
env = {
|
||||
"RANK": str(rank),
|
||||
"LOCAL_RANK": str(rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"LOCAL_WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": port,
|
||||
**extra_env,
|
||||
}
|
||||
process = mp.Process(target=fn, args=(env,))
|
||||
processes.append(process)
|
||||
process.start()
|
||||
|
||||
for process in processes:
|
||||
process.join(timeout=120)
|
||||
|
||||
for process in processes:
|
||||
if process.is_alive():
|
||||
process.kill()
|
||||
process.join()
|
||||
assert process.exitcode == 0
|
||||
|
||||
|
||||
class TestDCPCommBackendConfig:
|
||||
"""Test --dcp-comm-backend config validation."""
|
||||
|
||||
def test_default_is_ag_rs(self):
|
||||
"""Default comm backend is ag_rs."""
|
||||
config = ParallelConfig()
|
||||
assert config.dcp_comm_backend == "ag_rs"
|
||||
|
||||
def test_a2a_requires_dcp_greater_than_1(self):
|
||||
"""A2A backend requires decode_context_parallel_size > 1."""
|
||||
with pytest.raises(
|
||||
ValueError, match="requires decode_context_parallel_size > 1"
|
||||
):
|
||||
ParallelConfig(
|
||||
dcp_comm_backend="a2a",
|
||||
decode_context_parallel_size=1,
|
||||
)
|
||||
|
||||
def test_a2a_with_dcp_valid(self):
|
||||
"""A2A backend is valid when DCP > 1."""
|
||||
config = ParallelConfig(
|
||||
dcp_comm_backend="a2a",
|
||||
tensor_parallel_size=4,
|
||||
decode_context_parallel_size=4,
|
||||
)
|
||||
assert config.dcp_comm_backend == "a2a"
|
||||
|
||||
def test_invalid_backend_rejected(self):
|
||||
"""Invalid backend values are rejected."""
|
||||
with pytest.raises(ValueError, match="must be one of|Input should be"):
|
||||
ParallelConfig(
|
||||
dcp_comm_backend="invalid",
|
||||
)
|
||||
|
||||
def test_ag_rs_with_dcp_1_valid(self):
|
||||
"""ag_rs backend is valid with DCP=1 (no DCP)."""
|
||||
config = ParallelConfig(
|
||||
dcp_comm_backend="ag_rs",
|
||||
decode_context_parallel_size=1,
|
||||
)
|
||||
assert config.dcp_comm_backend == "ag_rs"
|
||||
|
||||
|
||||
class TestLSEWeightedCombine:
|
||||
"""Test LSE-weighted combination logic (CPU only, no GPU).
|
||||
|
||||
The _lse_weighted_combine function is the reference implementation
|
||||
that verifies the Triton kernel's correctness. It computes:
|
||||
|
||||
result[b,h,d] = sum_n(w_n * output_n[b,h,d])
|
||||
|
||||
where w_n = softmax(lse_n) = exp(lse_n) / sum_k(exp(lse_k))
|
||||
"""
|
||||
|
||||
def test_importable(self):
|
||||
"""Verify _lse_weighted_combine is importable."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
assert callable(_lse_weighted_combine)
|
||||
|
||||
def test_single_rank(self):
|
||||
"""Single rank: output unchanged."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
# N=1, B=2, H=4, D=8
|
||||
outputs = torch.randn(1, 2, 4, 8)
|
||||
lses = torch.randn(1, 2, 4)
|
||||
|
||||
result = _lse_weighted_combine(outputs, lses)
|
||||
|
||||
assert result.shape == (2, 4, 8)
|
||||
torch.testing.assert_close(result, outputs.squeeze(0), rtol=1e-5, atol=1e-5)
|
||||
|
||||
def test_equal_lse(self):
|
||||
"""Equal LSE values: outputs averaged equally."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
_N, B, H, D = 2, 1, 1, 4
|
||||
outputs = torch.tensor(
|
||||
[
|
||||
[[[1.0, 2.0, 3.0, 4.0]]], # Rank 0
|
||||
[[[5.0, 6.0, 7.0, 8.0]]], # Rank 1
|
||||
]
|
||||
)
|
||||
lses = torch.tensor(
|
||||
[
|
||||
[[0.0]], # Rank 0
|
||||
[[0.0]], # Rank 1
|
||||
]
|
||||
)
|
||||
|
||||
result = _lse_weighted_combine(outputs, lses)
|
||||
|
||||
expected = (outputs[0] + outputs[1]) / 2
|
||||
assert result.shape == (B, H, D)
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
|
||||
|
||||
def test_dominant_rank(self):
|
||||
"""Different LSE values: larger LSE gets more weight."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
B, H, D = 1, 1, 2
|
||||
outputs = torch.tensor(
|
||||
[
|
||||
[[[0.0, 0.0]]], # Rank 0
|
||||
[[[1.0, 1.0]]], # Rank 1
|
||||
]
|
||||
)
|
||||
lses = torch.tensor(
|
||||
[
|
||||
[[-100.0]], # Rank 0: negligible contribution
|
||||
[[0.0]], # Rank 1: dominant
|
||||
]
|
||||
)
|
||||
|
||||
result = _lse_weighted_combine(outputs, lses)
|
||||
|
||||
assert result.shape == (B, H, D)
|
||||
torch.testing.assert_close(result, outputs[1], atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_mathematically_correct(self):
|
||||
"""Verify mathematical correctness of LSE combination."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
outputs = torch.tensor(
|
||||
[
|
||||
[[[2.0, 4.0]]],
|
||||
[[[6.0, 8.0]]],
|
||||
]
|
||||
)
|
||||
lses = torch.tensor(
|
||||
[
|
||||
[[1.0]], # exp(1) ≈ 2.718
|
||||
[[2.0]], # exp(2) ≈ 7.389
|
||||
]
|
||||
)
|
||||
|
||||
result = _lse_weighted_combine(outputs, lses)
|
||||
|
||||
w0 = math.exp(1) / (math.exp(1) + math.exp(2))
|
||||
w1 = math.exp(2) / (math.exp(1) + math.exp(2))
|
||||
expected = torch.tensor([[[w0 * 2.0 + w1 * 6.0, w0 * 4.0 + w1 * 8.0]]])
|
||||
|
||||
torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
|
||||
|
||||
def test_return_lse(self):
|
||||
"""return_lse=True returns global LSE (logsumexp of inputs)."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
B, H, D = 1, 1, 2
|
||||
outputs = torch.tensor(
|
||||
[
|
||||
[[[1.0, 2.0]]],
|
||||
[[[3.0, 4.0]]],
|
||||
]
|
||||
)
|
||||
lses = torch.tensor(
|
||||
[
|
||||
[[1.0]],
|
||||
[[2.0]],
|
||||
]
|
||||
)
|
||||
|
||||
result, global_lse = _lse_weighted_combine(outputs, lses, return_lse=True)
|
||||
|
||||
expected_global_lse = math.log(math.exp(1) + math.exp(2))
|
||||
|
||||
assert result.shape == (B, H, D)
|
||||
assert global_lse.shape == (B, H)
|
||||
assert abs(global_lse.item() - expected_global_lse) < 1e-5
|
||||
|
||||
def test_base2_return_lse(self):
|
||||
"""Base-2 LSE mode returns log2-sum-exp2 global LSE."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
outputs = torch.tensor(
|
||||
[
|
||||
[[[1.0, 2.0]]],
|
||||
[[[3.0, 4.0]]],
|
||||
]
|
||||
)
|
||||
lses = torch.tensor(
|
||||
[
|
||||
[[1.0]],
|
||||
[[2.0]],
|
||||
]
|
||||
)
|
||||
|
||||
result, global_lse = _lse_weighted_combine(
|
||||
outputs,
|
||||
lses,
|
||||
return_lse=True,
|
||||
is_lse_base_on_e=False,
|
||||
)
|
||||
|
||||
expected_global_lse = math.log2(2**1 + 2**2)
|
||||
w0 = 2**1 / (2**1 + 2**2)
|
||||
w1 = 2**2 / (2**1 + 2**2)
|
||||
expected = torch.tensor([[[w0 * 1.0 + w1 * 3.0, w0 * 2.0 + w1 * 4.0]]])
|
||||
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
|
||||
torch.testing.assert_close(
|
||||
global_lse,
|
||||
torch.tensor([[expected_global_lse]]),
|
||||
rtol=1e-5,
|
||||
atol=1e-5,
|
||||
)
|
||||
|
||||
def test_lse_pack_dim(self):
|
||||
"""Packed A2A stores one fp32 LSE in output-dtype lanes."""
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _dcp_a2a_lse_pack_dim
|
||||
|
||||
assert _dcp_a2a_lse_pack_dim(torch.bfloat16) == 2
|
||||
assert _dcp_a2a_lse_pack_dim(torch.float16) == 2
|
||||
assert _dcp_a2a_lse_pack_dim(torch.float32) == 1
|
||||
|
||||
|
||||
class TestPackedA2AKernels:
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 1, reason="CUDA is required."
|
||||
)
|
||||
@pytest.mark.parametrize("dtype_name", ["float16", "bfloat16", "float32"])
|
||||
@pytest.mark.parametrize("return_lse", [False, True])
|
||||
@pytest.mark.parametrize("is_lse_base_on_e", [False, True])
|
||||
def test_pack_unpack_combine_matches_reference(
|
||||
self,
|
||||
dtype_name: str,
|
||||
return_lse: bool,
|
||||
is_lse_base_on_e: bool,
|
||||
):
|
||||
from vllm.v1.attention.ops.dcp_alltoall import (
|
||||
_dcp_a2a_lse_pack_dim,
|
||||
_dcp_a2a_pack_send,
|
||||
_dcp_a2a_unpack_combine,
|
||||
)
|
||||
|
||||
torch.manual_seed(0)
|
||||
dtype = _dtype_from_name(dtype_name)
|
||||
device = torch.device("cuda")
|
||||
world_size, B, h_per_rank, D = 4, 7, 2, 32
|
||||
H = world_size * h_per_rank
|
||||
cp_attn_out = torch.randn(B, H, D, device=device, dtype=dtype)
|
||||
cp_attn_lse = torch.randn(B, H, device=device, dtype=torch.float32)
|
||||
lse_pack_dim = _dcp_a2a_lse_pack_dim(dtype)
|
||||
send_buffer = torch.empty(
|
||||
(world_size, B, h_per_rank, D + lse_pack_dim),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
_dcp_a2a_pack_send(
|
||||
cp_attn_out,
|
||||
cp_attn_lse,
|
||||
send_buffer,
|
||||
world_size,
|
||||
h_per_rank,
|
||||
D,
|
||||
lse_pack_dim,
|
||||
)
|
||||
actual = _dcp_a2a_unpack_combine(
|
||||
send_buffer, D, lse_pack_dim, return_lse, is_lse_base_on_e
|
||||
)
|
||||
expected_out, expected_lse = _packed_a2a_reference(
|
||||
cp_attn_out, cp_attn_lse, world_size, h_per_rank, is_lse_base_on_e
|
||||
)
|
||||
|
||||
if return_lse:
|
||||
actual_out, actual_lse = actual
|
||||
_assert_packed_a2a_close(actual_out, expected_out, dtype)
|
||||
torch.testing.assert_close(actual_lse, expected_lse, rtol=1e-4, atol=1e-4)
|
||||
else:
|
||||
_assert_packed_a2a_close(actual, expected_out, dtype)
|
||||
|
||||
|
||||
def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:
|
||||
update_environment_variables(env)
|
||||
local_rank = int(env["LOCAL_RANK"])
|
||||
torch.accelerator.set_device_index(local_rank)
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
dist.init_process_group(
|
||||
backend="cpu:gloo,cuda:nccl",
|
||||
device_id=torch.device(f"cuda:{local_rank}"),
|
||||
)
|
||||
else:
|
||||
dist.init_process_group(backend="nccl")
|
||||
use_workspace = env.get("USE_WORKSPACE") == "1"
|
||||
if use_workspace:
|
||||
from vllm.v1.worker.workspace import init_workspace_manager
|
||||
|
||||
init_workspace_manager(torch.device(f"cuda:{local_rank}"))
|
||||
try:
|
||||
from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
|
||||
|
||||
dtype = _dtype_from_name(env["TEST_DTYPE"])
|
||||
return_lse = env["RETURN_LSE"] == "1"
|
||||
is_lse_base_on_e = env["LSE_BASE_E"] == "1"
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
B, h_per_rank, D = 5, 2, 32
|
||||
H = world_size * h_per_rank
|
||||
|
||||
generator = torch.Generator(device=f"cuda:{local_rank}")
|
||||
generator.manual_seed(1234 + rank)
|
||||
cp_attn_out = torch.randn(
|
||||
B,
|
||||
H,
|
||||
D,
|
||||
device=f"cuda:{local_rank}",
|
||||
dtype=dtype,
|
||||
generator=generator,
|
||||
)
|
||||
cp_attn_lse = torch.randn(
|
||||
B,
|
||||
H,
|
||||
device=f"cuda:{local_rank}",
|
||||
dtype=torch.float32,
|
||||
generator=generator,
|
||||
)
|
||||
actual = dcp_a2a_lse_reduce(
|
||||
cp_attn_out,
|
||||
cp_attn_lse,
|
||||
_FakeCPGroup(world_size, dist.group.WORLD),
|
||||
return_lse=return_lse,
|
||||
is_lse_base_on_e=is_lse_base_on_e,
|
||||
)
|
||||
|
||||
gathered_out = [torch.empty_like(cp_attn_out) for _ in range(world_size)]
|
||||
gathered_lse = [torch.empty_like(cp_attn_lse) for _ in range(world_size)]
|
||||
dist.all_gather(gathered_out, cp_attn_out)
|
||||
dist.all_gather(gathered_lse, cp_attn_lse)
|
||||
outputs = torch.stack(
|
||||
[
|
||||
t[:, rank * h_per_rank : (rank + 1) * h_per_rank, :]
|
||||
for t in gathered_out
|
||||
],
|
||||
dim=0,
|
||||
).float()
|
||||
lses = torch.stack(
|
||||
[t[:, rank * h_per_rank : (rank + 1) * h_per_rank] for t in gathered_lse],
|
||||
dim=0,
|
||||
)
|
||||
from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
|
||||
|
||||
expected_out, expected_lse = _lse_weighted_combine(
|
||||
outputs,
|
||||
lses,
|
||||
return_lse=True,
|
||||
is_lse_base_on_e=is_lse_base_on_e,
|
||||
)
|
||||
|
||||
if return_lse:
|
||||
actual_out, actual_lse = actual
|
||||
_assert_packed_a2a_close(actual_out, expected_out, dtype)
|
||||
torch.testing.assert_close(actual_lse, expected_lse, rtol=1e-4, atol=1e-4)
|
||||
else:
|
||||
_assert_packed_a2a_close(actual, expected_out, dtype)
|
||||
finally:
|
||||
if use_workspace:
|
||||
from vllm.v1.worker.workspace import reset_workspace_manager
|
||||
|
||||
reset_workspace_manager()
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs."
|
||||
)
|
||||
@pytest.mark.parametrize("dtype_name", ["float16", "bfloat16", "float32"])
|
||||
def test_distributed_packed_a2a_matches_reference(dtype_name: str):
|
||||
_distributed_run(
|
||||
_distributed_packed_a2a_worker,
|
||||
world_size=4,
|
||||
extra_env={
|
||||
"TEST_DTYPE": dtype_name,
|
||||
"RETURN_LSE": "1",
|
||||
"LSE_BASE_E": "1",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs."
|
||||
)
|
||||
def test_distributed_packed_a2a_with_workspace_matches_reference():
|
||||
_distributed_run(
|
||||
_distributed_packed_a2a_worker,
|
||||
world_size=4,
|
||||
extra_env={
|
||||
"TEST_DTYPE": "bfloat16",
|
||||
"RETURN_LSE": "1",
|
||||
"LSE_BASE_E": "1",
|
||||
"USE_WORKSPACE": "1",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from tests.plugins_tests.test_oot_registration_online import (
|
||||
run_and_test_dummy_opt_api_server,
|
||||
)
|
||||
|
||||
|
||||
def test_distributed_oot(dummy_opt_path: str):
|
||||
run_and_test_dummy_opt_api_server(dummy_opt_path, tp=2)
|
||||
@@ -0,0 +1,205 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from ..evals.gsm8k.gsm8k_eval import evaluate_gsm8k
|
||||
from ..utils import RemoteOpenAIServer, multi_gpu_test
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_ray_between_tests():
|
||||
"""Force-stop any lingering Ray processes between tests."""
|
||||
subprocess.run(["ray", "stop", "--force"], timeout=30, capture_output=True)
|
||||
time.sleep(5)
|
||||
yield
|
||||
|
||||
|
||||
MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat"
|
||||
|
||||
NUM_GSM8K_QUESTIONS = 256
|
||||
EXPECTED_ACCURACY = 0.58
|
||||
ACCURACY_TOL = 0.08
|
||||
MAX_NUM_SEQS = 32
|
||||
|
||||
|
||||
def _send_scale_command(server: RemoteOpenAIServer, new_dp_size: int) -> bool:
|
||||
url = server.url_for("scale_elastic_ep")
|
||||
payload = {"new_data_parallel_size": new_dp_size}
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
try:
|
||||
response = requests.post(url, json=payload, headers=headers, timeout=300)
|
||||
return response.status_code == 200
|
||||
except requests.exceptions.RequestException:
|
||||
return False
|
||||
|
||||
|
||||
def _run_gsm8k_eval(server: RemoteOpenAIServer, stage: str) -> float:
|
||||
assert server.port is not None
|
||||
result = evaluate_gsm8k(
|
||||
num_questions=NUM_GSM8K_QUESTIONS,
|
||||
host=f"http://{server.host}",
|
||||
port=server.port,
|
||||
)
|
||||
accuracy = result["accuracy"]
|
||||
print(
|
||||
f"[{stage}] GSM8K accuracy: {accuracy:.3f} "
|
||||
f"({result['num_questions']} questions)"
|
||||
)
|
||||
assert accuracy >= EXPECTED_ACCURACY, (
|
||||
f"[{stage}] GSM8K accuracy {accuracy:.3f} is below "
|
||||
f"expected threshold {EXPECTED_ACCURACY}"
|
||||
)
|
||||
return accuracy
|
||||
|
||||
|
||||
def _base_serve_args(use_async_eplb: bool = False) -> list[str]:
|
||||
args = [
|
||||
"--trust-remote-code",
|
||||
"--tensor-parallel-size",
|
||||
"1",
|
||||
"--gpu-memory-utilization",
|
||||
"0.8",
|
||||
"--max-model-len",
|
||||
"4096",
|
||||
"--max-num-seqs",
|
||||
str(MAX_NUM_SEQS),
|
||||
"--enable-expert-parallel",
|
||||
"--all2all-backend",
|
||||
"allgather_reducescatter",
|
||||
"--enable-elastic-ep",
|
||||
"--enable-eplb",
|
||||
"--eplb-config.num_redundant_experts",
|
||||
"0",
|
||||
"--eplb-config.use_async",
|
||||
"true" if use_async_eplb else "false",
|
||||
"--eplb-config.step_interval",
|
||||
"10",
|
||||
"--eplb-config.window_size",
|
||||
"5",
|
||||
"--data-parallel-backend",
|
||||
"ray",
|
||||
"--data-parallel-size",
|
||||
"2",
|
||||
"--api-server-count",
|
||||
"1",
|
||||
]
|
||||
|
||||
leader_address = os.environ.get("LEADER_ADDRESS")
|
||||
if leader_address:
|
||||
args.extend(["--data-parallel-address", leader_address])
|
||||
|
||||
return args
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
|
||||
)
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_elastic_ep_scaling(use_async_eplb: bool):
|
||||
if use_async_eplb:
|
||||
from vllm.distributed.eplb.eplb_communicator import has_nixl
|
||||
|
||||
if not has_nixl():
|
||||
pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
|
||||
|
||||
vllm_serve_args = _base_serve_args(use_async_eplb)
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
|
||||
) as server:
|
||||
initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")
|
||||
|
||||
assert _send_scale_command(server, 4)
|
||||
time.sleep(10)
|
||||
scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (4 GPUs)")
|
||||
|
||||
assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
|
||||
f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "
|
||||
f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
|
||||
)
|
||||
|
||||
assert _send_scale_command(server, 2)
|
||||
time.sleep(5)
|
||||
scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
|
||||
|
||||
assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
|
||||
f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "
|
||||
f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
|
||||
)
|
||||
|
||||
print("\nAccuracy Summary:")
|
||||
print(f" Initial: {initial_accuracy:.3f}")
|
||||
print(
|
||||
f" Scale up: {scale_up_accuracy:.3f} "
|
||||
f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})"
|
||||
)
|
||||
print(
|
||||
f" Scale down: {scale_down_accuracy:.3f} "
|
||||
f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"
|
||||
)
|
||||
print(f" Tolerance: {ACCURACY_TOL:.3f}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"use_async_eplb", [False, True], ids=["sync_eplb", "async_eplb"]
|
||||
)
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_elastic_ep_scaling_uneven(use_async_eplb: bool):
|
||||
"""Test scale up with uneven worker distribution.
|
||||
|
||||
This tests the case where num_new_workers % old_dp_size != 0,
|
||||
specifically 2 -> 3 where remainder = 1 % 2 = 1.
|
||||
This exercises the remainder handling in sender-receiver pairing.
|
||||
"""
|
||||
if use_async_eplb:
|
||||
from vllm.distributed.eplb.eplb_communicator import has_nixl
|
||||
|
||||
if not has_nixl():
|
||||
pytest.skip("Async EPLB with elastic EP requires NIXL (not installed)")
|
||||
|
||||
vllm_serve_args = _base_serve_args(use_async_eplb)
|
||||
|
||||
with RemoteOpenAIServer(
|
||||
MODEL_NAME, vllm_serve_args, env_dict={}, max_wait_seconds=1200
|
||||
) as server:
|
||||
initial_accuracy = _run_gsm8k_eval(server, "Initial (2 GPUs)")
|
||||
|
||||
# Scale 2 -> 3: This has remainder = 1 % 2 = 1
|
||||
# Tests uneven sender-receiver pairing
|
||||
assert _send_scale_command(server, 3)
|
||||
time.sleep(10)
|
||||
scale_up_accuracy = _run_gsm8k_eval(server, "After scale up (3 GPUs)")
|
||||
|
||||
assert scale_up_accuracy >= initial_accuracy - ACCURACY_TOL, (
|
||||
f"Scale up accuracy {scale_up_accuracy:.3f} dropped more than "
|
||||
f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
|
||||
)
|
||||
|
||||
# Scale back down to 2
|
||||
assert _send_scale_command(server, 2)
|
||||
time.sleep(5)
|
||||
scale_down_accuracy = _run_gsm8k_eval(server, "After scale down (2 GPUs)")
|
||||
|
||||
assert scale_down_accuracy >= initial_accuracy - ACCURACY_TOL, (
|
||||
f"Scale down accuracy {scale_down_accuracy:.3f} dropped more than "
|
||||
f"{ACCURACY_TOL} below initial accuracy {initial_accuracy:.3f}"
|
||||
)
|
||||
|
||||
print("\nAccuracy Summary (Uneven Scaling):")
|
||||
print(f" Initial: {initial_accuracy:.3f}")
|
||||
print(
|
||||
f" Scale up: {scale_up_accuracy:.3f} "
|
||||
f"(diff: {scale_up_accuracy - initial_accuracy:+.3f})"
|
||||
)
|
||||
print(
|
||||
f" Scale down: {scale_down_accuracy:.3f} "
|
||||
f"(diff: {scale_down_accuracy - initial_accuracy:+.3f})"
|
||||
)
|
||||
print(f" Tolerance: {ACCURACY_TOL:.3f}")
|
||||
@@ -0,0 +1,498 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.distributed.eplb.eplb_state import compute_logical_maps
|
||||
from vllm.distributed.eplb.policy.default import DefaultEplbPolicy
|
||||
|
||||
|
||||
def test_basic_rebalance():
|
||||
"""Test basic rebalancing functionality"""
|
||||
# Example from https://github.com/deepseek-ai/eplb
|
||||
weight = torch.tensor(
|
||||
[
|
||||
[90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
|
||||
[20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
|
||||
]
|
||||
)
|
||||
|
||||
num_layers = weight.shape[0]
|
||||
num_replicas = 16
|
||||
num_groups = 4
|
||||
num_nodes = 2
|
||||
num_gpus = 8
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Verify output shapes
|
||||
assert phy2log.shape == (
|
||||
2,
|
||||
16,
|
||||
), f"Expected `phy2log` shape (2, 16), got {phy2log.shape}"
|
||||
assert log2phy.shape[0] == 2, (
|
||||
f"Expected `log2phy` first dimension 2, got {log2phy.shape[0]}"
|
||||
)
|
||||
assert log2phy.shape[1] == 12, (
|
||||
f"Expected `log2phy` second dimension 12, got {log2phy.shape[1]}"
|
||||
)
|
||||
assert logcnt.shape == (
|
||||
2,
|
||||
12,
|
||||
), f"Expected `logcnt` shape (2, 12), got {logcnt.shape}"
|
||||
|
||||
# Verify physical to logical expert mapping range is correct
|
||||
assert torch.all(phy2log >= 0) and torch.all(phy2log < 12), (
|
||||
"Physical to logical mapping should be in range [0, 12)"
|
||||
)
|
||||
|
||||
# Verify expert count reasonableness
|
||||
assert torch.all(logcnt >= 1), "Each logical expert should have at least 1 replica"
|
||||
assert torch.sum(logcnt, dim=1).sum() == num_replicas * num_layers, (
|
||||
f"Total replicas should be {num_replicas * num_layers}"
|
||||
)
|
||||
|
||||
# Verify expected output
|
||||
expected_phy2log = torch.tensor(
|
||||
[
|
||||
[5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1],
|
||||
[7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1],
|
||||
]
|
||||
)
|
||||
assert torch.all(phy2log == expected_phy2log)
|
||||
|
||||
expected_logcnt = torch.tensor(
|
||||
[[1, 2, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1], [1, 2, 1, 1, 1, 2, 2, 1, 2, 1, 1, 1]]
|
||||
)
|
||||
assert torch.all(logcnt == expected_logcnt)
|
||||
|
||||
|
||||
def test_single_gpu_case():
|
||||
"""Test single GPU case"""
|
||||
weight = torch.tensor([[10, 20, 30, 40]])
|
||||
num_replicas = 4
|
||||
num_groups = 1
|
||||
num_nodes = 1
|
||||
num_gpus = 1
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Verify shapes
|
||||
assert phy2log.shape == (1, 4)
|
||||
assert log2phy.shape[0] == 1
|
||||
assert log2phy.shape[1] == 4
|
||||
assert logcnt.shape == (1, 4)
|
||||
|
||||
# Verify all logical experts are mapped
|
||||
assert set(phy2log[0].tolist()) == {0, 1, 2, 3}
|
||||
|
||||
|
||||
def test_equal_weights():
|
||||
"""Test case with equal weights"""
|
||||
weight = torch.tensor([[50, 50, 50, 50, 50, 50, 50, 50]])
|
||||
num_replicas = 8
|
||||
num_groups = 2
|
||||
num_nodes = 2
|
||||
num_gpus = 4
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Verify shapes
|
||||
assert phy2log.shape == (1, 8)
|
||||
assert logcnt.shape == (1, 8)
|
||||
|
||||
# With equal weights, each expert should have exactly one replica
|
||||
assert torch.all(logcnt == 1), (
|
||||
"With equal weights and no replication, "
|
||||
"each expert should have exactly 1 replica"
|
||||
)
|
||||
|
||||
|
||||
def test_extreme_weight_imbalance():
|
||||
"""Test extreme weight imbalance case"""
|
||||
weight = torch.tensor([[1000, 1, 1, 1, 1, 1, 1, 1]])
|
||||
num_replicas = 12
|
||||
num_groups = 2
|
||||
num_nodes = 2
|
||||
num_gpus = 4
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Verify shapes
|
||||
assert phy2log.shape == (1, 12)
|
||||
assert logcnt.shape == (1, 8)
|
||||
|
||||
# Expert with highest weight (index 0) should have more replicas
|
||||
assert logcnt[0, 0] > logcnt[0, 1], (
|
||||
"Expert with highest weight should have more replicas"
|
||||
)
|
||||
|
||||
|
||||
def test_multiple_layers():
|
||||
"""Test multiple layers case"""
|
||||
weight = torch.tensor(
|
||||
[
|
||||
[10, 20, 30, 40, 50, 60], # First layer
|
||||
[60, 50, 40, 30, 20, 10], # Second layer (opposite weight pattern)
|
||||
[25, 25, 25, 25, 25, 25], # Third layer (equal weights)
|
||||
]
|
||||
)
|
||||
num_replicas = 8
|
||||
num_groups = 2
|
||||
num_nodes = 2
|
||||
num_gpus = 4
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Verify shapes
|
||||
assert phy2log.shape == (3, 8)
|
||||
assert logcnt.shape == (3, 6)
|
||||
|
||||
# Verify expert allocation is reasonable for each layer
|
||||
for layer in range(3):
|
||||
assert torch.all(phy2log[layer] >= 0) and torch.all(phy2log[layer] < 6), (
|
||||
f"Layer {layer} physical to logical mappingshould be in range [0, 6)"
|
||||
)
|
||||
assert torch.sum(logcnt[layer]) == num_replicas, (
|
||||
f"Layer {layer} total replicas should be {num_replicas}"
|
||||
)
|
||||
|
||||
|
||||
def test_parameter_validation():
|
||||
"""Test parameter validation"""
|
||||
weight = torch.tensor([[10, 20, 30, 40]])
|
||||
|
||||
# Test non-divisible case - this should handle normally without throwing
|
||||
# errors because the function will fall back to global load balancing
|
||||
# strategy
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
assert phy2log.shape == (1, 8)
|
||||
assert logcnt.shape == (1, 4)
|
||||
|
||||
# Test cases that will actually cause errors:
|
||||
# num_physical_experts not divisible by num_gpus
|
||||
with pytest.raises(AssertionError):
|
||||
DefaultEplbPolicy.rebalance_experts(weight, 7, 2, 2, 4) # 7 not divisible by 4
|
||||
|
||||
|
||||
def test_small_scale_hierarchical():
|
||||
"""Test small-scale hierarchical load balancing"""
|
||||
weight = torch.tensor(
|
||||
[
|
||||
[100, 50, 200, 75, 150, 25, 300, 80], # 8 experts
|
||||
]
|
||||
)
|
||||
num_replicas = 12
|
||||
num_groups = 4 # 4 groups, 2 experts each
|
||||
num_nodes = 2 # 2 nodes
|
||||
num_gpus = 4 # 4 GPUs
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Verify basic constraints
|
||||
assert phy2log.shape == (1, 12)
|
||||
assert logcnt.shape == (1, 8)
|
||||
assert torch.sum(logcnt) == num_replicas
|
||||
assert torch.all(logcnt >= 1)
|
||||
|
||||
# Expert with highest weight should have more replicas
|
||||
max_weight_expert = torch.argmax(weight[0])
|
||||
assert logcnt[0, max_weight_expert] >= 2, (
|
||||
"Highest weight expert should have multiple replicas"
|
||||
)
|
||||
|
||||
|
||||
def test_global_load_balance_fallback():
|
||||
"""Test global load balancing fallback case"""
|
||||
# When num_groups % num_nodes != 0, should fall back to global load
|
||||
# balancing
|
||||
weight = torch.tensor([[10, 20, 30, 40, 50, 60]])
|
||||
num_replicas = 8
|
||||
num_groups = 3 # Cannot be divided evenly by num_nodes=2
|
||||
num_nodes = 2
|
||||
num_gpus = 4
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Should work normally, just using global load balancing strategy
|
||||
assert phy2log.shape == (1, 8)
|
||||
assert logcnt.shape == (1, 6)
|
||||
assert torch.sum(logcnt) == num_replicas
|
||||
|
||||
|
||||
@pytest.mark.parametrize("device", ["cpu", "cuda"])
|
||||
def test_device_compatibility(device):
|
||||
"""Test device compatibility"""
|
||||
if device == "cuda" and not torch.cuda.is_available():
|
||||
pytest.skip("CUDA not available")
|
||||
|
||||
weight = torch.tensor([[10, 20, 30, 40]], device=device)
|
||||
num_replicas = 6
|
||||
num_groups = 2
|
||||
num_nodes = 1
|
||||
num_gpus = 2
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
_, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
|
||||
|
||||
# Function will convert to CPU internally, but should handle different
|
||||
# device inputs normally
|
||||
assert phy2log.shape == (1, 6)
|
||||
assert logcnt.shape == (1, 4)
|
||||
|
||||
|
||||
def test_additional_cases():
|
||||
"""Test more edge cases and different parameter combinations"""
|
||||
|
||||
# Test case 1: Large-scale distributed setup
|
||||
weight1 = torch.tensor(
|
||||
[[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
|
||||
)
|
||||
phy2log1 = DefaultEplbPolicy.rebalance_experts(weight1, 24, 8, 4, 8)
|
||||
_, logcnt1 = compute_logical_maps(phy2log1, weight1.shape[-1])
|
||||
|
||||
assert phy2log1.shape == (1, 24)
|
||||
assert logcnt1.shape == (1, 16)
|
||||
assert torch.sum(logcnt1) == 24
|
||||
|
||||
# Test case 2: Different weight distributions
|
||||
weight2 = torch.tensor(
|
||||
[
|
||||
[200, 150, 100, 50, 25, 12], # Decreasing weights
|
||||
[12, 25, 50, 100, 150, 200], # Increasing weights
|
||||
]
|
||||
)
|
||||
phy2log2 = DefaultEplbPolicy.rebalance_experts(weight2, 10, 3, 1, 2)
|
||||
_, logcnt2 = compute_logical_maps(phy2log2, weight2.shape[-1])
|
||||
|
||||
assert phy2log2.shape == (2, 10)
|
||||
assert logcnt2.shape == (2, 6)
|
||||
|
||||
# Verify high-weight experts have more replicas
|
||||
for layer in range(2):
|
||||
max_weight_idx = torch.argmax(weight2[layer])
|
||||
assert logcnt2[layer, max_weight_idx] >= 2
|
||||
|
||||
|
||||
def test_compute_logical_maps_with_negative_indices():
|
||||
"""
|
||||
Test that compute_logical_maps correctly handles physical slots containing
|
||||
-1 (unused slots).
|
||||
"""
|
||||
# 2 layers, 6 physical slots, 4 logical experts.
|
||||
# Slots 2 and 5 are unused (-1).
|
||||
phy2log = torch.tensor(
|
||||
[
|
||||
[0, 1, -1, 2, 3, -1],
|
||||
[3, -1, 2, 1, 0, -1],
|
||||
]
|
||||
)
|
||||
num_layers = 2
|
||||
num_logical_experts = 4
|
||||
|
||||
log2phy, logcnt = compute_logical_maps(phy2log, num_logical_experts)
|
||||
|
||||
assert logcnt.shape == (num_layers, num_logical_experts)
|
||||
assert log2phy.shape == (num_layers, num_logical_experts, 1)
|
||||
|
||||
expected_logcnt = torch.ones(num_layers, num_logical_experts, dtype=phy2log.dtype)
|
||||
assert torch.all(logcnt == expected_logcnt), (
|
||||
f"Expected that all replica counts == 1, got {logcnt}"
|
||||
)
|
||||
|
||||
assert torch.all(log2phy >= 0), (
|
||||
"log2phy should only contain valid physical indices, not -1"
|
||||
)
|
||||
|
||||
assert log2phy[0, 0, 0] == 0
|
||||
assert log2phy[0, 1, 0] == 1
|
||||
assert log2phy[0, 2, 0] == 3
|
||||
assert log2phy[0, 3, 0] == 4
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
weight = torch.tensor(
|
||||
[
|
||||
[90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
|
||||
[20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27],
|
||||
]
|
||||
)
|
||||
|
||||
num_replicas = 16
|
||||
num_groups = 4
|
||||
num_nodes = 2
|
||||
num_gpus = 8
|
||||
|
||||
phy2log = DefaultEplbPolicy.rebalance_experts(
|
||||
weight, num_replicas, num_groups, num_nodes, num_gpus
|
||||
)
|
||||
print(phy2log)
|
||||
|
||||
test_basic_rebalance()
|
||||
|
||||
|
||||
def _make_phy_replicas_idx_from_phy2log(phy2log: np.ndarray) -> np.ndarray:
|
||||
"""Create replicas indices mapping from phy2log."""
|
||||
pr = np.zeros_like(phy2log, dtype=np.int64)
|
||||
for layer in range(phy2log.shape[0]):
|
||||
seen: dict[int, int] = {}
|
||||
row = phy2log[layer].tolist()
|
||||
for i, expert in enumerate(row):
|
||||
r = seen.get(expert, 0)
|
||||
pr[layer, i] = r
|
||||
seen[expert] = r + 1
|
||||
return pr
|
||||
|
||||
|
||||
def _validate_intragpu_rearrangement(
|
||||
old_global_expert_indices: np.ndarray,
|
||||
new_phy2log: np.ndarray,
|
||||
new_phy_replicas_idx: np.ndarray,
|
||||
post_phy2log: np.ndarray,
|
||||
post_phy_replicas_idx: np.ndarray,
|
||||
num_ranks: int,
|
||||
slots_per_gpu: int,
|
||||
):
|
||||
# Per-GPU checks
|
||||
for gpu_idx in range(num_ranks):
|
||||
start = gpu_idx * slots_per_gpu
|
||||
end = start + slots_per_gpu
|
||||
old_seg = old_global_expert_indices[0, start:end]
|
||||
new_seg = new_phy2log[0, start:end]
|
||||
new_rnk = new_phy_replicas_idx[0, start:end]
|
||||
post_seg = post_phy2log[0, start:end]
|
||||
post_rnk = post_phy_replicas_idx[0, start:end]
|
||||
|
||||
# Pairwise equality for (expert, rank) pairs to ensure nothing is lost
|
||||
def sorted_pairs(seg, rnk):
|
||||
pairs = list(zip(seg.tolist(), rnk.tolist()))
|
||||
pairs.sort()
|
||||
return pairs
|
||||
|
||||
assert sorted_pairs(post_seg, post_rnk) == sorted_pairs(new_seg, new_rnk), (
|
||||
f"Per-GPU pairs of (expert,rank) must match new mapping for GPU {gpu_idx}"
|
||||
)
|
||||
|
||||
# For experts that remain on the same GPU, the old slot is preserved
|
||||
# for at least one occurrence; rank at that slot must be valid for that expert
|
||||
old_list = old_seg.tolist()
|
||||
new_list = new_seg.tolist()
|
||||
post_list = post_seg.tolist()
|
||||
remained = set(old_list) & set(new_list)
|
||||
new_ranks_for_expert: dict[int, list[int]] = {}
|
||||
for v, r in zip(new_list, new_rnk.tolist()):
|
||||
new_ranks_for_expert.setdefault(v, []).append(r)
|
||||
for expert in remained:
|
||||
old_pos = old_list.index(expert)
|
||||
assert post_list[old_pos] == expert, (
|
||||
f"Expert {expert} on GPU {gpu_idx} should stay at old slot {old_pos}"
|
||||
)
|
||||
# Rank at preserved slot must be one of the ranks
|
||||
# the expert has in new mapping
|
||||
assert post_rnk.tolist()[old_pos] in new_ranks_for_expert[expert], (
|
||||
f"Rank for expert {expert} at preserved slot on GPU {gpu_idx} "
|
||||
"must come from new mapping"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_ranks, slots_per_gpu, old_phy2log, new_phy2log",
|
||||
[
|
||||
pytest.param(
|
||||
# Setup: 2 GPUs, 4 slots each, 1 layer
|
||||
# Old mapping: GPU0 -> [0,1,2,3], GPU1 -> [4,5,6,7]
|
||||
# New mapping shuffles within GPU0 and brings 4,5 into GPU0.
|
||||
# GPU0 new -> [1,5,0,4]; GPU1 new -> [6,2,7,3]
|
||||
2,
|
||||
4,
|
||||
np.array([[0, 1, 2, 3, 4, 5, 6, 7]]),
|
||||
np.array([[1, 5, 0, 4, 6, 2, 7, 3]]),
|
||||
id="simple",
|
||||
),
|
||||
pytest.param(
|
||||
# Setup: 2 GPUs, 5 slots each (total 10 physical experts), 1 layer
|
||||
# Old mapping:
|
||||
# GPU0 -> [0, 1, 0, 2, 3] (expert 0 duplicated)
|
||||
# GPU1 -> [4, 5, 6, 1, 2]
|
||||
# New mapping reorders within GPUs and moves some experts across GPUs,
|
||||
# while still including duplicates:
|
||||
# GPU0 new -> [0, 5, 4, 0, 1] (expert 0 duplicated, 4/5 incoming)
|
||||
# GPU1 new -> [6, 2, 3, 2, 1] (expert 2 duplicated)
|
||||
2,
|
||||
5,
|
||||
np.array([[0, 1, 0, 2, 3, 4, 5, 6, 1, 2]]),
|
||||
np.array([[0, 5, 4, 0, 1, 6, 2, 3, 2, 1]]),
|
||||
id="duplicates",
|
||||
),
|
||||
pytest.param(
|
||||
# Setup: 3 GPUs, 4 slots each (total 12 physical experts), 1 layer
|
||||
# Old mapping:
|
||||
# GPU0 -> [0, 1, 2, 3]
|
||||
# GPU1 -> [0, 1, 2, 3]
|
||||
# GPU2 -> [0, 1, 2, 3]
|
||||
# New mapping decides to use one expert on 2 GPUs and shuffles
|
||||
# experts on the third GPU,
|
||||
# GPU0 new -> [0, 0, 0, 0]
|
||||
# GPU1 new -> [0, 0, 0, 0]
|
||||
# GPU2 new -> [1, 2, 3, 0]
|
||||
3,
|
||||
4,
|
||||
np.array([[0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]]),
|
||||
np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 0]]),
|
||||
id="skewed_expert",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_preserve_intragpu_slots(
|
||||
num_ranks: int,
|
||||
slots_per_gpu: int,
|
||||
old_phy2log: torch.Tensor,
|
||||
new_phy2log: torch.Tensor,
|
||||
):
|
||||
"""Experts that stay on a GPU keep their old slots; incoming not lost."""
|
||||
phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(new_phy2log)
|
||||
|
||||
post_phy2log = DefaultEplbPolicy.preserve_intragpu_slots(
|
||||
new_phy2log, num_ranks, old_phy2log
|
||||
)
|
||||
post_phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(post_phy2log)
|
||||
|
||||
# Shapes preserved
|
||||
assert post_phy2log.shape == new_phy2log.shape
|
||||
assert post_phy_replicas_idx.shape == phy_replicas_idx.shape
|
||||
|
||||
_validate_intragpu_rearrangement(
|
||||
old_phy2log,
|
||||
new_phy2log,
|
||||
phy_replicas_idx,
|
||||
post_phy2log,
|
||||
post_phy_replicas_idx,
|
||||
num_ranks,
|
||||
slots_per_gpu,
|
||||
)
|
||||
@@ -0,0 +1,98 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import threading
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.distributed.eplb.eplb_utils import CpuGpuEvent
|
||||
|
||||
|
||||
def test_wait_blocks_until_record():
|
||||
event = CpuGpuEvent()
|
||||
record_stream = torch.cuda.Stream()
|
||||
wait_stream = torch.cuda.Stream()
|
||||
wait_returned = threading.Event()
|
||||
|
||||
def waiter():
|
||||
event.wait(stream=wait_stream)
|
||||
wait_returned.set()
|
||||
|
||||
t = threading.Thread(target=waiter)
|
||||
t.start()
|
||||
|
||||
time.sleep(0.05)
|
||||
assert not wait_returned.is_set(), "wait() returned before record() was called"
|
||||
|
||||
event.record(stream=record_stream)
|
||||
t.join(timeout=5.0)
|
||||
|
||||
assert not event._recorded.is_set()
|
||||
|
||||
|
||||
def test_reuse_across_multiple_cycles():
|
||||
wrapper = CpuGpuEvent()
|
||||
record_stream = torch.cuda.Stream()
|
||||
wait_stream = torch.cuda.Stream()
|
||||
NUM_CYCLES = 8
|
||||
completed_cycles = []
|
||||
barriers = [threading.Barrier(2) for _ in range(NUM_CYCLES)]
|
||||
|
||||
def waiter():
|
||||
for i in range(NUM_CYCLES):
|
||||
wrapper.wait(stream=wait_stream)
|
||||
completed_cycles.append(True)
|
||||
barriers[i].wait()
|
||||
|
||||
t = threading.Thread(target=waiter)
|
||||
t.start()
|
||||
|
||||
for i in range(NUM_CYCLES):
|
||||
wrapper.record(stream=record_stream)
|
||||
barriers[i].wait()
|
||||
|
||||
t.join(timeout=10.0)
|
||||
assert len(completed_cycles) == NUM_CYCLES
|
||||
|
||||
|
||||
def test_producer_consumer():
|
||||
"""
|
||||
This test uses the CpuGpuEvent to synchronize reads and writes to/from a shared GPU
|
||||
tensor on multiple CPU threads.
|
||||
"""
|
||||
worker_stream = torch.cuda.Stream()
|
||||
# Create a single element counter that will be shared between two threads
|
||||
buf = torch.zeros(1, device="cuda")
|
||||
NUM_ROUNDS = 5
|
||||
|
||||
ready_cpu = [threading.Event() for _ in range(NUM_ROUNDS)]
|
||||
events = [CpuGpuEvent() for _ in range(NUM_ROUNDS)]
|
||||
errors: list[str] = []
|
||||
|
||||
# For each round, the worker thread (writer) sets the counter in buf and waits for
|
||||
# the main thread to read it.
|
||||
def worker():
|
||||
for i in range(NUM_ROUNDS):
|
||||
if i > 0:
|
||||
events[i - 1].wait(stream=worker_stream)
|
||||
|
||||
with torch.cuda.stream(worker_stream):
|
||||
buf.fill_(float(i + 1))
|
||||
|
||||
worker_stream.synchronize()
|
||||
ready_cpu[i].set()
|
||||
|
||||
t = threading.Thread(target=worker)
|
||||
t.start()
|
||||
|
||||
for i in range(NUM_ROUNDS):
|
||||
ready_cpu[i].wait()
|
||||
snapshot = buf.clone()
|
||||
events[i].record()
|
||||
val = snapshot.item()
|
||||
if val != float(i + 1):
|
||||
errors.append(f"round {i}: expected {i + 1:.1f}, got {val:.1f}")
|
||||
|
||||
t.join(timeout=10.0)
|
||||
assert not errors, f"Buffer ordering errors: {errors}"
|
||||
@@ -0,0 +1,906 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.eplb.eplb_communicator import (
|
||||
create_eplb_communicator,
|
||||
has_nixl,
|
||||
)
|
||||
from vllm.distributed.eplb.rebalance_execute import (
|
||||
move_from_buffer,
|
||||
rearrange_expert_weights_inplace,
|
||||
transfer_layer,
|
||||
)
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
get_tp_group,
|
||||
)
|
||||
|
||||
from .eplb_utils import distributed_run, set_env_vars_and_device
|
||||
|
||||
|
||||
def create_expert_indices_with_redundancy(
|
||||
num_layers: int,
|
||||
num_logical_experts: int,
|
||||
total_physical_experts: int,
|
||||
redundancy_config: list[int], # redundancy for each logical expert
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Create expert indices with redundancy.
|
||||
|
||||
Args:
|
||||
num_layers: number of layers
|
||||
num_logical_experts: number of logical experts
|
||||
total_physical_experts: total number of physical experts
|
||||
redundancy_config: redundancy for each logical expert
|
||||
|
||||
Returns:
|
||||
indices: Shape (num_layers, total_physical_experts)
|
||||
"""
|
||||
assert sum(redundancy_config) == total_physical_experts
|
||||
assert len(redundancy_config) == num_logical_experts
|
||||
|
||||
indices = torch.zeros(num_layers, total_physical_experts, dtype=torch.long)
|
||||
|
||||
for layer in range(num_layers):
|
||||
physical_pos = 0
|
||||
for logical_expert_id, redundancy in enumerate(redundancy_config):
|
||||
for _ in range(redundancy):
|
||||
indices[layer, physical_pos] = logical_expert_id
|
||||
physical_pos += 1
|
||||
|
||||
# Shuffle the indices at dim 1
|
||||
for layer in range(num_layers):
|
||||
indices[layer] = indices[layer][torch.randperm(indices.shape[1])]
|
||||
|
||||
return indices
|
||||
|
||||
|
||||
def create_expert_weights(
|
||||
num_layers: int,
|
||||
num_local_experts: int,
|
||||
hidden_sizes: list[int],
|
||||
rank: int,
|
||||
device: torch.device,
|
||||
physical_to_logical_mapping: torch.Tensor,
|
||||
) -> list[list[torch.Tensor]]:
|
||||
"""
|
||||
Create fake expert weights tensor for testing.
|
||||
|
||||
Use `arange` to generate predictable weights values, based on logical
|
||||
expert ID.
|
||||
All replicas of the same logical expert should have the same weights.
|
||||
|
||||
Args:
|
||||
physical_to_logical_mapping: Shape (num_layers, num_local_experts)
|
||||
mapping[layer, physical_pos] = logical_expert_id
|
||||
"""
|
||||
expert_weights = []
|
||||
|
||||
for layer in range(num_layers):
|
||||
layer_weights = []
|
||||
for weight_idx, hidden_size in enumerate(hidden_sizes):
|
||||
weight_tensor = torch.zeros(
|
||||
num_local_experts, hidden_size, device=device, dtype=torch.float32
|
||||
)
|
||||
|
||||
for local_expert in range(num_local_experts):
|
||||
# Get the logical expert ID for this physical expert
|
||||
global_pos = rank * num_local_experts + local_expert
|
||||
logical_expert_id = physical_to_logical_mapping[
|
||||
layer, global_pos
|
||||
].item()
|
||||
|
||||
# Generate weights based on logical expert ID
|
||||
# (so that all replicas of the same logical expert have the
|
||||
# same weights)
|
||||
base_value = logical_expert_id * 1000 + layer * 100 + weight_idx * 10
|
||||
weight_tensor[local_expert] = torch.arange(
|
||||
base_value,
|
||||
base_value + hidden_size,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
layer_weights.append(weight_tensor)
|
||||
expert_weights.append(layer_weights)
|
||||
|
||||
return expert_weights
|
||||
|
||||
|
||||
def create_redundancy_config(
|
||||
num_logical_experts: int,
|
||||
num_physical_experts: int,
|
||||
) -> list[int]:
|
||||
"""Create a redundancy configuration."""
|
||||
redundancy_config = [1] * num_logical_experts
|
||||
remaining = num_physical_experts - num_logical_experts
|
||||
# Randomly assign the remaining physical experts to the logical experts
|
||||
for _ in range(remaining):
|
||||
redundancy_config[random.choice(range(num_logical_experts))] += 1
|
||||
return redundancy_config
|
||||
|
||||
|
||||
def verify_expert_weights_after_shuffle(
|
||||
expert_weights: list[list[torch.Tensor]],
|
||||
new_indices: torch.Tensor,
|
||||
hidden_sizes: list[int],
|
||||
ep_rank: int,
|
||||
num_local_experts: int,
|
||||
) -> bool:
|
||||
"""Verify the weights after shuffling are correct."""
|
||||
num_layers = len(expert_weights)
|
||||
ok = True
|
||||
|
||||
for layer in range(num_layers):
|
||||
for weight_idx, hidden_size in enumerate(hidden_sizes):
|
||||
weight_tensor = expert_weights[layer][weight_idx]
|
||||
|
||||
for local_expert in range(num_local_experts):
|
||||
# Calculate the global expert ID for this local expert
|
||||
global_pos = ep_rank * num_local_experts + local_expert
|
||||
expected_logical_expert = new_indices[layer, global_pos].item()
|
||||
|
||||
# Check if the weights are correct
|
||||
actual_weights = weight_tensor[local_expert]
|
||||
expected_base = (
|
||||
expected_logical_expert * 1000 + layer * 100 + weight_idx * 10
|
||||
)
|
||||
expected_weights = torch.arange(
|
||||
expected_base,
|
||||
expected_base + hidden_size,
|
||||
device=actual_weights.device,
|
||||
dtype=actual_weights.dtype,
|
||||
)
|
||||
|
||||
if not torch.equal(actual_weights, expected_weights):
|
||||
ok = False
|
||||
actual_head = actual_weights[:8].detach().cpu().tolist()
|
||||
expected_head = expected_weights[:8].detach().cpu().tolist()
|
||||
print(
|
||||
"verify_expert_weights_after_shuffle failed: "
|
||||
f"rank={ep_rank}, "
|
||||
f"layer={layer}, weight_idx={weight_idx}, "
|
||||
f"local_expert={local_expert}, "
|
||||
f"expected_logical_expert={expected_logical_expert}, "
|
||||
f"actual_head={actual_head}, expected_head={expected_head}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
return ok
|
||||
|
||||
|
||||
def verify_redundant_experts_have_same_weights(
|
||||
expert_weights: list[list[torch.Tensor]],
|
||||
indices: torch.Tensor,
|
||||
hidden_sizes: list[int],
|
||||
ep_rank: int,
|
||||
world_size: int,
|
||||
num_local_experts: int,
|
||||
) -> bool:
|
||||
"""
|
||||
Verify that all replicas of the same logical expert have the same weights.
|
||||
"""
|
||||
num_layers = len(expert_weights)
|
||||
total_physical_experts = world_size * num_local_experts
|
||||
|
||||
ok = True
|
||||
for layer in range(num_layers):
|
||||
# Collect weights for all physical experts for each weight matrix
|
||||
all_weights: list[torch.Tensor] = []
|
||||
|
||||
for weight_idx, hidden_size in enumerate(hidden_sizes):
|
||||
# Create tensor to store all expert weights
|
||||
# Shape: [total_physical_experts, hidden_size]
|
||||
gathered_weights = torch.zeros(
|
||||
total_physical_experts,
|
||||
hidden_size,
|
||||
device=expert_weights[layer][weight_idx].device,
|
||||
dtype=expert_weights[layer][weight_idx].dtype,
|
||||
)
|
||||
|
||||
# Use all_gather to collect expert weights from current node
|
||||
# expert_weights[layer][weight_idx] shape:
|
||||
# [num_local_experts, hidden_size]
|
||||
local_weights = expert_weights[layer][
|
||||
weight_idx
|
||||
] # [num_local_experts, hidden_size]
|
||||
|
||||
# Split tensor along dim 0 into a list for all_gather
|
||||
gathered_weights_list = torch.chunk(gathered_weights, world_size, dim=0)
|
||||
|
||||
torch.distributed.all_gather(
|
||||
# Output list: each element corresponds to one rank's weights
|
||||
list(gathered_weights_list),
|
||||
local_weights, # Input: current rank's local weights
|
||||
)
|
||||
|
||||
all_weights.append(gathered_weights)
|
||||
|
||||
# Verify that all replicas of the same logical expert have the same
|
||||
# weights
|
||||
logical_expert_weights: dict[int, dict[int, torch.Tensor]] = {}
|
||||
|
||||
for physical_pos in range(total_physical_experts):
|
||||
logical_expert_id = int(indices[layer, physical_pos].item())
|
||||
|
||||
if logical_expert_id not in logical_expert_weights:
|
||||
# First time encountering this logical expert, save its weights
|
||||
logical_expert_weights[logical_expert_id] = {
|
||||
weight_idx: all_weights[weight_idx][physical_pos]
|
||||
for weight_idx in range(len(hidden_sizes))
|
||||
}
|
||||
else:
|
||||
# Verify that current physical expert's weights match the
|
||||
# previously saved logical expert weights
|
||||
for weight_idx in range(len(hidden_sizes)):
|
||||
if not torch.equal(
|
||||
all_weights[weight_idx][physical_pos],
|
||||
logical_expert_weights[logical_expert_id][weight_idx],
|
||||
):
|
||||
ok = False
|
||||
actual_head = (
|
||||
all_weights[weight_idx][physical_pos][:8]
|
||||
.detach()
|
||||
.cpu()
|
||||
.tolist()
|
||||
)
|
||||
reference_head = (
|
||||
logical_expert_weights[logical_expert_id][weight_idx][:8]
|
||||
.detach()
|
||||
.cpu()
|
||||
.tolist()
|
||||
)
|
||||
print(
|
||||
"verify_redundant_experts_have_same_weights failed: "
|
||||
f"rank={ep_rank}, "
|
||||
f"layer={layer}, weight_idx={weight_idx}, "
|
||||
f"logical_expert={logical_expert_id}, "
|
||||
f"physical_pos={physical_pos}, "
|
||||
f"actual_head={actual_head}, "
|
||||
f"reference_head={reference_head}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
return ok
|
||||
|
||||
|
||||
def assert_verification_synced(local_ok: bool, msg: str) -> None:
|
||||
ok_tensor = torch.tensor([1 if local_ok else 0], device="cuda", dtype=torch.int32)
|
||||
torch.distributed.all_reduce(ok_tensor, op=torch.distributed.ReduceOp.MIN)
|
||||
assert bool(ok_tensor.item()), msg
|
||||
|
||||
|
||||
def create_eplb_communicator_or_raise(
|
||||
*, group_coordinator, backend, expert_weights, expert_buffer
|
||||
):
|
||||
try:
|
||||
return create_eplb_communicator(
|
||||
group_coordinator=group_coordinator,
|
||||
backend=backend,
|
||||
expert_weights=expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
f"Failed to create EPLB communicator for backend={backend}: {exc}"
|
||||
) from exc
|
||||
|
||||
|
||||
def _test_async_transfer_layer_without_mtp_worker(
|
||||
env,
|
||||
world_size: int,
|
||||
num_layers: int,
|
||||
num_local_experts: int,
|
||||
num_logical_experts: int,
|
||||
eplb_communicator: str,
|
||||
) -> None:
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.tensor_parallel_size = world_size
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group_coordinator = get_tp_group()
|
||||
ep_group = ep_group_coordinator.device_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
device = torch.device(f"cuda:{ep_rank}")
|
||||
|
||||
total_physical_experts = world_size * num_local_experts
|
||||
hidden_sizes = [16, 32]
|
||||
|
||||
redundancy_config = create_redundancy_config(
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
)
|
||||
old_indices = create_expert_indices_with_redundancy(
|
||||
num_layers,
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
redundancy_config,
|
||||
)
|
||||
|
||||
new_redundancy_config = create_redundancy_config(
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
)
|
||||
new_indices = create_expert_indices_with_redundancy(
|
||||
num_layers,
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
new_redundancy_config,
|
||||
)
|
||||
|
||||
expert_weights = create_expert_weights(
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
device,
|
||||
old_indices,
|
||||
)
|
||||
old_indices_cpu = old_indices.cpu()
|
||||
new_indices_cpu = new_indices.cpu()
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
|
||||
cuda_stream = torch.cuda.Stream(device=device)
|
||||
|
||||
communicator = create_eplb_communicator_or_raise(
|
||||
group_coordinator=ep_group_coordinator,
|
||||
backend=eplb_communicator,
|
||||
expert_weights=expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
communicator.set_stream(cuda_stream)
|
||||
|
||||
for layer_idx in range(num_layers):
|
||||
transfer_metadata = transfer_layer(
|
||||
old_layer_indices=old_indices_cpu[layer_idx],
|
||||
new_layer_indices=new_indices_cpu[layer_idx],
|
||||
expert_weights=expert_weights[layer_idx],
|
||||
expert_weights_buffer=expert_buffer,
|
||||
ep_group=ep_group,
|
||||
communicator=communicator,
|
||||
cuda_stream=cuda_stream,
|
||||
layer_idx=layer_idx,
|
||||
)
|
||||
cuda_stream.synchronize()
|
||||
move_from_buffer(
|
||||
expert_weights=expert_weights[layer_idx],
|
||||
expert_weights_buffers=expert_buffer,
|
||||
transfer_metadata=transfer_metadata,
|
||||
new_indices=new_indices_cpu[layer_idx].numpy(),
|
||||
ep_rank=ep_rank,
|
||||
)
|
||||
|
||||
local_ok = verify_expert_weights_after_shuffle(
|
||||
expert_weights,
|
||||
new_indices,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
num_local_experts,
|
||||
)
|
||||
local_ok = (
|
||||
verify_redundant_experts_have_same_weights(
|
||||
expert_weights,
|
||||
new_indices,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
world_size,
|
||||
num_local_experts,
|
||||
)
|
||||
and local_ok
|
||||
)
|
||||
assert_verification_synced(
|
||||
local_ok,
|
||||
"Async transfer verification failed on at least one rank. "
|
||||
"See logs for details.",
|
||||
)
|
||||
|
||||
|
||||
def _test_rearrange_expert_weights_with_redundancy(
|
||||
env,
|
||||
world_size,
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
num_logical_experts,
|
||||
eplb_communicator: str,
|
||||
) -> None:
|
||||
# Initialize model parallel (using tensor parallel as an entrypoint
|
||||
# to expert parallel)
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.tensor_parallel_size = world_size
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group_coordinator = get_tp_group()
|
||||
ep_group = ep_group_coordinator.cpu_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
device = torch.device(f"cuda:{ep_rank}")
|
||||
|
||||
# Test parameters
|
||||
total_physical_experts = world_size * num_local_experts
|
||||
hidden_sizes = [32, 64] # Two different weight matrices
|
||||
|
||||
# Create old expert indices (with redundancy)
|
||||
redundancy_config = create_redundancy_config(
|
||||
num_logical_experts, total_physical_experts
|
||||
)
|
||||
|
||||
old_indices = create_expert_indices_with_redundancy(
|
||||
num_layers,
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
redundancy_config,
|
||||
)
|
||||
|
||||
# Create new expert indices (with redundancy)
|
||||
new_redundancy_config = create_redundancy_config(
|
||||
num_logical_experts, total_physical_experts
|
||||
)
|
||||
new_indices = create_expert_indices_with_redundancy(
|
||||
num_layers,
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
new_redundancy_config,
|
||||
)
|
||||
|
||||
# Create expert weights
|
||||
expert_weights = create_expert_weights(
|
||||
num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
|
||||
)
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
|
||||
communicator = create_eplb_communicator_or_raise(
|
||||
group_coordinator=ep_group_coordinator,
|
||||
backend=eplb_communicator,
|
||||
expert_weights=expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
|
||||
# Execute weight rearrangement
|
||||
rearrange_expert_weights_inplace(
|
||||
old_indices,
|
||||
new_indices,
|
||||
expert_weights,
|
||||
expert_buffer,
|
||||
ep_group,
|
||||
communicator,
|
||||
)
|
||||
|
||||
# Verify the rearrangement result
|
||||
local_ok = verify_expert_weights_after_shuffle(
|
||||
expert_weights,
|
||||
new_indices,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
num_local_experts,
|
||||
)
|
||||
|
||||
local_ok = (
|
||||
verify_redundant_experts_have_same_weights(
|
||||
expert_weights,
|
||||
new_indices,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
world_size,
|
||||
num_local_experts,
|
||||
)
|
||||
and local_ok
|
||||
)
|
||||
assert_verification_synced(
|
||||
local_ok,
|
||||
"Rearrange verification failed on at least one rank. See logs for details.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"world_size,num_layers,num_local_experts,num_logical_experts",
|
||||
[
|
||||
# 2 GPU, 2 experts per GPU
|
||||
# 3 logical experts, 4 physical experts, 1 redundant experts
|
||||
(2, 1, 2, 3),
|
||||
# 2 GPU, 3 experts per GPU
|
||||
# 4 logical experts, 6 physical experts, 2 redundant experts
|
||||
(2, 2, 3, 4),
|
||||
# 2 GPU, 8 experts per GPU
|
||||
# 16 logical experts, 16 physical experts, 0 redundant experts
|
||||
(2, 4, 8, 16),
|
||||
# 4 GPU, 2 experts per GPU
|
||||
# 6 logical experts, 8 physical experts, 2 redundant experts
|
||||
(4, 1, 2, 6),
|
||||
# 4 GPU, 2 experts per GPU
|
||||
# 5 logical experts, 8 physical experts, 3 redundant experts
|
||||
(4, 2, 2, 5),
|
||||
# 4 GPU, 8 experts per GPU
|
||||
# 16 logical experts, 32 physical experts, 16 redundant experts
|
||||
(4, 8, 8, 16),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"eplb_communicator", ["torch_nccl", "torch_gloo", "pynccl", "nixl"]
|
||||
)
|
||||
def test_rearrange_expert_weights_with_redundancy(
|
||||
world_size,
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
num_logical_experts,
|
||||
eplb_communicator,
|
||||
):
|
||||
"""Test the functionality of rearranging expert weights with redundancy."""
|
||||
|
||||
if eplb_communicator == "nixl" and not has_nixl():
|
||||
pytest.skip("NIXL is not available")
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
distributed_run(
|
||||
_test_rearrange_expert_weights_with_redundancy,
|
||||
world_size,
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
num_logical_experts,
|
||||
eplb_communicator,
|
||||
)
|
||||
|
||||
|
||||
def _test_rearrange_expert_weights_no_change(env, world_size) -> None:
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.tensor_parallel_size = world_size
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group_coordinator = get_tp_group()
|
||||
ep_group = ep_group_coordinator.cpu_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
device = torch.device(f"cuda:{ep_rank}")
|
||||
|
||||
num_layers = 2
|
||||
num_local_experts = 2
|
||||
total_physical_experts = world_size * num_local_experts
|
||||
num_logical_experts = total_physical_experts // 2 # Some redundancy
|
||||
hidden_sizes = [32, 64]
|
||||
|
||||
# Create redundancy configuration
|
||||
redundancy_config = [2] * num_logical_experts
|
||||
|
||||
# Same indices - no change
|
||||
indices = create_expert_indices_with_redundancy(
|
||||
num_layers, num_logical_experts, total_physical_experts, redundancy_config
|
||||
)
|
||||
|
||||
expert_weights = create_expert_weights(
|
||||
num_layers, num_local_experts, hidden_sizes, ep_rank, device, indices
|
||||
)
|
||||
|
||||
# Save original weights
|
||||
original_weights = []
|
||||
for layer_weights in expert_weights:
|
||||
layer_copy = []
|
||||
for weight in layer_weights:
|
||||
layer_copy.append(weight.clone())
|
||||
original_weights.append(layer_copy)
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
|
||||
communicator = create_eplb_communicator_or_raise(
|
||||
group_coordinator=ep_group_coordinator,
|
||||
backend="torch_nccl",
|
||||
expert_weights=expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
|
||||
# Execute rearrangement (should be no change)
|
||||
rearrange_expert_weights_inplace(
|
||||
indices,
|
||||
indices, # Same indices
|
||||
expert_weights,
|
||||
expert_buffer,
|
||||
ep_group,
|
||||
communicator,
|
||||
)
|
||||
|
||||
# Verify that the weights have not changed
|
||||
local_ok = True
|
||||
for layer in range(num_layers):
|
||||
for weight_idx in range(len(hidden_sizes)):
|
||||
if not torch.equal(
|
||||
expert_weights[layer][weight_idx],
|
||||
original_weights[layer][weight_idx],
|
||||
):
|
||||
local_ok = False
|
||||
print(
|
||||
"test_rearrange_expert_weights_no_change failed: "
|
||||
f"layer={layer}, weight_idx={weight_idx}",
|
||||
flush=True,
|
||||
)
|
||||
assert_verification_synced(
|
||||
local_ok,
|
||||
"No-change EPLB verification failed on at least one rank.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"world_size,num_layers,num_local_experts,num_logical_experts",
|
||||
[
|
||||
(2, 2, 2, 3),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("eplb_communicator", ["torch_gloo", "nixl"])
|
||||
def test_async_transfer_layer_without_mtp(
|
||||
world_size: int,
|
||||
num_layers: int,
|
||||
num_local_experts: int,
|
||||
num_logical_experts: int,
|
||||
eplb_communicator: str,
|
||||
):
|
||||
"""Exercise async EPLB transfer path without MTP/spec decode."""
|
||||
|
||||
if eplb_communicator == "nixl" and not has_nixl():
|
||||
pytest.skip("NIXL is not available")
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
|
||||
distributed_run(
|
||||
_test_async_transfer_layer_without_mtp_worker,
|
||||
world_size,
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
num_logical_experts,
|
||||
eplb_communicator,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
def test_rearrange_expert_weights_no_change(world_size):
|
||||
"""
|
||||
Test that when the indices do not change, the weights should remain
|
||||
unchanged.
|
||||
"""
|
||||
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
distributed_run(
|
||||
_test_rearrange_expert_weights_no_change,
|
||||
world_size,
|
||||
)
|
||||
|
||||
|
||||
def _test_rearrange_expert_weights_profile_mode(env, world_size) -> None:
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.tensor_parallel_size = world_size
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group_coordinator = get_tp_group()
|
||||
ep_group = ep_group_coordinator.cpu_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
device = torch.device(f"cuda:{ep_rank}")
|
||||
|
||||
num_layers = 1
|
||||
num_local_experts = 2
|
||||
total_physical_experts = world_size * num_local_experts
|
||||
num_logical_experts = total_physical_experts // 2
|
||||
hidden_sizes = [32]
|
||||
|
||||
# Create different index distributions
|
||||
old_redundancy = create_redundancy_config(
|
||||
num_logical_experts, total_physical_experts
|
||||
)
|
||||
new_redundancy = create_redundancy_config(
|
||||
num_logical_experts, total_physical_experts
|
||||
)
|
||||
|
||||
old_indices = create_expert_indices_with_redundancy(
|
||||
num_layers, num_logical_experts, total_physical_experts, old_redundancy
|
||||
)
|
||||
new_indices = create_expert_indices_with_redundancy(
|
||||
num_layers, num_logical_experts, total_physical_experts, new_redundancy
|
||||
)
|
||||
|
||||
expert_weights = create_expert_weights(
|
||||
num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
|
||||
)
|
||||
|
||||
# Save original weights
|
||||
original_weights = []
|
||||
for layer_weights in expert_weights:
|
||||
layer_copy = []
|
||||
for weight in layer_weights:
|
||||
layer_copy.append(weight.clone())
|
||||
original_weights.append(layer_copy)
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
|
||||
communicator = create_eplb_communicator_or_raise(
|
||||
group_coordinator=ep_group_coordinator,
|
||||
backend="torch_nccl",
|
||||
expert_weights=expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
|
||||
# Execute profile mode rearrangement
|
||||
rearrange_expert_weights_inplace(
|
||||
old_indices,
|
||||
new_indices,
|
||||
expert_weights,
|
||||
expert_buffer,
|
||||
ep_group,
|
||||
communicator,
|
||||
is_profile=True,
|
||||
)
|
||||
|
||||
# In profile mode, the weights should remain unchanged
|
||||
local_ok = True
|
||||
for layer in range(num_layers):
|
||||
for weight_idx in range(len(hidden_sizes)):
|
||||
if not torch.equal(
|
||||
expert_weights[layer][weight_idx],
|
||||
original_weights[layer][weight_idx],
|
||||
):
|
||||
local_ok = False
|
||||
print(
|
||||
"test_rearrange_expert_weights_profile_mode failed: "
|
||||
f"layer={layer}, weight_idx={weight_idx}",
|
||||
flush=True,
|
||||
)
|
||||
assert_verification_synced(
|
||||
local_ok,
|
||||
"Profile-mode EPLB verification failed on at least one rank.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
def test_rearrange_expert_weights_profile_mode(world_size):
|
||||
"""Test profile mode (should not copy actual weights)"""
|
||||
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
distributed_run(
|
||||
_test_rearrange_expert_weights_profile_mode,
|
||||
world_size,
|
||||
)
|
||||
|
||||
|
||||
def _test_nixl_deferred_init_worker(
|
||||
env,
|
||||
world_size: int,
|
||||
num_layers: int,
|
||||
num_local_experts: int,
|
||||
num_logical_experts: int,
|
||||
) -> None:
|
||||
"""Exercise NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""
|
||||
from vllm.distributed.eplb.eplb_communicator import NixlEplbCommunicator
|
||||
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.tensor_parallel_size = world_size
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group_coordinator = get_tp_group()
|
||||
ep_group = ep_group_coordinator.cpu_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
device = torch.device(f"cuda:{ep_rank}")
|
||||
|
||||
total_physical_experts = world_size * num_local_experts
|
||||
hidden_sizes = [32, 64]
|
||||
|
||||
redundancy_config = create_redundancy_config(
|
||||
num_logical_experts, total_physical_experts
|
||||
)
|
||||
old_indices = create_expert_indices_with_redundancy(
|
||||
num_layers,
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
redundancy_config,
|
||||
)
|
||||
|
||||
new_redundancy_config = create_redundancy_config(
|
||||
num_logical_experts, total_physical_experts
|
||||
)
|
||||
new_indices = create_expert_indices_with_redundancy(
|
||||
num_layers,
|
||||
num_logical_experts,
|
||||
total_physical_experts,
|
||||
new_redundancy_config,
|
||||
)
|
||||
|
||||
expert_weights = create_expert_weights(
|
||||
num_layers, num_local_experts, hidden_sizes, ep_rank, device, old_indices
|
||||
)
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in expert_weights[0]]
|
||||
|
||||
communicator = NixlEplbCommunicator(
|
||||
cpu_group=ep_group_coordinator.cpu_group,
|
||||
all_expert_weights=expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
defer_remote_setup=True,
|
||||
)
|
||||
assert not communicator._remote_state_initialized
|
||||
|
||||
rearrange_expert_weights_inplace(
|
||||
old_indices,
|
||||
new_indices,
|
||||
expert_weights,
|
||||
expert_buffer,
|
||||
ep_group,
|
||||
communicator,
|
||||
)
|
||||
|
||||
assert communicator._remote_state_initialized
|
||||
|
||||
local_ok = verify_expert_weights_after_shuffle(
|
||||
expert_weights,
|
||||
new_indices,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
num_local_experts,
|
||||
)
|
||||
|
||||
local_ok = (
|
||||
verify_redundant_experts_have_same_weights(
|
||||
expert_weights,
|
||||
new_indices,
|
||||
hidden_sizes,
|
||||
ep_rank,
|
||||
world_size,
|
||||
num_local_experts,
|
||||
)
|
||||
and local_ok
|
||||
)
|
||||
assert_verification_synced(
|
||||
local_ok,
|
||||
"Deferred NIXL init verification failed on at least one rank.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not has_nixl(), reason="NIXL is not available")
|
||||
@pytest.mark.parametrize(
|
||||
"world_size,num_layers,num_local_experts,num_logical_experts",
|
||||
[(2, 2, 3, 4)],
|
||||
)
|
||||
def test_nixl_deferred_init(
|
||||
world_size,
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
num_logical_experts,
|
||||
):
|
||||
"""Test NixlEplbCommunicator with defer_remote_setup=True (elastic EP path)."""
|
||||
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
distributed_run(
|
||||
_test_nixl_deferred_init_worker,
|
||||
world_size,
|
||||
num_layers,
|
||||
num_local_experts,
|
||||
num_logical_experts,
|
||||
)
|
||||
@@ -0,0 +1,295 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Test that the interaction between EPLB and FusedMoE Layer is okay
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
|
||||
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
get_eplb_group,
|
||||
get_tp_group,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
|
||||
|
||||
from .eplb_utils import distributed_run, set_env_vars_and_device
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestConfig:
|
||||
num_layers: int
|
||||
num_experts: int
|
||||
num_local_experts: int
|
||||
num_topk: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
weight_dtype: torch.dtype
|
||||
weight_scale_dtype: torch.dtype | None
|
||||
column_major_scales: bool
|
||||
|
||||
|
||||
def make_expert_weights(
|
||||
layer_idx: int,
|
||||
global_expert_idx: int,
|
||||
global_num_experts: int,
|
||||
tensor_shape: tuple[int, ...],
|
||||
tensor_dtype: torch.dtype,
|
||||
tensor_device: torch.device,
|
||||
is_column_major: bool,
|
||||
) -> torch.Tensor:
|
||||
assert len(tensor_shape) == 2
|
||||
|
||||
if is_column_major:
|
||||
tensor_shape = (tensor_shape[1], tensor_shape[0])
|
||||
|
||||
x = torch.empty(tensor_shape, dtype=tensor_dtype, device=tensor_device)
|
||||
value_offset = (layer_idx * global_num_experts + global_expert_idx) * x.numel()
|
||||
x.view(-1).copy_(
|
||||
torch.arange(
|
||||
value_offset,
|
||||
value_offset + x.numel(),
|
||||
dtype=tensor_dtype,
|
||||
device=tensor_device,
|
||||
)
|
||||
)
|
||||
|
||||
if is_column_major:
|
||||
x = torch.transpose(x, 1, 0)
|
||||
assert not x.is_contiguous()
|
||||
return x
|
||||
|
||||
|
||||
def make_fused_moe_layer(
|
||||
rank: int,
|
||||
layer_idx: int,
|
||||
test_config: TestConfig,
|
||||
) -> FusedMoE:
|
||||
fml = FusedMoE(
|
||||
num_experts=test_config.num_experts,
|
||||
top_k=test_config.num_topk,
|
||||
hidden_size=test_config.hidden_size,
|
||||
intermediate_size=test_config.intermediate_size,
|
||||
prefix=f"dummy_layer_{layer_idx}",
|
||||
activation="silu",
|
||||
params_dtype=test_config.weight_dtype,
|
||||
)
|
||||
re = fml.routed_experts
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
|
||||
from functools import partial
|
||||
|
||||
_make_expert_weights = partial(
|
||||
make_expert_weights,
|
||||
layer_idx=layer_idx,
|
||||
global_num_experts=test_config.num_experts,
|
||||
tensor_device=device,
|
||||
)
|
||||
|
||||
assert isinstance(re.w13_weight.data, torch.Tensor)
|
||||
assert isinstance(re.w2_weight.data, torch.Tensor)
|
||||
re.w13_weight.data = re.w13_weight.data.to(device=device)
|
||||
re.w2_weight.data = re.w2_weight.data.to(device=device)
|
||||
w13_weight = re.w13_weight.data
|
||||
w2_weight = re.w2_weight.data
|
||||
assert w13_weight.size(0) == test_config.num_local_experts
|
||||
for i in range(test_config.num_local_experts):
|
||||
g_i = rank * test_config.num_local_experts + i
|
||||
w13_weight_e = w13_weight[i]
|
||||
w2_weight_e = w2_weight[i]
|
||||
w13_weight_e.copy_(
|
||||
_make_expert_weights(
|
||||
global_expert_idx=g_i,
|
||||
tensor_shape=w13_weight_e.shape,
|
||||
tensor_dtype=w13_weight_e.dtype,
|
||||
is_column_major=False,
|
||||
)
|
||||
)
|
||||
w2_weight_e.copy_(
|
||||
_make_expert_weights(
|
||||
global_expert_idx=g_i,
|
||||
tensor_shape=w2_weight_e.shape,
|
||||
tensor_dtype=w2_weight_e.dtype,
|
||||
is_column_major=False,
|
||||
)
|
||||
)
|
||||
|
||||
block_size = 16
|
||||
|
||||
def block_quant_scales_shape(
|
||||
shape: tuple[int, ...], is_column_major: bool
|
||||
) -> tuple[int, ...]:
|
||||
assert len(shape) == 3
|
||||
if not is_column_major:
|
||||
return (shape[0], shape[1] // block_size, shape[2] // block_size)
|
||||
else:
|
||||
return (shape[0], shape[2] // block_size, shape[1] // block_size)
|
||||
|
||||
is_column_major = test_config.column_major_scales
|
||||
w13_weight_scale_inv = torch.empty(
|
||||
block_quant_scales_shape(w13_weight.shape, is_column_major),
|
||||
dtype=test_config.weight_dtype,
|
||||
device=device,
|
||||
)
|
||||
w2_weight_scale_inv = torch.empty(
|
||||
block_quant_scales_shape(w2_weight.shape, is_column_major),
|
||||
dtype=test_config.weight_dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
for i in range(test_config.num_local_experts):
|
||||
g_i = rank * test_config.num_local_experts + i
|
||||
w13_s_e = w13_weight_scale_inv[i]
|
||||
w2_s_e = w2_weight_scale_inv[i]
|
||||
w13_s_e.copy_(
|
||||
_make_expert_weights(
|
||||
global_expert_idx=g_i,
|
||||
tensor_shape=w13_s_e.shape,
|
||||
tensor_dtype=w13_s_e.dtype,
|
||||
# Fill data in row-major and then
|
||||
# transpose if test_config requires col-major.
|
||||
is_column_major=False,
|
||||
)
|
||||
)
|
||||
w2_s_e.copy_(
|
||||
_make_expert_weights(
|
||||
global_expert_idx=g_i,
|
||||
tensor_shape=w2_s_e.shape,
|
||||
tensor_dtype=w2_s_e.dtype,
|
||||
is_column_major=False,
|
||||
)
|
||||
)
|
||||
if is_column_major:
|
||||
w13_weight_scale_inv = torch.transpose(w13_weight_scale_inv, 1, 2)
|
||||
w2_weight_scale_inv = torch.transpose(w2_weight_scale_inv, 1, 2)
|
||||
assert not w13_weight_scale_inv.is_contiguous()
|
||||
assert not w2_weight_scale_inv.is_contiguous()
|
||||
|
||||
# Add scales to the parameter list
|
||||
re.w13_weight_scale_inv = torch.nn.Parameter(
|
||||
w13_weight_scale_inv, requires_grad=False
|
||||
)
|
||||
re.w2_weight_scale_inv = torch.nn.Parameter(
|
||||
w2_weight_scale_inv, requires_grad=False
|
||||
)
|
||||
|
||||
return fml
|
||||
|
||||
|
||||
def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
|
||||
# Initialize model parallel (using tensor parallel as an entrypoint
|
||||
# to expert parallel)
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.tensor_parallel_size = world_size
|
||||
vllm_config.parallel_config.enable_expert_parallel = True
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=world_size, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group = get_tp_group().cpu_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
|
||||
fml_layers = [
|
||||
make_fused_moe_layer(ep_rank, layer_idx, test_config)
|
||||
for layer_idx in range(test_config.num_layers)
|
||||
]
|
||||
rank_expert_weights = [fml.get_expert_weights() for fml in fml_layers]
|
||||
|
||||
indices = torch.zeros(
|
||||
test_config.num_layers, test_config.num_experts, dtype=torch.long
|
||||
)
|
||||
for lidx in range(test_config.num_layers):
|
||||
indices[lidx] = torch.Tensor(range(test_config.num_experts))
|
||||
|
||||
shuffled_indices = torch.zeros_like(indices)
|
||||
for lidx in range(test_config.num_layers):
|
||||
shuffled_indices[lidx] = torch.randperm(test_config.num_experts)
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]]
|
||||
communicator = create_eplb_communicator(
|
||||
group_coordinator=get_eplb_group(),
|
||||
backend="torch_nccl",
|
||||
expert_weights=rank_expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
rearrange_expert_weights_inplace(
|
||||
indices,
|
||||
shuffled_indices,
|
||||
rank_expert_weights,
|
||||
expert_buffer,
|
||||
ep_group,
|
||||
communicator,
|
||||
)
|
||||
|
||||
num_local_experts = test_config.num_local_experts
|
||||
num_global_experts = test_config.num_experts
|
||||
for lidx, fml in enumerate(fml_layers):
|
||||
for name, w in fml.named_parameters():
|
||||
for e in range(num_local_experts):
|
||||
g_e = shuffled_indices[lidx][ep_rank * num_local_experts + e]
|
||||
ref = make_expert_weights(
|
||||
layer_idx=lidx,
|
||||
global_expert_idx=int(g_e.item()),
|
||||
global_num_experts=num_global_experts,
|
||||
tensor_shape=w[e].shape,
|
||||
tensor_dtype=w[e].dtype,
|
||||
tensor_device=w[e].device,
|
||||
is_column_major=not w[e].is_contiguous(),
|
||||
)
|
||||
assert w[e].shape == ref.shape and w[e].stride() == ref.stride(), (
|
||||
f"w[{e}] {w[e].size()} {w[e].stride()} vs "
|
||||
f"ref {ref.size()} {ref.stride()}"
|
||||
)
|
||||
torch.testing.assert_close(w[e], ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.parametrize("num_layers", [4])
|
||||
@pytest.mark.parametrize("num_experts", [16])
|
||||
@pytest.mark.parametrize("hidden_size", [256])
|
||||
@pytest.mark.parametrize("intermediate_size", [256])
|
||||
@pytest.mark.parametrize("column_major_scales", [True, False])
|
||||
def test_eplb_fml(
|
||||
world_size: int,
|
||||
num_layers: int,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
column_major_scales: bool,
|
||||
):
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
|
||||
num_local_experts = num_experts // world_size
|
||||
num_topk = 4
|
||||
# The dtypes are fine as we are essentially just checking data-copies
|
||||
weight_dtype = torch.bfloat16
|
||||
weight_scale_dtype = torch.bfloat16
|
||||
|
||||
test_config = TestConfig(
|
||||
num_layers=num_layers,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
num_topk=num_topk,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
weight_dtype=weight_dtype,
|
||||
weight_scale_dtype=weight_scale_dtype,
|
||||
column_major_scales=column_major_scales,
|
||||
)
|
||||
|
||||
distributed_run(
|
||||
_test_eplb_fml,
|
||||
world_size,
|
||||
test_config,
|
||||
)
|
||||
@@ -0,0 +1,292 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Test that the interaction between EPLB and FusedMoE Layer is okay for DP w/ NVFP4
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.moe.utils import make_test_quant_config
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
|
||||
from vllm.distributed.eplb.eplb_state import EplbLayerState
|
||||
from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
get_dp_group,
|
||||
get_eplb_group,
|
||||
)
|
||||
from vllm.forward_context import set_forward_context
|
||||
from vllm.model_executor.layers.fused_moe.layer import FusedMoE
|
||||
from vllm.model_executor.layers.quantization.modelopt import (
|
||||
ModelOptNvFp4Config,
|
||||
ModelOptNvFp4FusedMoE,
|
||||
)
|
||||
|
||||
from .eplb_utils import distributed_run, set_env_vars_and_device
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestConfig:
|
||||
num_layers: int
|
||||
num_experts: int
|
||||
num_local_experts: int
|
||||
num_topk: int
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_tokens: int
|
||||
moe_backend: str
|
||||
|
||||
|
||||
def make_fused_moe_layer(
|
||||
rank: int,
|
||||
layer_idx: int,
|
||||
test_config: TestConfig,
|
||||
) -> FusedMoE:
|
||||
quant_config = None
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
|
||||
quant_config = ModelOptNvFp4Config(
|
||||
is_checkpoint_nvfp4_serialized=True,
|
||||
kv_cache_quant_algo=None,
|
||||
exclude_modules=[],
|
||||
)
|
||||
|
||||
fml = FusedMoE(
|
||||
num_experts=test_config.num_experts,
|
||||
top_k=test_config.num_topk,
|
||||
hidden_size=test_config.hidden_size,
|
||||
intermediate_size=test_config.intermediate_size,
|
||||
prefix=f"dummy_layer_{layer_idx}",
|
||||
activation="silu",
|
||||
params_dtype=torch.bfloat16,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
nvfp4_fused_moe = ModelOptNvFp4FusedMoE(quant_config, fml)
|
||||
nvfp4_fused_moe.create_weights(
|
||||
fml,
|
||||
test_config.num_local_experts,
|
||||
test_config.hidden_size,
|
||||
test_config.intermediate_size,
|
||||
params_dtype=torch.uint8,
|
||||
global_num_experts=test_config.num_experts,
|
||||
)
|
||||
|
||||
fml = fml.to(device)
|
||||
re = fml.routed_experts
|
||||
w1_q, w2_q, quant_config = make_test_quant_config(
|
||||
test_config.num_local_experts,
|
||||
test_config.intermediate_size,
|
||||
test_config.hidden_size,
|
||||
in_dtype=torch.bfloat16,
|
||||
quant_dtype="nvfp4",
|
||||
block_shape=None,
|
||||
per_act_token_quant=False,
|
||||
)
|
||||
|
||||
re.w13_weight.data = w1_q
|
||||
re.w2_weight.data = w2_q
|
||||
|
||||
re.w2_input_scale.data = torch.randn_like(re.w2_input_scale.data) / 5
|
||||
re.w13_input_scale.data = torch.randn_like(re.w13_input_scale.data) / 5
|
||||
re.w2_weight_scale_2.data = torch.randn_like(re.w2_weight_scale_2.data) / 5
|
||||
re.w13_weight_scale_2.data = torch.randn_like(re.w13_weight_scale_2.data) / 5
|
||||
re.w2_weight_scale.data = (
|
||||
torch.randn(re.w2_weight_scale.data.shape, device=device) / 5
|
||||
).to(re.w2_weight_scale.data.dtype)
|
||||
re.w13_weight_scale.data = (
|
||||
torch.randn(re.w13_weight_scale.data.shape, device=device) / 5
|
||||
).to(re.w13_weight_scale.data.dtype)
|
||||
|
||||
nvfp4_fused_moe.process_weights_after_loading(re)
|
||||
|
||||
fml.maybe_init_modular_kernel()
|
||||
|
||||
return fml
|
||||
|
||||
|
||||
def _test_eplb_fml(env, world_size: int, test_config: TestConfig):
|
||||
set_env_vars_and_device(env)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config.data_parallel_size = world_size
|
||||
vllm_config.parallel_config.enable_expert_parallel = True
|
||||
vllm_config.kernel_config.moe_backend = test_config.moe_backend
|
||||
|
||||
with set_current_vllm_config(vllm_config):
|
||||
ensure_model_parallel_initialized(
|
||||
tensor_model_parallel_size=1, pipeline_model_parallel_size=1
|
||||
)
|
||||
|
||||
ep_group = get_dp_group().cpu_group
|
||||
ep_rank = torch.distributed.get_rank()
|
||||
|
||||
device = torch.device(f"cuda:{ep_rank}")
|
||||
|
||||
fml_layers = [
|
||||
make_fused_moe_layer(ep_rank, layer_idx, test_config).to(device)
|
||||
for layer_idx in range(test_config.num_layers)
|
||||
]
|
||||
rank_expert_weights = [fml.get_expert_weights() for fml in fml_layers]
|
||||
|
||||
hidden_states = []
|
||||
router_logits = []
|
||||
for layer_idx in range(test_config.num_layers):
|
||||
hidden_states.append(
|
||||
torch.randn(
|
||||
(test_config.num_tokens, test_config.hidden_size),
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
router_logits.append(
|
||||
torch.randn(
|
||||
(test_config.num_tokens, test_config.num_experts),
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
)
|
||||
|
||||
out_before_shuffle = []
|
||||
with set_forward_context(
|
||||
{},
|
||||
num_tokens=test_config.num_tokens,
|
||||
num_tokens_across_dp=torch.tensor(
|
||||
[test_config.num_tokens] * world_size, device="cpu", dtype=torch.int
|
||||
),
|
||||
vllm_config=vllm_config,
|
||||
):
|
||||
for lidx, fml in enumerate(fml_layers):
|
||||
out_before_shuffle.append(
|
||||
fml(hidden_states[lidx].clone(), router_logits[lidx].clone())
|
||||
)
|
||||
|
||||
indices = torch.zeros(
|
||||
test_config.num_layers, test_config.num_experts, dtype=torch.long
|
||||
)
|
||||
for lidx in range(test_config.num_layers):
|
||||
indices[lidx] = torch.Tensor(range(test_config.num_experts))
|
||||
|
||||
shuffled_indices = torch.zeros_like(indices)
|
||||
for lidx in range(test_config.num_layers):
|
||||
shuffled_indices[lidx] = torch.randperm(test_config.num_experts)
|
||||
|
||||
expert_buffer = [torch.empty_like(w) for w in rank_expert_weights[0]]
|
||||
communicator = create_eplb_communicator(
|
||||
group_coordinator=get_eplb_group(),
|
||||
backend="torch_nccl",
|
||||
expert_weights=rank_expert_weights,
|
||||
expert_buffer=expert_buffer,
|
||||
)
|
||||
rearrange_expert_weights_inplace(
|
||||
indices,
|
||||
shuffled_indices,
|
||||
rank_expert_weights,
|
||||
expert_buffer,
|
||||
ep_group,
|
||||
communicator,
|
||||
)
|
||||
|
||||
num_global_experts = test_config.num_experts
|
||||
|
||||
logical_to_physical_map_list = []
|
||||
for lidx, fml in enumerate(fml_layers):
|
||||
physical_to_logical_map = shuffled_indices[lidx].to(device)
|
||||
logical_to_physical_map = torch.empty(
|
||||
(num_global_experts,), dtype=torch.int32, device=device
|
||||
)
|
||||
logical_to_physical_map[physical_to_logical_map] = torch.arange(
|
||||
0, num_global_experts, dtype=torch.int32, device=device
|
||||
)
|
||||
logical_to_physical_map_list.append(
|
||||
logical_to_physical_map.reshape(num_global_experts, 1)
|
||||
)
|
||||
|
||||
logical_to_physical_map = torch.stack(logical_to_physical_map_list)
|
||||
|
||||
for lidx, fml in enumerate(fml_layers):
|
||||
logical_replica_count = torch.ones(
|
||||
(test_config.num_layers, num_global_experts),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
fml.eplb_state = EplbLayerState()
|
||||
fml.set_eplb_state(
|
||||
lidx,
|
||||
torch.zeros(
|
||||
(test_config.num_layers, num_global_experts),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
),
|
||||
logical_to_physical_map,
|
||||
logical_replica_count,
|
||||
)
|
||||
fml.router.eplb_state.should_record_tensor = torch.ones(
|
||||
(), dtype=torch.bool, device=device
|
||||
)
|
||||
fml.router.eplb_state.num_unpadded_tokens_tensors = [
|
||||
torch.tensor(0, dtype=torch.int32, device=device)
|
||||
]
|
||||
|
||||
out_after_shuffle = []
|
||||
with set_forward_context(
|
||||
{},
|
||||
num_tokens=test_config.num_tokens,
|
||||
num_tokens_across_dp=torch.tensor(
|
||||
[test_config.num_tokens] * world_size, device="cpu", dtype=torch.int
|
||||
),
|
||||
vllm_config=vllm_config,
|
||||
):
|
||||
for lidx, fml in enumerate(fml_layers):
|
||||
out_after_shuffle.append(
|
||||
fml(hidden_states[lidx].clone(), router_logits[lidx].clone())
|
||||
)
|
||||
|
||||
for lidx in range(test_config.num_layers):
|
||||
torch.testing.assert_close(
|
||||
out_before_shuffle[lidx], out_after_shuffle[lidx], atol=1e-1, rtol=1e-1
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
@pytest.mark.parametrize("num_layers", [8])
|
||||
@pytest.mark.parametrize("num_experts", [32])
|
||||
@pytest.mark.parametrize("hidden_size", [256])
|
||||
@pytest.mark.parametrize("intermediate_size", [256])
|
||||
@pytest.mark.parametrize("num_tokens", [256])
|
||||
@pytest.mark.parametrize("moe_backend", ["flashinfer_trtllm", "flashinfer_cutlass"])
|
||||
def test_eplb_fml(
|
||||
world_size: int,
|
||||
num_layers: int,
|
||||
num_experts: int,
|
||||
hidden_size: int,
|
||||
intermediate_size: int,
|
||||
num_tokens: int,
|
||||
moe_backend: str,
|
||||
):
|
||||
if torch.accelerator.device_count() < world_size:
|
||||
pytest.skip(f"Need at least {world_size} GPUs to run the test")
|
||||
|
||||
num_local_experts = num_experts // world_size
|
||||
num_topk = 4
|
||||
|
||||
test_config = TestConfig(
|
||||
num_layers=num_layers,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
num_topk=num_topk,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size=intermediate_size,
|
||||
num_tokens=num_tokens,
|
||||
moe_backend=moe_backend,
|
||||
)
|
||||
|
||||
distributed_run(
|
||||
_test_eplb_fml,
|
||||
world_size,
|
||||
test_config,
|
||||
)
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from __future__ import annotations
|
||||
|
||||
import lm_eval
|
||||
import pytest
|
||||
|
||||
from tests.utils import large_gpu_mark
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def get_model_args(
|
||||
model_name: str,
|
||||
spec_model_name: str | None,
|
||||
spec_method: str,
|
||||
tp_size: int,
|
||||
model_max_len: int,
|
||||
use_async: bool = True,
|
||||
) -> dict:
|
||||
speculative_config = {
|
||||
"method": spec_method,
|
||||
"model": spec_model_name,
|
||||
"num_speculative_tokens": 1,
|
||||
"max_model_len": model_max_len,
|
||||
}
|
||||
eplb_config = {
|
||||
"num_redundant_experts": tp_size,
|
||||
"window_size": 128,
|
||||
"step_interval": 1024,
|
||||
"log_balancedness": False,
|
||||
"use_async": use_async,
|
||||
}
|
||||
model_args = {
|
||||
"pretrained": model_name,
|
||||
"dtype": "auto",
|
||||
"add_bos_token": True,
|
||||
"tensor_parallel_size": tp_size,
|
||||
"gpu_memory_utilization": 0.7,
|
||||
"speculative_config": speculative_config,
|
||||
"enable_expert_parallel": True,
|
||||
"eplb_config": eplb_config,
|
||||
"enable_eplb": True,
|
||||
"max_model_len": model_max_len,
|
||||
}
|
||||
return model_args
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
current_platform.is_rocm(),
|
||||
reason="EPLB with Spec Decode is a work in progress on ROCm.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_setup",
|
||||
[
|
||||
pytest.param(
|
||||
("mtp", "Qwen/Qwen3-Next-80B-A3B-Instruct", None, 4, 0.86),
|
||||
marks=large_gpu_mark(min_gb=80),
|
||||
),
|
||||
pytest.param(
|
||||
(
|
||||
"eagle",
|
||||
"meta-llama/Llama-4-Scout-17B-16E-Instruct",
|
||||
"morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct",
|
||||
4,
|
||||
0.92,
|
||||
),
|
||||
marks=pytest.mark.skip(reason="Skipping due to CI OOM issues"),
|
||||
),
|
||||
],
|
||||
ids=["qwen3_next_mtp", "llama4_eagle"],
|
||||
)
|
||||
def test_eplb_spec_decode(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
model_setup: tuple[str, str, str, int, float],
|
||||
):
|
||||
"""
|
||||
Test the correctness of EPLB speculative decoding with GSM8K dataset.
|
||||
Applicable to MoE models with mtp or eagle spec decode.
|
||||
"""
|
||||
method, model_name, spec_model_name, tp_size, expected_gsm8k_value = model_setup
|
||||
|
||||
TASK = "gsm8k"
|
||||
FILTER = "exact_match,strict-match"
|
||||
RTOL = 0.03
|
||||
|
||||
model_args = get_model_args(
|
||||
model_name=model_name,
|
||||
spec_model_name=spec_model_name,
|
||||
spec_method=method,
|
||||
tp_size=tp_size,
|
||||
model_max_len=4096,
|
||||
)
|
||||
|
||||
results = lm_eval.simple_evaluate(
|
||||
model="vllm",
|
||||
model_args=model_args,
|
||||
tasks=TASK,
|
||||
batch_size=64,
|
||||
num_fewshot=8,
|
||||
)
|
||||
measured_value = results["results"][TASK][FILTER]
|
||||
assert (
|
||||
measured_value - RTOL < expected_gsm8k_value
|
||||
and measured_value + RTOL > expected_gsm8k_value
|
||||
), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"
|
||||
|
||||
|
||||
@large_gpu_mark(min_gb=80)
|
||||
def test_eplb_spec_decode_qwen3_next_mtp_async() -> None:
|
||||
"""
|
||||
Ensure async EPLB works with MTP speculative decoding for Qwen3-Next.
|
||||
"""
|
||||
|
||||
TASK = "gsm8k"
|
||||
FILTER = "exact_match,strict-match"
|
||||
RTOL = 0.03
|
||||
expected_gsm8k_value = 0.86
|
||||
|
||||
model_args = get_model_args(
|
||||
model_name="Qwen/Qwen3-Next-80B-A3B-Instruct",
|
||||
spec_model_name=None,
|
||||
spec_method="mtp",
|
||||
tp_size=4,
|
||||
model_max_len=4096,
|
||||
use_async=True,
|
||||
)
|
||||
|
||||
results = lm_eval.simple_evaluate(
|
||||
model="vllm",
|
||||
model_args=model_args,
|
||||
tasks=TASK,
|
||||
batch_size=64,
|
||||
num_fewshot=8,
|
||||
)
|
||||
measured_value = results["results"][TASK][FILTER]
|
||||
assert (
|
||||
measured_value - RTOL < expected_gsm8k_value
|
||||
and measured_value + RTOL > expected_gsm8k_value
|
||||
), f"Expected: {expected_gsm8k_value} | Measured: {measured_value}"
|
||||
@@ -0,0 +1,154 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.distributed.eplb.eplb_state import (
|
||||
_commit_eplb_maps,
|
||||
_commit_eplb_maps_for_layer,
|
||||
)
|
||||
|
||||
|
||||
def _make_model_state(
|
||||
phy2log: torch.Tensor,
|
||||
log2phy: torch.Tensor,
|
||||
logcnt: torch.Tensor,
|
||||
) -> MagicMock:
|
||||
"""Build a minimal EplbModelState mock with only the three map tensors."""
|
||||
state = MagicMock()
|
||||
state.physical_to_logical_map = phy2log
|
||||
state.logical_to_physical_map = log2phy
|
||||
state.logical_replica_count = logcnt
|
||||
return state
|
||||
|
||||
|
||||
def test_commit_eplb_maps_shape_change():
|
||||
"""
|
||||
The normal path copies the physical_to_logical map in-place. When the number of
|
||||
physical experts changes, the old map should be replaced entirely.
|
||||
"""
|
||||
num_layers, num_logical, num_physical = 2, 4, 6
|
||||
max_replicas = 3
|
||||
|
||||
# Build current state tensors
|
||||
model_state = _make_model_state(
|
||||
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
|
||||
log2phy=torch.full(
|
||||
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
|
||||
),
|
||||
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
|
||||
)
|
||||
|
||||
# The new map has two more physical experts. These new physical experts will
|
||||
# automatically map to the first two logical experts
|
||||
new_phy2log_larger = (
|
||||
(torch.arange(num_physical + 2, dtype=torch.long) % num_logical)
|
||||
.unsqueeze(0)
|
||||
.expand(num_layers, -1)
|
||||
)
|
||||
_commit_eplb_maps(model_state, new_phy2log_larger)
|
||||
|
||||
# Check that the number of physical experts has been updated and that the values
|
||||
# match
|
||||
assert model_state.physical_to_logical_map.shape[1] == num_physical + 2
|
||||
assert torch.equal(model_state.physical_to_logical_map, new_phy2log_larger)
|
||||
|
||||
|
||||
def test_commit_eplb_maps_for_layer_logical_padding():
|
||||
"""
|
||||
Test that logical_to_physical_map is padded with -1 to fill the
|
||||
pre-allocated slots when the new map has fewer replicas than the max.
|
||||
"""
|
||||
num_layers, num_logical, num_physical = 2, 4, 6
|
||||
max_replicas = 3
|
||||
|
||||
model_state = _make_model_state(
|
||||
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
|
||||
log2phy=torch.full(
|
||||
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
|
||||
),
|
||||
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
|
||||
)
|
||||
|
||||
new_phy2log = (
|
||||
(torch.arange(num_physical, dtype=torch.long) % num_logical)
|
||||
.unsqueeze(0)
|
||||
.expand(num_layers, -1)
|
||||
.contiguous()
|
||||
)
|
||||
layer = 0
|
||||
_commit_eplb_maps_for_layer(model_state, new_phy2log[layer], layer)
|
||||
|
||||
assert torch.all(model_state.logical_to_physical_map[layer, :, 2] == -1)
|
||||
|
||||
|
||||
def test_commit_eplb_maps_for_layer_shape_assert():
|
||||
"""Test that a mismatched number of physical experts triggers an assertion error."""
|
||||
num_layers, num_logical, num_physical = 2, 4, 6
|
||||
|
||||
model_state = _make_model_state(
|
||||
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
|
||||
log2phy=torch.full((num_layers, num_logical, 2), -1, dtype=torch.long),
|
||||
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
|
||||
)
|
||||
bad_phy2log = torch.zeros(num_layers, num_physical + 1, dtype=torch.long)
|
||||
with pytest.raises(AssertionError):
|
||||
_commit_eplb_maps_for_layer(model_state, bad_phy2log, layer=0)
|
||||
|
||||
|
||||
def test_commit_eplb_maps():
|
||||
"""Test that all values are copied correctly into model_state."""
|
||||
num_layers, num_logical, num_physical, max_replicas = 2, 3, 4, 2
|
||||
|
||||
model_state = _make_model_state(
|
||||
phy2log=torch.zeros(num_layers, num_physical, dtype=torch.long),
|
||||
log2phy=torch.full(
|
||||
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
|
||||
),
|
||||
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
|
||||
)
|
||||
|
||||
new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
|
||||
new_log2phy = torch.tensor(
|
||||
[[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
|
||||
)
|
||||
new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
|
||||
|
||||
_commit_eplb_maps(model_state, new_phy2log)
|
||||
|
||||
assert torch.equal(model_state.physical_to_logical_map, new_phy2log)
|
||||
assert torch.equal(model_state.logical_to_physical_map, new_log2phy)
|
||||
assert torch.equal(model_state.logical_replica_count, new_logcnt)
|
||||
|
||||
|
||||
def test_commit_eplb_maps_for_layer():
|
||||
"""Test that only the target layer is updated"""
|
||||
num_layers, num_logical, max_replicas = 2, 3, 2
|
||||
|
||||
original_phy2log = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
|
||||
model_state = _make_model_state(
|
||||
phy2log=original_phy2log.clone(),
|
||||
log2phy=torch.full(
|
||||
(num_layers, num_logical, max_replicas), -1, dtype=torch.long
|
||||
),
|
||||
logcnt=torch.zeros(num_layers, num_logical, dtype=torch.long),
|
||||
)
|
||||
|
||||
new_phy2log = torch.tensor([[0, 1, 2, 0], [1, 2, 0, 1]], dtype=torch.long)
|
||||
new_log2phy = torch.tensor(
|
||||
[[[0, 3], [1, -1], [2, -1]], [[2, -1], [0, 3], [1, -1]]], dtype=torch.long
|
||||
)
|
||||
new_logcnt = torch.tensor([[2, 1, 1], [1, 2, 1]], dtype=torch.long)
|
||||
|
||||
_commit_eplb_maps_for_layer(model_state, new_phy2log[0], layer=0)
|
||||
|
||||
# Layer 0 updated
|
||||
assert torch.equal(model_state.physical_to_logical_map[0], new_phy2log[0])
|
||||
assert torch.equal(model_state.logical_to_physical_map[0], new_log2phy[0])
|
||||
assert torch.equal(model_state.logical_replica_count[0], new_logcnt[0])
|
||||
|
||||
# Layer 1 untouched
|
||||
assert torch.equal(model_state.physical_to_logical_map[1], original_phy2log[1])
|
||||
@@ -0,0 +1,333 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import threading
|
||||
import time
|
||||
|
||||
import msgspec
|
||||
import pytest
|
||||
|
||||
from vllm.distributed.kv_events import (
|
||||
EventBatch,
|
||||
EventPublisherFactory,
|
||||
NullEventPublisher,
|
||||
)
|
||||
|
||||
DP_RANK = 0
|
||||
|
||||
|
||||
class EventSample(
|
||||
msgspec.Struct,
|
||||
tag=True, # type: ignore
|
||||
array_like=True, # type: ignore
|
||||
):
|
||||
"""Test event for publisher testing"""
|
||||
|
||||
id: int
|
||||
value: str
|
||||
|
||||
|
||||
class SampleBatch(EventBatch):
|
||||
"""Test event batch for publisher testing"""
|
||||
|
||||
events: list[EventSample]
|
||||
|
||||
|
||||
def create_test_events(count: int) -> SampleBatch:
|
||||
"""Create a batch of test events"""
|
||||
events = [EventSample(id=i, value=f"test-{i}") for i in range(count)]
|
||||
return SampleBatch(ts=time.time(), events=events)
|
||||
|
||||
|
||||
def test_basic_publishing(publisher, subscriber):
|
||||
"""Test basic event publishing works"""
|
||||
|
||||
test_batch = create_test_events(5)
|
||||
publisher.publish(test_batch)
|
||||
|
||||
result = subscriber.receive_one(timeout=1000)
|
||||
assert result is not None, "No message received"
|
||||
|
||||
seq, received = result
|
||||
assert seq == 0, "Sequence number mismatch"
|
||||
assert received.ts == pytest.approx(test_batch.ts, abs=0.1), "Timestamp mismatch"
|
||||
assert len(received.events) == len(test_batch.events), "Number of events mismatch"
|
||||
|
||||
for i, event in enumerate(received.events):
|
||||
assert event.id == i, "Event id mismatch"
|
||||
assert event.value == f"test-{i}", "Event value mismatch"
|
||||
|
||||
|
||||
def test_multiple_events(publisher, subscriber):
|
||||
"""Test publishing and receiving multiple event batches"""
|
||||
for _ in range(10):
|
||||
batch = create_test_events(2)
|
||||
publisher.publish(batch)
|
||||
|
||||
received = []
|
||||
for _ in range(10):
|
||||
data = subscriber.receive_one(timeout=100)
|
||||
if data:
|
||||
received.append(data)
|
||||
|
||||
assert len(received) == 10, "Number of messages mismatch"
|
||||
seqs = [seq for seq, _ in received]
|
||||
assert seqs == list(range(10)), "Sequence numbers mismatch"
|
||||
|
||||
|
||||
def test_replay_mechanism(publisher, subscriber):
|
||||
"""Test the replay mechanism works correctly"""
|
||||
for _ in range(19):
|
||||
batch = create_test_events(1)
|
||||
publisher.publish(batch)
|
||||
|
||||
# Drain live events to ensure publisher has buffered them.
|
||||
for _ in range(19):
|
||||
assert subscriber.receive_one(timeout=1000) is not None
|
||||
|
||||
subscriber.request_replay(10)
|
||||
|
||||
replayed = subscriber.receive_replay()
|
||||
|
||||
assert len(replayed) == 9, (
|
||||
f"Expected 9 replayed messages (seq 10-18), got {len(replayed)}"
|
||||
)
|
||||
seqs = [seq for seq, _ in replayed]
|
||||
assert seqs == list(range(10, 19)), "Replayed sequences should be 10-18"
|
||||
|
||||
|
||||
def test_replay_includes_topic(publisher, subscriber, publisher_config):
|
||||
"""Test that replay responses include the topic, matching PUB format"""
|
||||
for _ in range(5):
|
||||
publisher.publish(create_test_events(1))
|
||||
|
||||
# Drain live events to ensure publisher has processed them.
|
||||
for _ in range(5):
|
||||
assert subscriber.receive_one(timeout=1000) is not None
|
||||
|
||||
subscriber.request_replay(0)
|
||||
|
||||
# receive_replay unpacks (topic, seq, payload) and asserts
|
||||
# topic == publisher topic for each message.
|
||||
replayed = subscriber.receive_replay()
|
||||
assert len(replayed) == 5, f"Expected 5 replayed messages, got {len(replayed)}"
|
||||
seqs = [seq for seq, _ in replayed]
|
||||
assert seqs == list(range(5)), "Replayed sequences should be 0-4"
|
||||
|
||||
|
||||
def test_buffer_limit(publisher, subscriber, publisher_config):
|
||||
"""Test buffer limit behavior"""
|
||||
buffer_size = publisher_config.buffer_steps
|
||||
|
||||
# Publish more events than the buffer can hold
|
||||
for i in range(buffer_size + 10):
|
||||
batch = create_test_events(1)
|
||||
publisher.publish(batch)
|
||||
|
||||
time.sleep(0.5) # Need publisher to process above requests
|
||||
subscriber.request_replay(0)
|
||||
|
||||
replayed = subscriber.receive_replay()
|
||||
|
||||
assert len(replayed) == buffer_size, (
|
||||
f"Expected {buffer_size} replayed messages, got {len(replayed)}"
|
||||
)
|
||||
|
||||
seqs = [seq for seq, _ in replayed]
|
||||
assert seqs == list(range(10, buffer_size + 10)), (
|
||||
"Should replay seq 11 through buffer_size+10"
|
||||
)
|
||||
|
||||
|
||||
def test_topic_filtering(publisher_config):
|
||||
"""
|
||||
Test that a subscriber only receives messages matching its topic filter
|
||||
"""
|
||||
publisher_config.replay_endpoint = None
|
||||
|
||||
publisher_config.topic = "foo"
|
||||
pub = EventPublisherFactory.create(publisher_config, DP_RANK)
|
||||
|
||||
from .conftest import MockSubscriber
|
||||
|
||||
sub_foo = MockSubscriber(publisher_config.endpoint, None, "foo")
|
||||
sub_bar = MockSubscriber(publisher_config.endpoint, None, "bar")
|
||||
|
||||
try:
|
||||
time.sleep(0.1)
|
||||
|
||||
for _ in range(3):
|
||||
pub.publish(create_test_events(1))
|
||||
|
||||
foo_received = [sub_foo.receive_one(timeout=200) for _ in range(3)]
|
||||
assert all(msg is not None for msg in foo_received), (
|
||||
"Subscriber with matching topic should receive messages"
|
||||
)
|
||||
|
||||
bar_received = [sub_bar.receive_one(timeout=200) for _ in range(3)]
|
||||
assert all(msg is None for msg in bar_received), (
|
||||
"Subscriber with non-matching topic should receive no messages"
|
||||
)
|
||||
finally:
|
||||
pub.shutdown()
|
||||
sub_foo.close()
|
||||
sub_bar.close()
|
||||
|
||||
|
||||
def test_high_volume(publisher, subscriber):
|
||||
"""Test publishing and receiving a high volume of events"""
|
||||
num_batches = 10_000
|
||||
events_per_batch = 100
|
||||
|
||||
# Publish events in a separate thread to not block
|
||||
def publish_events():
|
||||
for i in range(num_batches):
|
||||
batch = create_test_events(events_per_batch)
|
||||
publisher.publish(batch)
|
||||
# Small delay to avoid overwhelming
|
||||
if i % 100 == 0:
|
||||
time.sleep(0.01)
|
||||
|
||||
received: list[tuple[int, SampleBatch]] = []
|
||||
|
||||
publisher_thread = threading.Thread(target=publish_events)
|
||||
publisher_thread.start()
|
||||
|
||||
start_time = time.time()
|
||||
while len(received) < num_batches:
|
||||
if time.time() - start_time > 10: # Timeout after 10 seconds
|
||||
break
|
||||
|
||||
result = subscriber.receive_one(timeout=100)
|
||||
if result:
|
||||
received.append(result)
|
||||
|
||||
publisher_thread.join()
|
||||
|
||||
assert len(received) >= num_batches * 0.9, "We should have received most messages"
|
||||
|
||||
seqs = [seq for seq, _ in received]
|
||||
assert sorted(seqs) == seqs, "Sequence numbers should be in order"
|
||||
|
||||
|
||||
def test_null_publisher():
|
||||
"""Test that NullEventPublisher can be used without errors"""
|
||||
publisher = NullEventPublisher(DP_RANK)
|
||||
|
||||
# This should not raise any errors
|
||||
batch = create_test_events(5)
|
||||
publisher.publish(batch)
|
||||
publisher.shutdown()
|
||||
|
||||
|
||||
def test_data_parallel_rank_tagging(publisher_config):
|
||||
"""Test that events are properly tagged with their data parallel rank"""
|
||||
|
||||
publisher_config.topic = "foo"
|
||||
pub_0 = EventPublisherFactory.create(publisher_config, DP_RANK)
|
||||
pub_1 = EventPublisherFactory.create(publisher_config, DP_RANK + 1)
|
||||
|
||||
# Hardcode the expected endpoints based on port offsetting behavior
|
||||
# Both ranks get offsets according to _offset_endpoint_port function
|
||||
base_endpoint = publisher_config.endpoint
|
||||
if "tcp://" in base_endpoint:
|
||||
# For TCP endpoints: tcp://localhost:5557 -> tcp://localhost:5557, tcp://localhost:5558
|
||||
expected_endpoint_0 = base_endpoint # rank 0 gets port + 0 = same port
|
||||
expected_endpoint_1 = base_endpoint.replace(
|
||||
":5557", ":5558"
|
||||
) # rank 1 gets port + 1
|
||||
else:
|
||||
# For inproc endpoints: inproc://test -> inproc://test_dp0, inproc://test_dp1
|
||||
expected_endpoint_0 = base_endpoint # rank 0 gets base
|
||||
expected_endpoint_1 = base_endpoint + "_dp1" # rank 1 gets _dp1
|
||||
|
||||
from .conftest import MockSubscriber
|
||||
|
||||
sub_0 = MockSubscriber(expected_endpoint_0, None, publisher_config.topic)
|
||||
sub_1 = MockSubscriber(expected_endpoint_1, None, publisher_config.topic)
|
||||
|
||||
try:
|
||||
time.sleep(0.1) # Let publishers start up
|
||||
|
||||
# Publish events from different ranks
|
||||
batch_0 = create_test_events(2)
|
||||
batch_1 = create_test_events(3)
|
||||
|
||||
pub_0.publish(batch_0)
|
||||
pub_1.publish(batch_1)
|
||||
|
||||
# Receive events from rank 0
|
||||
result_0 = sub_0.receive_one(timeout=200)
|
||||
assert result_0 is not None, "No message received from rank 0"
|
||||
seq_0, received_0 = result_0
|
||||
|
||||
# Receive events from rank 1
|
||||
result_1 = sub_1.receive_one(timeout=200)
|
||||
assert result_1 is not None, "No message received from rank 1"
|
||||
seq_1, received_1 = result_1
|
||||
|
||||
# Verify DP rank tagging
|
||||
assert received_0.data_parallel_rank == 0, (
|
||||
f"Expected DP rank 0, got {received_0.data_parallel_rank}"
|
||||
)
|
||||
assert received_1.data_parallel_rank == 1, (
|
||||
f"Expected DP rank 1, got {received_1.data_parallel_rank}"
|
||||
)
|
||||
|
||||
# Verify event content is correct
|
||||
assert len(received_0.events) == 2, "Wrong number of events from rank 0"
|
||||
assert len(received_1.events) == 3, "Wrong number of events from rank 1"
|
||||
|
||||
finally:
|
||||
pub_0.shutdown()
|
||||
pub_1.shutdown()
|
||||
sub_0.close()
|
||||
sub_1.close()
|
||||
|
||||
|
||||
def test_event_publisher_factory():
|
||||
"""Test event publisher factory creation behavior under different configurations"""
|
||||
from vllm.config.kv_events import KVEventsConfig
|
||||
from vllm.distributed.kv_events import ZmqEventPublisher
|
||||
|
||||
# test config is None
|
||||
publisher = EventPublisherFactory.create(None, DP_RANK)
|
||||
assert isinstance(publisher, NullEventPublisher)
|
||||
publisher.shutdown()
|
||||
|
||||
# test disable kv cache events
|
||||
config = KVEventsConfig(
|
||||
enable_kv_cache_events=False,
|
||||
publisher="zmq", # Even if zmq is specified, should return NullEventPublisher
|
||||
endpoint="tcp://localhost:5557",
|
||||
)
|
||||
publisher = EventPublisherFactory.create(config, DP_RANK)
|
||||
assert isinstance(publisher, NullEventPublisher)
|
||||
publisher.shutdown()
|
||||
|
||||
# test zmq publisher
|
||||
config = KVEventsConfig(
|
||||
enable_kv_cache_events=True,
|
||||
publisher="zmq",
|
||||
endpoint="inproc://test-factory-true",
|
||||
)
|
||||
publisher = EventPublisherFactory.create(config, DP_RANK)
|
||||
assert isinstance(publisher, ZmqEventPublisher)
|
||||
publisher.shutdown()
|
||||
|
||||
# test unknown publisher
|
||||
with pytest.raises(ValueError, match="Input should be"):
|
||||
KVEventsConfig(
|
||||
enable_kv_cache_events=True,
|
||||
publisher="unknown_publisher",
|
||||
endpoint="tcp://localhost:5557",
|
||||
)
|
||||
|
||||
# test publisher not specified
|
||||
config = KVEventsConfig(
|
||||
enable_kv_cache_events=True,
|
||||
# publisher not specified, should default to "zmq"
|
||||
endpoint="tcp://localhost:5557",
|
||||
)
|
||||
publisher = EventPublisherFactory.create(config, DP_RANK)
|
||||
assert isinstance(publisher, ZmqEventPublisher)
|
||||
publisher.shutdown()
|
||||
@@ -0,0 +1,231 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, NamedTuple
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config.model import RunnerOption
|
||||
from vllm.logger import init_logger
|
||||
|
||||
from ..utils import compare_two_settings, create_new_process_for_each_test
|
||||
|
||||
logger = init_logger("test_expert_parallel")
|
||||
|
||||
|
||||
class ParallelSetup(NamedTuple):
|
||||
tp_size: int
|
||||
eager_mode: bool
|
||||
chunked_prefill: bool
|
||||
|
||||
|
||||
class EPTestOptions(NamedTuple):
|
||||
trust_remote_code: bool
|
||||
tokenizer_mode: str | None
|
||||
load_format: str | None = None
|
||||
hf_overrides: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class EPTestSettings:
|
||||
parallel_setups: list[ParallelSetup]
|
||||
distributed_backends: list[str]
|
||||
runner: RunnerOption
|
||||
test_options: EPTestOptions
|
||||
|
||||
@staticmethod
|
||||
def detailed(
|
||||
*,
|
||||
tp_base: int = 2,
|
||||
runner: RunnerOption = "auto",
|
||||
trust_remote_code: bool = False,
|
||||
tokenizer_mode: str | None = None,
|
||||
load_format: str | None = None,
|
||||
hf_overrides: str | None = None,
|
||||
):
|
||||
return EPTestSettings(
|
||||
parallel_setups=[
|
||||
ParallelSetup(tp_size=tp_base, eager_mode=False, chunked_prefill=False),
|
||||
ParallelSetup(tp_size=tp_base, eager_mode=False, chunked_prefill=True),
|
||||
ParallelSetup(tp_size=tp_base, eager_mode=True, chunked_prefill=False),
|
||||
ParallelSetup(
|
||||
tp_size=2 * tp_base, eager_mode=False, chunked_prefill=True
|
||||
),
|
||||
ParallelSetup(
|
||||
tp_size=2 * tp_base, eager_mode=True, chunked_prefill=False
|
||||
),
|
||||
],
|
||||
distributed_backends=["mp", "ray"],
|
||||
runner=runner,
|
||||
test_options=EPTestOptions(
|
||||
trust_remote_code=trust_remote_code,
|
||||
tokenizer_mode=tokenizer_mode,
|
||||
load_format=load_format,
|
||||
hf_overrides=hf_overrides,
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def fast(
|
||||
*,
|
||||
tp_base: int = 2,
|
||||
runner: RunnerOption = "auto",
|
||||
trust_remote_code: bool = False,
|
||||
tokenizer_mode: str | None = None,
|
||||
load_format: str | None = None,
|
||||
hf_overrides: str | None = None,
|
||||
):
|
||||
return EPTestSettings(
|
||||
parallel_setups=[
|
||||
ParallelSetup(tp_size=tp_base, eager_mode=True, chunked_prefill=False),
|
||||
],
|
||||
distributed_backends=["mp"],
|
||||
runner=runner,
|
||||
test_options=EPTestOptions(
|
||||
trust_remote_code=trust_remote_code,
|
||||
tokenizer_mode=tokenizer_mode,
|
||||
load_format=load_format,
|
||||
hf_overrides=hf_overrides,
|
||||
),
|
||||
)
|
||||
|
||||
def iter_params(self, model_name: str):
|
||||
opts = self.test_options
|
||||
|
||||
for parallel_setup in self.parallel_setups:
|
||||
for distributed_backend in self.distributed_backends:
|
||||
yield (
|
||||
model_name,
|
||||
parallel_setup,
|
||||
distributed_backend,
|
||||
self.runner,
|
||||
opts,
|
||||
)
|
||||
|
||||
|
||||
# NOTE: You can adjust tp_base locally to fit the model in GPU
|
||||
# The values displayed here are only a rough indicator of the size of the model
|
||||
|
||||
TEST_MODELS = {
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": EPTestSettings.fast(trust_remote_code=True),
|
||||
"mistralai/Mixtral-8x7B-Instruct-v0.1": EPTestSettings.fast(tp_base=4),
|
||||
}
|
||||
|
||||
|
||||
def _compare_tp(
|
||||
model_name: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: EPTestOptions,
|
||||
num_gpus_available: int,
|
||||
*,
|
||||
method: Literal["generate"],
|
||||
):
|
||||
(
|
||||
tp_size,
|
||||
eager_mode,
|
||||
chunked_prefill,
|
||||
) = parallel_setup
|
||||
(
|
||||
trust_remote_code,
|
||||
tokenizer_mode,
|
||||
load_format,
|
||||
hf_overrides,
|
||||
) = test_options
|
||||
|
||||
if num_gpus_available < tp_size:
|
||||
pytest.skip(f"Need at least {tp_size} GPUs")
|
||||
|
||||
common_args = [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"float16",
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"8",
|
||||
"--load-format",
|
||||
"auto",
|
||||
]
|
||||
if chunked_prefill:
|
||||
common_args.append("--enable-chunked-prefill")
|
||||
if eager_mode:
|
||||
common_args.append("--enforce-eager")
|
||||
if runner != "auto":
|
||||
common_args.extend(["--runner", runner])
|
||||
if trust_remote_code:
|
||||
common_args.append("--trust-remote-code")
|
||||
if tokenizer_mode:
|
||||
common_args.extend(["--tokenizer-mode", tokenizer_mode])
|
||||
if load_format:
|
||||
common_args.extend(["--load-format", load_format])
|
||||
if hf_overrides:
|
||||
common_args.extend(["--hf-overrides", hf_overrides])
|
||||
|
||||
ep_env = {
|
||||
"VLLM_TEST_ENABLE_EP": "1",
|
||||
}
|
||||
|
||||
ep_args = [
|
||||
*common_args,
|
||||
"--tensor-parallel-size",
|
||||
str(tp_size),
|
||||
"--distributed-executor-backend",
|
||||
distributed_backend,
|
||||
]
|
||||
|
||||
# compare without expert parallelism
|
||||
tp_env = {
|
||||
"VLLM_TEST_ENABLE_EP": "0",
|
||||
}
|
||||
|
||||
tp_args = [
|
||||
*common_args,
|
||||
"--tensor-parallel-size",
|
||||
str(tp_size),
|
||||
"--distributed-executor-backend",
|
||||
"mp",
|
||||
]
|
||||
|
||||
try:
|
||||
compare_two_settings(
|
||||
model_name,
|
||||
ep_args,
|
||||
tp_args,
|
||||
ep_env,
|
||||
tp_env,
|
||||
method=method,
|
||||
max_wait_seconds=360,
|
||||
)
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_name", "parallel_setup", "distributed_backend", "runner", "test_options"),
|
||||
[
|
||||
params
|
||||
for model_name, settings in TEST_MODELS.items()
|
||||
for params in settings.iter_params(model_name)
|
||||
],
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_ep(
|
||||
model_name: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: EPTestOptions,
|
||||
num_gpus_available,
|
||||
):
|
||||
_compare_tp(
|
||||
model_name,
|
||||
parallel_setup,
|
||||
distributed_backend,
|
||||
runner,
|
||||
test_options,
|
||||
num_gpus_available,
|
||||
method="generate",
|
||||
)
|
||||
@@ -0,0 +1,246 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.model_executor.layers.fused_moe.expert_map_manager import (
|
||||
determine_expert_map,
|
||||
)
|
||||
|
||||
|
||||
def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts):
|
||||
"""Verify that the expert map follows the round_robin pattern."""
|
||||
# Calculate expected local experts (supporting non-divisible cases)
|
||||
base_experts = global_num_experts // ep_size
|
||||
remainder = global_num_experts % ep_size
|
||||
|
||||
local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts
|
||||
|
||||
# Expected expert IDs for this rank in round_robin pattern
|
||||
# For non-divisible cases, ranks with extra experts start earlier
|
||||
expected_expert_ids = []
|
||||
for expert_idx in range(local_num_experts):
|
||||
global_expert_id = ep_rank + expert_idx * ep_size
|
||||
expected_expert_ids.append(global_expert_id)
|
||||
|
||||
# Check that only expected experts are mapped to this rank
|
||||
for global_expert_id in range(global_num_experts):
|
||||
if global_expert_id in expected_expert_ids:
|
||||
local_expert_id = expert_map[global_expert_id]
|
||||
expected_local_id = expected_expert_ids.index(global_expert_id)
|
||||
assert local_expert_id == expected_local_id, (
|
||||
f"Global expert {global_expert_id} should map to local expert "
|
||||
f"{expected_local_id}, got {local_expert_id}"
|
||||
)
|
||||
else:
|
||||
assert expert_map[global_expert_id] == -1, (
|
||||
f"Global expert {global_expert_id} should not be mapped to this rank"
|
||||
)
|
||||
|
||||
# Verify that all local expert IDs are consecutive starting from 0
|
||||
local_expert_ids = [expert_map[global_id] for global_id in expected_expert_ids]
|
||||
expected_local_ids = list(range(local_num_experts))
|
||||
assert local_expert_ids == expected_local_ids, (
|
||||
f"Expected local expert IDs {expected_local_ids}, got {local_expert_ids}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expert_placement_strategy", ["round_robin"])
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
def test_expert_placement_various_sizes(expert_placement_strategy, world_size):
|
||||
"""Test round_robin expert placement with various expert counts."""
|
||||
|
||||
# Test with different global_num_experts values
|
||||
# Include both divisible and non-divisible cases
|
||||
if world_size == 2:
|
||||
test_cases = [
|
||||
(4, 2), # 4 experts (divisible)
|
||||
(8, 2), # 8 experts (divisible)
|
||||
(9, 2), # 9 experts (non-divisible)
|
||||
(16, 2), # 16 experts (divisible)
|
||||
(17, 2), # 17 experts (non-divisible)
|
||||
]
|
||||
elif world_size == 4:
|
||||
test_cases = [
|
||||
(8, 4), # 8 experts (divisible)
|
||||
(16, 4), # 16 experts (divisible)
|
||||
(18, 4), # 18 experts (non-divisible)
|
||||
(32, 4), # 32 experts (divisible)
|
||||
(33, 4), # 33 experts (non-divisible)
|
||||
]
|
||||
else:
|
||||
test_cases = []
|
||||
|
||||
for test_global_experts, test_ep_size in test_cases:
|
||||
# Ensure ep_size matches world_size
|
||||
assert test_ep_size == world_size, (
|
||||
f"ep_size {test_ep_size} must equal world_size {world_size}"
|
||||
)
|
||||
|
||||
# Test each rank
|
||||
for ep_rank in range(world_size):
|
||||
# Calculate expected local experts
|
||||
base_experts = test_global_experts // test_ep_size
|
||||
remainder = test_global_experts % test_ep_size
|
||||
if ep_rank < remainder:
|
||||
expected_test_local = base_experts + 1
|
||||
else:
|
||||
expected_test_local = base_experts
|
||||
|
||||
test_local_experts, test_expert_map, _ = determine_expert_map(
|
||||
ep_size=test_ep_size,
|
||||
ep_rank=ep_rank,
|
||||
global_num_experts=test_global_experts,
|
||||
expert_placement_strategy=expert_placement_strategy,
|
||||
)
|
||||
|
||||
assert test_local_experts == expected_test_local, (
|
||||
f"For {test_global_experts} experts on {test_ep_size} ranks, "
|
||||
f"rank {ep_rank}: expected {expected_test_local} local"
|
||||
f"experts, got {test_local_experts}"
|
||||
)
|
||||
|
||||
if test_expert_map is not None:
|
||||
assert test_expert_map.shape == (test_global_experts,), (
|
||||
f"Expected expert map shape ({test_global_experts},), "
|
||||
f"got {test_expert_map.shape}"
|
||||
)
|
||||
|
||||
# Verify round_robin pattern for this test case
|
||||
verify_round_robin_pattern(
|
||||
test_expert_map, ep_rank, test_ep_size, test_global_experts
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expert_placement_strategy", ["round_robin"])
|
||||
@pytest.mark.parametrize("world_size", [2, 4])
|
||||
def test_expert_placement_edge_cases(expert_placement_strategy, world_size):
|
||||
"""Test edge cases for round_robin expert placement."""
|
||||
|
||||
# Test case 1: ep_size = 1 (should return None for expert_map)
|
||||
local_num_experts, expert_map, _ = determine_expert_map(
|
||||
ep_size=1,
|
||||
ep_rank=0,
|
||||
global_num_experts=8,
|
||||
expert_placement_strategy=expert_placement_strategy,
|
||||
)
|
||||
assert local_num_experts == 8, "For ep_size=1, should get all experts"
|
||||
assert expert_map is None, "For ep_size=1, expert_map should be None"
|
||||
|
||||
# Test case 2: ep_size = 0 (should raise assertion)
|
||||
with pytest.raises(AssertionError):
|
||||
determine_expert_map(
|
||||
ep_size=0,
|
||||
ep_rank=0,
|
||||
global_num_experts=8,
|
||||
expert_placement_strategy=expert_placement_strategy,
|
||||
)
|
||||
|
||||
|
||||
def test_determine_expert_map_comprehensive():
|
||||
"""Test of determine_expert_map function with various configurations."""
|
||||
|
||||
# Test cases: (ep_size, ep_rank, global_num_experts,
|
||||
# expert_placement_strategy, expected_local, expected_map_pattern)
|
||||
test_cases = [
|
||||
# Round robin placement tests
|
||||
(
|
||||
2,
|
||||
0,
|
||||
8,
|
||||
"round_robin",
|
||||
4,
|
||||
[0, -1, 1, -1, 2, -1, 3, -1],
|
||||
), # rank 0 gets even experts
|
||||
(
|
||||
2,
|
||||
1,
|
||||
8,
|
||||
"round_robin",
|
||||
4,
|
||||
[-1, 0, -1, 1, -1, 2, -1, 3],
|
||||
), # rank 1 gets odd experts
|
||||
(
|
||||
2,
|
||||
0,
|
||||
9,
|
||||
"round_robin",
|
||||
5,
|
||||
[0, -1, 1, -1, 2, -1, 3, -1, 4],
|
||||
), # rank 0 gets 5 experts (even + last)
|
||||
(
|
||||
2,
|
||||
1,
|
||||
9,
|
||||
"round_robin",
|
||||
4,
|
||||
[-1, 0, -1, 1, -1, 2, -1, 3, -1],
|
||||
), # rank 1 gets 4 experts (odd)
|
||||
# 4-rank tests
|
||||
(
|
||||
4,
|
||||
0,
|
||||
8,
|
||||
"round_robin",
|
||||
2,
|
||||
[0, -1, -1, -1, 1, -1, -1, -1],
|
||||
), # rank 0 gets experts 0, 4
|
||||
(
|
||||
4,
|
||||
1,
|
||||
8,
|
||||
"round_robin",
|
||||
2,
|
||||
[-1, 0, -1, -1, -1, 1, -1, -1],
|
||||
), # rank 1 gets experts 1, 5
|
||||
(
|
||||
4,
|
||||
2,
|
||||
8,
|
||||
"round_robin",
|
||||
2,
|
||||
[-1, -1, 0, -1, -1, -1, 1, -1],
|
||||
), # rank 2 gets experts 2, 6
|
||||
(
|
||||
4,
|
||||
3,
|
||||
8,
|
||||
"round_robin",
|
||||
2,
|
||||
[-1, -1, -1, 0, -1, -1, -1, 1],
|
||||
), # rank 3 gets experts 3, 7
|
||||
]
|
||||
|
||||
for (
|
||||
ep_size,
|
||||
ep_rank,
|
||||
global_num_experts,
|
||||
expert_placement_strategy,
|
||||
expected_local,
|
||||
expected_map_pattern,
|
||||
) in test_cases:
|
||||
local_num_experts, expert_map, _ = determine_expert_map(
|
||||
ep_size=ep_size,
|
||||
ep_rank=ep_rank,
|
||||
global_num_experts=global_num_experts,
|
||||
expert_placement_strategy=expert_placement_strategy,
|
||||
)
|
||||
|
||||
assert local_num_experts == expected_local, (
|
||||
f"ep_size={ep_size}, ep_rank={ep_rank}, "
|
||||
f"global_num_experts={global_num_experts}, "
|
||||
f"expert_placement_strategy={expert_placement_strategy}: "
|
||||
f"expected {expected_local} local experts, got {local_num_experts}"
|
||||
)
|
||||
|
||||
if expected_map_pattern is None:
|
||||
assert expert_map is None, "Expected expert_map to be None"
|
||||
else:
|
||||
assert expert_map is not None, "Expected expert_map to not be None"
|
||||
actual_map = expert_map.tolist()
|
||||
assert actual_map == expected_map_pattern, (
|
||||
f"ep_size={ep_size}, ep_rank={ep_rank}, "
|
||||
f"global_num_experts={global_num_experts}, "
|
||||
f"expert_placement_strategy={expert_placement_strategy}: "
|
||||
f"expected map {expected_map_pattern}, got {actual_map}"
|
||||
)
|
||||
@@ -0,0 +1,86 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.distributed.kv_events import BlockRemoved, BlockStored
|
||||
|
||||
# Minimal ExternalBlockHash for testing (bytes are a valid ExternalBlockHash).
|
||||
_FAKE_HASH: bytes = b"\xab" * 32
|
||||
|
||||
|
||||
def _make_block_stored(
|
||||
group_idx: int | None = None,
|
||||
kv_cache_spec_sliding_window: int | None = None,
|
||||
) -> BlockStored:
|
||||
return BlockStored(
|
||||
block_hashes=[_FAKE_HASH],
|
||||
parent_block_hash=None,
|
||||
token_ids=[1, 2, 3, 4],
|
||||
block_size=4,
|
||||
lora_id=None,
|
||||
medium="GPU",
|
||||
lora_name=None,
|
||||
group_idx=group_idx,
|
||||
kv_cache_spec_sliding_window=kv_cache_spec_sliding_window,
|
||||
)
|
||||
|
||||
|
||||
def _make_block_removed(
|
||||
group_idx: int | None = None,
|
||||
) -> BlockRemoved:
|
||||
return BlockRemoved(
|
||||
block_hashes=[_FAKE_HASH],
|
||||
medium="GPU",
|
||||
group_idx=group_idx,
|
||||
)
|
||||
|
||||
|
||||
def test_block_stored_default_group_idx_is_none():
|
||||
"""group_idx defaults to None when not provided."""
|
||||
event = _make_block_stored()
|
||||
assert event.group_idx is None
|
||||
|
||||
|
||||
def test_block_removed_default_group_idx_is_none():
|
||||
"""group_idx defaults to None when not provided."""
|
||||
event = _make_block_removed()
|
||||
assert event.group_idx is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_idx", [1, 2, 3])
|
||||
def test_block_stored_hash_differs_by_group_idx(group_idx: int):
|
||||
"""BlockStored events that differ only in group_idx must hash differently."""
|
||||
other_group_idx = group_idx + 1
|
||||
event_a = _make_block_stored(group_idx=group_idx)
|
||||
event_b = _make_block_stored(group_idx=other_group_idx)
|
||||
assert hash(event_a) != hash(event_b)
|
||||
|
||||
|
||||
def test_block_stored_hash_same_for_equal_group_idx():
|
||||
"""Two BlockStored events with identical fields produce the same hash."""
|
||||
event_a = _make_block_stored(group_idx=1)
|
||||
event_b = _make_block_stored(group_idx=1)
|
||||
assert hash(event_a) == hash(event_b)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_idx", [1, 2, 3])
|
||||
def test_block_removed_hash_differs_by_group_idx(group_idx: int):
|
||||
"""BlockRemoved events that differ only in group_idx must hash differently."""
|
||||
other_group_idx = group_idx + 1
|
||||
event_a = _make_block_removed(group_idx=group_idx)
|
||||
event_b = _make_block_removed(group_idx=other_group_idx)
|
||||
assert hash(event_a) != hash(event_b)
|
||||
|
||||
|
||||
def test_block_removed_hash_same_for_equal_group_idx():
|
||||
"""Two BlockRemoved events with identical fields produce the same hash."""
|
||||
event_a = _make_block_removed(group_idx=1)
|
||||
event_b = _make_block_removed(group_idx=1)
|
||||
assert hash(event_a) == hash(event_b)
|
||||
|
||||
|
||||
def test_block_stored_hash_differs_by_sliding_window():
|
||||
event_a = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=128)
|
||||
event_b = _make_block_stored(group_idx=1, kv_cache_spec_sliding_window=256)
|
||||
assert hash(event_a) != hash(event_b)
|
||||
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm.config import (
|
||||
DeviceConfig,
|
||||
KVTransferConfig,
|
||||
ModelConfig,
|
||||
VllmConfig,
|
||||
set_current_vllm_config,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.utils import (
|
||||
get_kv_connector_cache_layout,
|
||||
)
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger("test_expert_parallel")
|
||||
|
||||
|
||||
def test_get_kv_connector_cache_layout_without_kv_connector():
|
||||
vllm_config = VllmConfig(device_config=DeviceConfig("cpu"))
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Test with default settings
|
||||
layout = get_kv_connector_cache_layout()
|
||||
assert layout == "NHD"
|
||||
|
||||
|
||||
def test_get_kv_connector_cache_layout_with_lmcache_connector():
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="LMCacheConnectorV1",
|
||||
kv_role="kv_both",
|
||||
)
|
||||
vllm_config = VllmConfig(
|
||||
device_config=DeviceConfig("cpu"), kv_transfer_config=kv_transfer_config
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Test with default settings
|
||||
layout = get_kv_connector_cache_layout()
|
||||
assert layout == "NHD"
|
||||
|
||||
|
||||
def test_get_kv_connector_cache_layout_with_nixl_connector():
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="NixlConnector",
|
||||
kv_role="kv_both",
|
||||
)
|
||||
model_config = ModelConfig()
|
||||
vllm_config = VllmConfig(
|
||||
device_config=DeviceConfig("cpu"),
|
||||
model_config=model_config,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Test with default settings
|
||||
layout = get_kv_connector_cache_layout()
|
||||
assert layout == "HND"
|
||||
|
||||
|
||||
def test_get_kv_connector_cache_layout_with_multi_connector():
|
||||
kv_transfer_config = KVTransferConfig(
|
||||
kv_connector="MultiConnector",
|
||||
kv_role="kv_both",
|
||||
kv_connector_extra_config={
|
||||
"connectors": [
|
||||
{"kv_connector": "ExampleConnector", "kv_role": "kv_both"},
|
||||
{"kv_connector": "NixlConnector", "kv_role": "kv_both"},
|
||||
]
|
||||
},
|
||||
)
|
||||
model_config = ModelConfig()
|
||||
vllm_config = VllmConfig(
|
||||
device_config=DeviceConfig("cpu"),
|
||||
model_config=model_config,
|
||||
kv_transfer_config=kv_transfer_config,
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# Test with default settings
|
||||
layout = get_kv_connector_cache_layout()
|
||||
assert layout == "HND"
|
||||
@@ -0,0 +1,941 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for MNNVL AllToAll operations.
|
||||
|
||||
Requires: docker run ... --cap-add=SYS_PTRACE ...
|
||||
Run: pytest tests/distributed/test_mnnvl_alltoall.py -v
|
||||
"""
|
||||
|
||||
import os
|
||||
import traceback
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from vllm.distributed import get_ep_group
|
||||
from vllm.utils.flashinfer import (
|
||||
has_flashinfer_nvlink_one_sided,
|
||||
has_flashinfer_nvlink_two_sided,
|
||||
)
|
||||
from vllm.utils.import_utils import has_deep_ep_v2
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
from ..utils import init_test_distributed_environment
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _has_sys_ptrace() -> bool:
|
||||
"""Check for SYS_PTRACE capability (bit 19 in CapEff)."""
|
||||
try:
|
||||
with open("/proc/self/status") as f:
|
||||
for line in f:
|
||||
if line.startswith("CapEff:"):
|
||||
return bool(int(line.split()[1], 16) & (1 << 19))
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _spawn_workers(worker_fn, world_size, *, dp_size=None):
|
||||
"""Spawn one process per GPU, run worker_fn, assert all succeed.
|
||||
|
||||
Uses an mp.Queue to propagate worker tracebacks back to the parent
|
||||
so pytest shows the actual failure, not just an exit code.
|
||||
"""
|
||||
if mp.get_start_method(allow_none=True) is None:
|
||||
mp.set_start_method("spawn")
|
||||
|
||||
port = str(get_open_port())
|
||||
# Allocate a second port for DP master when dp_size is set, so the
|
||||
# distributed init port and DP port can't collide even under xdist.
|
||||
dp_port = str(get_open_port()) if dp_size is not None else None
|
||||
err_queue: mp.Queue = mp.Queue()
|
||||
procs = []
|
||||
for rank in range(world_size):
|
||||
p = mp.Process(
|
||||
target=_run_worker,
|
||||
args=(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue),
|
||||
)
|
||||
p.start()
|
||||
procs.append(p)
|
||||
for p in procs:
|
||||
p.join()
|
||||
|
||||
# Collect any errors from workers before asserting.
|
||||
errors = []
|
||||
while not err_queue.empty():
|
||||
errors.append(err_queue.get_nowait())
|
||||
err_queue.close()
|
||||
err_queue.join_thread()
|
||||
if errors:
|
||||
pytest.fail("Worker(s) failed:\n" + "\n---\n".join(errors))
|
||||
|
||||
|
||||
def _run_worker(rank, world_size, port, worker_fn, dp_size, dp_port, err_queue):
|
||||
"""Per-process setup: device, distributed env, then call worker_fn.
|
||||
|
||||
Args:
|
||||
dp_size: If set, initialize with tp=1 and data_parallel_size=dp_size.
|
||||
Otherwise use tp=world_size (default for EP-based tests).
|
||||
dp_port: Separate port for the DP master (only used when dp_size is set).
|
||||
err_queue: Queue for propagating tracebacks to the parent process.
|
||||
"""
|
||||
try:
|
||||
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
torch.accelerator.set_device_index(rank)
|
||||
if dp_size is not None:
|
||||
_init_dp_environment(world_size, rank, port, dp_size, dp_port)
|
||||
else:
|
||||
init_test_distributed_environment(world_size, 1, rank, port)
|
||||
worker_fn(rank, world_size)
|
||||
torch.distributed.barrier()
|
||||
except Exception:
|
||||
err_queue.put(f"[Rank {rank}]\n{traceback.format_exc()}")
|
||||
# Don't re-raise: the parent reads errors from err_queue.
|
||||
# A non-zero exit from the re-raise would be redundant.
|
||||
import sys
|
||||
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def _init_dp_environment(world_size, rank, port, dp_size, dp_port):
|
||||
"""Initialize distributed env with data parallelism.
|
||||
|
||||
Sets up tp=1, pp=1, dp=dp_size. Each process is one DP rank
|
||||
with local rank 0 within its (trivial) tp*pp group.
|
||||
|
||||
Args:
|
||||
port: Port for torch.distributed init.
|
||||
dp_port: Separate port for the DP master group init.
|
||||
"""
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.config.parallel import ParallelConfig
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
init_distributed_environment,
|
||||
)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config = ParallelConfig(
|
||||
data_parallel_size=dp_size,
|
||||
data_parallel_rank=rank,
|
||||
# Pre-populate port list so __post_init__ doesn't auto-generate
|
||||
# random ports. All DP ranks must agree on the same port.
|
||||
_data_parallel_master_port_list=[int(dp_port)],
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
# rank=0 here because each DP rank has a single (tp=1,pp=1) process,
|
||||
# so the local rank within the tp*pp group is always 0.
|
||||
# init_distributed_environment will offset by data_parallel_rank.
|
||||
init_distributed_environment(
|
||||
world_size=1, # tp * pp = 1
|
||||
rank=0,
|
||||
distributed_init_method=f"tcp://localhost:{port}",
|
||||
local_rank=rank,
|
||||
)
|
||||
ensure_model_parallel_initialized(1, 1)
|
||||
|
||||
|
||||
def _make_forward_context(rank, world_size, num_tokens_per_rank):
|
||||
"""Create a forward context with mock DP metadata for AgRs tests.
|
||||
|
||||
Returns a context manager suitable for ``with`` statements.
|
||||
The real DPMetadata (with sp_local_sizes etc.) is created internally
|
||||
by set_forward_context from num_tokens_across_dp; the attn_metadata
|
||||
placeholder just satisfies the "attn_metadata is not None" guard.
|
||||
"""
|
||||
from vllm.config.parallel import ParallelConfig
|
||||
from vllm.config.vllm import VllmConfig
|
||||
from vllm.forward_context import set_forward_context
|
||||
|
||||
class _AttnMeta:
|
||||
"""Minimal placeholder so set_forward_context's
|
||||
``attn_metadata is not None`` guard (forward_context.py:334)
|
||||
is satisfied. The real DPMetadata is built from num_tokens_across_dp."""
|
||||
|
||||
dp_metadata = None
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
vllm_config.parallel_config = ParallelConfig(
|
||||
data_parallel_size=world_size,
|
||||
is_moe_model=True,
|
||||
data_parallel_rank=rank,
|
||||
)
|
||||
return set_forward_context(
|
||||
_AttnMeta(),
|
||||
vllm_config,
|
||||
num_tokens=num_tokens_per_rank,
|
||||
num_tokens_across_dp=torch.tensor(
|
||||
[num_tokens_per_rank] * world_size, dtype=torch.int
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Skip conditions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
requires_multi_gpu = pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need >= 2 GPUs"
|
||||
)
|
||||
requires_two_sided = pytest.mark.skipif(
|
||||
not has_flashinfer_nvlink_two_sided(),
|
||||
reason="FlashInfer NVLink two-sided not available",
|
||||
)
|
||||
requires_one_sided = pytest.mark.skipif(
|
||||
not has_flashinfer_nvlink_one_sided(),
|
||||
reason="FlashInfer NVLink one-sided not available",
|
||||
)
|
||||
requires_ptrace = pytest.mark.skipif(
|
||||
not _has_sys_ptrace(),
|
||||
reason="SYS_PTRACE required (docker run --cap-add=SYS_PTRACE)",
|
||||
)
|
||||
requires_deep_ep_v2 = pytest.mark.skipif(
|
||||
not has_deep_ep_v2(),
|
||||
reason="DeepEP v2 (ElasticBuffer) not available or NCCL < 2.30.4",
|
||||
)
|
||||
|
||||
# NOTE: No module-level pytestmark here. The FlashInfer lifecycle tests have
|
||||
# their own @requires_two_sided / @requires_one_sided decorators, and
|
||||
# test_args_dispatch_combine uses only standard torch.distributed ops and
|
||||
# should run even when FlashInfer NVLink backends are not installed.
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1: Two-sided manager lifecycle (init, cleanup, reinit, ensure_init)
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests FlashInferNVLinkTwoSidedManager which wraps FlashInfer's MnnvlMoe.
|
||||
# initialize() allocates MNNVL shared workspaces via MnnvlMoe.get_moe_workspaces,
|
||||
# which uses pidfd_getfd() to share memory file descriptors across processes —
|
||||
# hence the SYS_PTRACE requirement.
|
||||
#
|
||||
# Uses EP group (get_ep_group) because the two-sided manager is constructed
|
||||
# with an EP-scoped communicator in production. With tp=world_size the EP
|
||||
# group spans all ranks, giving us a multi-rank group for testing.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _two_sided_lifecycle_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
FlashInferNVLinkTwoSidedManager,
|
||||
)
|
||||
|
||||
cpu_group = get_ep_group().cpu_group
|
||||
num_gpus = torch.accelerator.device_count()
|
||||
manager = FlashInferNVLinkTwoSidedManager(cpu_group)
|
||||
|
||||
# Not initialized yet
|
||||
assert not manager.initialized
|
||||
assert manager.rank == rank
|
||||
assert manager.world_size == world_size
|
||||
|
||||
# Initialize
|
||||
manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)
|
||||
assert manager.initialized
|
||||
assert manager.workspace_tensor is not None
|
||||
assert manager.prepare_workspace_tensor is not None
|
||||
assert manager.mapping is not None
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Cleanup
|
||||
manager.cleanup()
|
||||
assert not manager.initialized
|
||||
assert manager.workspace_tensor is None
|
||||
assert manager.prepare_workspace_tensor is None
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Reinitialize
|
||||
manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)
|
||||
assert manager.initialized
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# ensure_alltoall_workspace_initialized is idempotent when already init'd
|
||||
assert manager.ensure_alltoall_workspace_initialized()
|
||||
assert manager.initialized
|
||||
|
||||
manager.cleanup()
|
||||
assert not manager.initialized
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_two_sided
|
||||
@requires_ptrace
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_two_sided_manager_lifecycle(world_size):
|
||||
"""Test init, cleanup, reinit, and ensure_initialized idempotency."""
|
||||
_spawn_workers(_two_sided_lifecycle_worker, world_size)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2: One-sided manager lifecycle (init, cleanup, reinit)
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests FlashInferNVLinkOneSidedManager which wraps FlashInfer's MoeAlltoAll.
|
||||
# initialize() creates MoeAlltoAll with an MnnvlConfig, which allocates MNNVL
|
||||
# shared workspaces — same cross-process memory sharing as two-sided, hence
|
||||
# the SYS_PTRACE requirement.
|
||||
#
|
||||
# Uses DP group (get_dp_group) because the one-sided manager's initialize()
|
||||
# internally calls get_dp_group() to set up the MnnvlConfig communicator.
|
||||
# We therefore need a real DP group with world_size > 1, which requires
|
||||
# dp_size=world_size via _init_dp_environment.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _one_sided_lifecycle_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
FlashInferNVLinkOneSidedManager,
|
||||
)
|
||||
from vllm.distributed.parallel_state import get_dp_group
|
||||
|
||||
cpu_group = get_dp_group().cpu_group
|
||||
manager = FlashInferNVLinkOneSidedManager(cpu_group)
|
||||
|
||||
assert not manager.initialized
|
||||
assert manager.rank == rank
|
||||
assert manager.world_size == world_size
|
||||
|
||||
init_kwargs = dict(
|
||||
max_num_tokens=1024,
|
||||
top_k=2,
|
||||
num_experts=world_size * 8,
|
||||
hidden_size=4096,
|
||||
)
|
||||
|
||||
# Initialize
|
||||
manager.initialize(**init_kwargs)
|
||||
assert manager.initialized
|
||||
assert manager.moe_alltoall is not None
|
||||
assert manager.mapping is not None
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Cleanup
|
||||
manager.cleanup()
|
||||
assert not manager.initialized
|
||||
assert manager.moe_alltoall is None
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Reinitialize with different token count
|
||||
manager.initialize(**{**init_kwargs, "max_num_tokens": 2048})
|
||||
assert manager.initialized
|
||||
|
||||
torch.distributed.barrier()
|
||||
manager.cleanup()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_one_sided
|
||||
@requires_ptrace
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_one_sided_manager_lifecycle(world_size):
|
||||
"""Test init, cleanup, and reinit with different params."""
|
||||
_spawn_workers(
|
||||
_one_sided_lifecycle_worker,
|
||||
world_size,
|
||||
dp_size=world_size,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2b: One-sided manager grows workspace across heterogeneous MoE layers
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Models with heterogeneous MoE quantization — most notably a quantized base
|
||||
# MoE combined with an unquantized MTP head — can call initialize() multiple
|
||||
# times with different per-token dispatch payload sizes. The shared workspace
|
||||
# must grow to the union and the MoeAlltoAll must be rebuilt; otherwise a
|
||||
# later layer's combine call overruns the workspace sized for the first
|
||||
# layer's smaller payload and trips FlashInfer's combinePayloadOffset assert.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _one_sided_workspace_grow_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
FlashInferNVLinkOneSidedManager,
|
||||
)
|
||||
from vllm.distributed.parallel_state import get_dp_group
|
||||
|
||||
cpu_group = get_dp_group().cpu_group
|
||||
manager = FlashInferNVLinkOneSidedManager(cpu_group)
|
||||
|
||||
base_kwargs = dict(
|
||||
max_num_tokens=1024,
|
||||
top_k=2,
|
||||
num_experts=world_size * 8,
|
||||
hidden_size=4096,
|
||||
)
|
||||
nvfp4_kwargs = dict(
|
||||
dispatch_dtype_bytes_per_elem=0,
|
||||
dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16,
|
||||
)
|
||||
bf16_kwargs = dict(
|
||||
dispatch_dtype_bytes_per_elem=2,
|
||||
dispatch_scale_bytes_per_token=0,
|
||||
)
|
||||
|
||||
# First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16).
|
||||
manager.initialize(**base_kwargs, **nvfp4_kwargs)
|
||||
assert manager.initialized
|
||||
nvfp4_workspace_size = manager.workspace_size
|
||||
nvfp4_moe_alltoall = manager.moe_alltoall
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Second init: bf16-like (hidden_bytes = hidden * 2). Models the case of
|
||||
# a quantized base MoE followed by an unquantized MoE layer (e.g. an MTP
|
||||
# head). Per-token dispatch payload is ~4x larger, so the union workspace
|
||||
# must grow and MoeAlltoAll must be rebuilt.
|
||||
manager.initialize(**base_kwargs, **bf16_kwargs)
|
||||
assert manager.initialized
|
||||
assert manager.workspace_size > nvfp4_workspace_size
|
||||
assert manager.moe_alltoall is not nvfp4_moe_alltoall
|
||||
bf16_workspace_size = manager.workspace_size
|
||||
bf16_moe_alltoall = manager.moe_alltoall
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Third init: back to NVFP4-like shape. Existing workspace already covers
|
||||
# it, so initialize() must no-op — no shrink, no rebuild.
|
||||
manager.initialize(**base_kwargs, **nvfp4_kwargs)
|
||||
assert manager.initialized
|
||||
assert manager.workspace_size == bf16_workspace_size
|
||||
assert manager.moe_alltoall is bf16_moe_alltoall
|
||||
|
||||
torch.distributed.barrier()
|
||||
manager.cleanup()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_one_sided
|
||||
@requires_ptrace
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_one_sided_manager_workspace_grow(world_size):
|
||||
"""A later initialize() with a larger per-token payload must grow the
|
||||
workspace and rebuild MoeAlltoAll; a later initialize() with a smaller
|
||||
payload must no-op."""
|
||||
_spawn_workers(
|
||||
_one_sided_workspace_grow_worker,
|
||||
world_size,
|
||||
dp_size=world_size,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3: AgRs dispatch/combine with value validation
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests AgRsAll2AllManager which uses only standard torch.distributed
|
||||
# all_gatherv / reduce_scatterv — no FlashInfer or MNNVL dependency.
|
||||
# This test validates the reference all-to-all implementation that other
|
||||
# backends are compared against.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _args_dispatch_combine_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import AgRsAll2AllManager
|
||||
from vllm.forward_context import get_forward_context
|
||||
|
||||
cpu_group = get_ep_group().cpu_group
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
|
||||
hidden_size = 64
|
||||
tokens_per_rank = 16
|
||||
experts_per_token = 2
|
||||
num_experts = world_size * 4
|
||||
total_tokens = world_size * tokens_per_rank
|
||||
|
||||
# Deterministic per-rank data: rank r has value (r + 1)
|
||||
hidden = torch.full(
|
||||
(tokens_per_rank, hidden_size),
|
||||
float(rank + 1),
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
router = torch.full(
|
||||
(tokens_per_rank, num_experts),
|
||||
float(rank + 1) * 10,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
weights = torch.full(
|
||||
(tokens_per_rank, experts_per_token),
|
||||
float(rank + 1) * 100,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
ids = torch.full(
|
||||
(tokens_per_rank, experts_per_token),
|
||||
rank,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
|
||||
with _make_forward_context(rank, world_size, tokens_per_rank):
|
||||
manager = AgRsAll2AllManager(cpu_group)
|
||||
dp_metadata = get_forward_context().dp_metadata
|
||||
|
||||
with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
|
||||
# -- dispatch_router_logits --
|
||||
d_hidden, d_router = manager.dispatch_router_logits(
|
||||
hidden.clone(),
|
||||
router.clone(),
|
||||
is_sequence_parallel=True,
|
||||
)
|
||||
assert d_hidden.shape == (total_tokens, hidden_size)
|
||||
assert d_router.shape == (total_tokens, num_experts)
|
||||
|
||||
for r in range(world_size):
|
||||
s = r * tokens_per_rank
|
||||
e = (r + 1) * tokens_per_rank
|
||||
torch.testing.assert_close(
|
||||
d_hidden[s:e],
|
||||
torch.full_like(d_hidden[s:e], float(r + 1)),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
d_router[s:e],
|
||||
torch.full_like(d_router[s:e], float(r + 1) * 10),
|
||||
)
|
||||
|
||||
# -- dispatch --
|
||||
d_hidden2, d_weights, d_ids = manager.dispatch(
|
||||
hidden.clone(),
|
||||
weights.clone(),
|
||||
ids.clone(),
|
||||
is_sequence_parallel=True,
|
||||
)
|
||||
assert d_hidden2.shape == (total_tokens, hidden_size)
|
||||
assert d_weights.shape == (total_tokens, experts_per_token)
|
||||
assert d_ids.shape == (total_tokens, experts_per_token)
|
||||
|
||||
for r in range(world_size):
|
||||
s = r * tokens_per_rank
|
||||
e = (r + 1) * tokens_per_rank
|
||||
torch.testing.assert_close(
|
||||
d_weights[s:e],
|
||||
torch.full_like(d_weights[s:e], float(r + 1) * 100),
|
||||
)
|
||||
assert (d_ids[s:e] == r).all()
|
||||
|
||||
# -- combine (reduce-scatter) --
|
||||
# Each token i has value i in all columns; after reduce-scatter
|
||||
# each rank gets its slice, summed across ranks.
|
||||
expert_out = (
|
||||
torch.arange(total_tokens, device=device, dtype=torch.float32)
|
||||
.unsqueeze(1)
|
||||
.expand(total_tokens, hidden_size)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
combined = manager.combine(expert_out, is_sequence_parallel=True)
|
||||
assert combined.shape == (tokens_per_rank, hidden_size)
|
||||
|
||||
for i in range(tokens_per_rank):
|
||||
expected_val = float(rank * tokens_per_rank + i) * world_size
|
||||
torch.testing.assert_close(
|
||||
combined[i],
|
||||
torch.full_like(combined[i], expected_val),
|
||||
)
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_args_dispatch_combine(world_size):
|
||||
"""Validate dispatch gathers all-rank data and combine reduces correctly."""
|
||||
_spawn_workers(_args_dispatch_combine_worker, world_size)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4: FlashInfer two-sided dispatch/combine data communication
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests actual data flow through the FlashInfer NVLink two-sided backend
|
||||
# by calling flashinfer_alltoall_dispatch (with defer_input_quant=True to
|
||||
# skip quantization) and flashinfer_alltoall_combine, then verifying exact
|
||||
# round-trip values. Dispatch sends each token once per distinct expert
|
||||
# rank, and combine performs an unweighted sum, so:
|
||||
# dispatch(hidden) → identity → combine = hidden * num_distinct_ranks(i)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _two_sided_data_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
FlashInferNVLinkTwoSidedManager,
|
||||
)
|
||||
from vllm.distributed.parallel_state import get_dp_group
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEQuantConfig,
|
||||
FusedMoEQuantDesc,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501
|
||||
flashinfer_alltoall_combine,
|
||||
flashinfer_alltoall_dispatch,
|
||||
)
|
||||
|
||||
# Use DP group because MnnvlMoe workspace allocation calls get_dp_group()
|
||||
# internally and requires dp_size == ep_size.
|
||||
cpu_group = get_dp_group().cpu_group
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
num_gpus = torch.accelerator.device_count()
|
||||
|
||||
hidden_size = 128
|
||||
tokens_per_rank = 32
|
||||
experts_per_token = 2
|
||||
num_experts = world_size * 4
|
||||
|
||||
# Initialize the FlashInfer two-sided manager
|
||||
manager = FlashInferNVLinkTwoSidedManager(cpu_group)
|
||||
manager.initialize(world_size=world_size, rank=rank, gpus_per_node=num_gpus)
|
||||
assert manager.initialized
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Create deterministic per-rank test data
|
||||
torch.manual_seed(rank + 42)
|
||||
hidden = torch.randn(
|
||||
tokens_per_rank,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
# Assign each token to experts spread across ranks so tokens move between GPUs
|
||||
topk_ids = torch.randint(
|
||||
0,
|
||||
num_experts,
|
||||
(tokens_per_rank, experts_per_token),
|
||||
device=device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
topk_weights = torch.rand(
|
||||
tokens_per_rank,
|
||||
experts_per_token,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
# Unquantized config: quant_dtype=None means moe_kernel_quantize_input is a no-op
|
||||
no_quant = FusedMoEQuantDesc()
|
||||
quant_config = FusedMoEQuantConfig(
|
||||
_a1=no_quant,
|
||||
_a2=no_quant,
|
||||
_w1=no_quant,
|
||||
_w2=no_quant,
|
||||
)
|
||||
assert quant_config.quant_dtype is None # sanity: no quantization
|
||||
|
||||
with _make_forward_context(rank, world_size, tokens_per_rank):
|
||||
dp_metadata = get_forward_context().dp_metadata
|
||||
|
||||
with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
|
||||
local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
|
||||
|
||||
# --- FlashInfer two-sided dispatch ---
|
||||
alltoall_info, fi_topk_ids, fi_topk_weights, fi_hidden, fi_scale = (
|
||||
flashinfer_alltoall_dispatch(
|
||||
manager,
|
||||
local_sizes,
|
||||
hidden.clone(),
|
||||
None, # no global scale
|
||||
topk_ids.clone(),
|
||||
topk_weights.clone(),
|
||||
experts_per_token,
|
||||
num_experts,
|
||||
quant_config,
|
||||
defer_input_quant=True,
|
||||
)
|
||||
)
|
||||
assert fi_scale is None # deferred quant: no scale produced
|
||||
assert fi_hidden is not None
|
||||
assert fi_hidden.shape[1] == hidden_size
|
||||
assert fi_hidden.numel() > 0
|
||||
|
||||
# --- Round-trip exact verification ---
|
||||
# The all-to-all sends each token once per *distinct* expert
|
||||
# rank. Combine performs an unweighted sum of the per-rank
|
||||
# contributions. With identity expert (feeding dispatched
|
||||
# hidden straight back):
|
||||
# result[i] = hidden[i] * num_distinct_expert_ranks(i)
|
||||
combined = flashinfer_alltoall_combine(
|
||||
manager,
|
||||
fi_hidden,
|
||||
top_k=experts_per_token,
|
||||
token_count=tokens_per_rank,
|
||||
alltoall_info=alltoall_info,
|
||||
)
|
||||
assert combined.shape == (tokens_per_rank, hidden_size)
|
||||
|
||||
experts_per_rank = num_experts // world_size
|
||||
expert_ranks = topk_ids // experts_per_rank # (tokens, top_k)
|
||||
num_distinct = torch.tensor(
|
||||
[len(set(row.tolist())) for row in expert_ranks],
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
).unsqueeze(1) # (tokens, 1)
|
||||
expected = (hidden.float() * num_distinct).to(hidden.dtype)
|
||||
torch.testing.assert_close(combined, expected)
|
||||
|
||||
# --- Linearity check with scaled expert output ---
|
||||
# Scaling the expert output by a constant should scale the
|
||||
# combined result by the same constant.
|
||||
scale = 3.0
|
||||
combined_scaled = flashinfer_alltoall_combine(
|
||||
manager,
|
||||
fi_hidden * scale,
|
||||
top_k=experts_per_token,
|
||||
token_count=tokens_per_rank,
|
||||
alltoall_info=alltoall_info,
|
||||
)
|
||||
expected_scaled = (hidden.float() * num_distinct * scale).to(hidden.dtype)
|
||||
torch.testing.assert_close(combined_scaled, expected_scaled)
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
manager.cleanup()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_two_sided
|
||||
@requires_ptrace
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_two_sided_dispatch_combine(world_size):
|
||||
"""Test FlashInfer two-sided dispatch/combine with exact value verification."""
|
||||
_spawn_workers(_two_sided_data_worker, world_size, dp_size=world_size)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5: FlashInfer one-sided dispatch/combine data communication
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests actual data flow through the FlashInfer NVLink one-sided backend
|
||||
# by calling MoeAlltoAll.dispatch() and MoeAlltoAll.combine() directly
|
||||
# with synthetic payloads, then verifying shapes and round-trip consistency.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _one_sided_data_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
FlashInferNVLinkOneSidedManager,
|
||||
)
|
||||
from vllm.distributed.parallel_state import get_dp_group
|
||||
from vllm.forward_context import get_forward_context
|
||||
|
||||
cpu_group = get_dp_group().cpu_group
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
|
||||
hidden_size = 256
|
||||
tokens_per_rank = 32
|
||||
experts_per_token = 2
|
||||
num_experts = world_size * 8
|
||||
|
||||
# Initialize the one-sided manager
|
||||
manager = FlashInferNVLinkOneSidedManager(cpu_group)
|
||||
manager.initialize(
|
||||
max_num_tokens=tokens_per_rank,
|
||||
top_k=experts_per_token,
|
||||
num_experts=num_experts,
|
||||
hidden_size=hidden_size,
|
||||
# Account for the fp8 block-scale payload (a1q_scale: hidden//16 bytes
|
||||
# per token) that is dispatched alongside the nvfp4 hidden states.
|
||||
# Without this the dispatch region is under-reserved and the combine
|
||||
# payload overflows the per-rank workspace.
|
||||
dispatch_scale_bytes_per_token=hidden_size // 16,
|
||||
)
|
||||
assert manager.initialized
|
||||
assert manager.moe_alltoall is not None
|
||||
|
||||
with _make_forward_context(rank, world_size, tokens_per_rank):
|
||||
dp_metadata = get_forward_context().dp_metadata
|
||||
|
||||
with dp_metadata.sp_local_sizes(sequence_parallel_size=1):
|
||||
local_sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
|
||||
runtime_max_tokens = max(local_sizes)
|
||||
|
||||
# Create test data with raw tensors matching the nvfp4 payload
|
||||
# sizes the workspace was allocated for:
|
||||
# a1q: (tokens, hidden_size // 2) — nvfp4 hidden states
|
||||
# a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors
|
||||
torch.manual_seed(rank + 42)
|
||||
a1q = torch.randint(
|
||||
0,
|
||||
256,
|
||||
(tokens_per_rank, hidden_size // 2),
|
||||
device=device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
a1q_scale = torch.randint(
|
||||
0,
|
||||
256,
|
||||
(tokens_per_rank, hidden_size // 16),
|
||||
device=device,
|
||||
dtype=torch.uint8,
|
||||
)
|
||||
topk_ids = torch.randint(
|
||||
0,
|
||||
num_experts,
|
||||
(tokens_per_rank, experts_per_token),
|
||||
device=device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
topk_weights = torch.rand(
|
||||
tokens_per_rank,
|
||||
experts_per_token,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
# --- One-sided dispatch ---
|
||||
payloads = [a1q, a1q_scale, topk_ids, topk_weights]
|
||||
recv_payloads = manager.moe_alltoall.dispatch(
|
||||
token_selected_experts=topk_ids,
|
||||
input_payloads=payloads,
|
||||
runtime_max_tokens_per_rank=runtime_max_tokens,
|
||||
)
|
||||
assert len(recv_payloads) == 4
|
||||
recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads
|
||||
assert recv_a1q.numel() > 0
|
||||
assert recv_ids.numel() > 0
|
||||
|
||||
# --- Round-trip exact verification ---
|
||||
# The dispatch routes each token once per *distinct* expert
|
||||
# rank. Combine performs an unweighted sum of per-rank
|
||||
# contributions. With constant expert output (all 1s):
|
||||
# result[i] = 1.0 * num_distinct_expert_ranks(i)
|
||||
expert_output = torch.ones(
|
||||
world_size,
|
||||
runtime_max_tokens,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
combined = manager.moe_alltoall.combine(
|
||||
payload=expert_output,
|
||||
runtime_max_tokens_per_rank=runtime_max_tokens,
|
||||
)
|
||||
assert combined.shape == (tokens_per_rank, hidden_size)
|
||||
|
||||
experts_per_rank = num_experts // world_size
|
||||
expert_ranks = topk_ids // experts_per_rank # (tokens, top_k)
|
||||
num_distinct = torch.tensor(
|
||||
[len(set(row.tolist())) for row in expert_ranks],
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
).unsqueeze(1) # (tokens, 1)
|
||||
expected = num_distinct.expand_as(combined)
|
||||
torch.testing.assert_close(combined, expected)
|
||||
|
||||
# --- Linearity check with scaled expert output ---
|
||||
# Scaling the expert output by a constant should scale the
|
||||
# combined result by the same constant.
|
||||
# Re-dispatch to reset internal state (one-sided requires a
|
||||
# fresh dispatch before each combine).
|
||||
manager.moe_alltoall.dispatch(
|
||||
token_selected_experts=topk_ids,
|
||||
input_payloads=payloads,
|
||||
runtime_max_tokens_per_rank=runtime_max_tokens,
|
||||
)
|
||||
scale = 3.0
|
||||
combined_scaled = manager.moe_alltoall.combine(
|
||||
payload=expert_output * scale,
|
||||
runtime_max_tokens_per_rank=runtime_max_tokens,
|
||||
)
|
||||
expected_scaled = (expected * scale).to(torch.bfloat16)
|
||||
torch.testing.assert_close(combined_scaled, expected_scaled)
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
manager.cleanup()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_one_sided
|
||||
@requires_ptrace
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_one_sided_dispatch_combine(world_size):
|
||||
"""Test FlashInfer one-sided dispatch/combine with actual data flow."""
|
||||
_spawn_workers(_one_sided_data_worker, world_size, dp_size=world_size)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6: DeepEP v2 (ElasticBuffer) manager lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
#
|
||||
# Tests DeepEPV2All2AllManager which wraps DeepEP's ElasticBuffer API using
|
||||
# the NCCL GIN backend. Requires DeepEP >= 2.0 and NCCL >= 2.30.4.
|
||||
#
|
||||
# Uses EP group because the DeepEP v2 manager is constructed with an
|
||||
# EP-scoped communicator in production. With tp=world_size the EP group
|
||||
# spans all ranks.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _deepep_v2_lifecycle_worker(rank, world_size):
|
||||
from vllm.distributed.device_communicators.all2all import (
|
||||
DeepEPV2All2AllManager,
|
||||
)
|
||||
|
||||
cpu_group = get_ep_group().cpu_group
|
||||
manager = DeepEPV2All2AllManager(cpu_group)
|
||||
|
||||
assert manager.rank == rank
|
||||
assert manager.world_size == world_size
|
||||
assert manager._num_sms is None
|
||||
|
||||
hidden_size = 7168
|
||||
num_experts = world_size * 32
|
||||
num_topk = 8
|
||||
max_tokens = 256
|
||||
|
||||
handle_kwargs = dict(
|
||||
num_max_tokens_per_rank=max_tokens,
|
||||
hidden=hidden_size,
|
||||
num_topk=num_topk,
|
||||
num_experts=num_experts,
|
||||
use_fp8_dispatch=False,
|
||||
)
|
||||
|
||||
handle = manager.get_handle(handle_kwargs)
|
||||
assert handle is not None
|
||||
assert manager._num_sms is not None
|
||||
assert manager._num_sms > 0
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# get_handle again with same args should return cached handle
|
||||
handle2 = manager.get_handle(dict(handle_kwargs))
|
||||
assert handle2 is handle
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Destroy clears the cache
|
||||
manager.destroy()
|
||||
assert len(manager.handle_cache._cache) == 0
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
# Re-create after destroy
|
||||
handle3 = manager.get_handle(dict(handle_kwargs))
|
||||
assert handle3 is not None
|
||||
|
||||
torch.distributed.barrier()
|
||||
manager.destroy()
|
||||
|
||||
|
||||
@requires_multi_gpu
|
||||
@requires_deep_ep_v2
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
def test_deepep_v2_manager_lifecycle(world_size):
|
||||
"""Test DeepEP v2 ElasticBuffer manager init, caching, and destroy."""
|
||||
_spawn_workers(_deepep_v2_lifecycle_worker, world_size)
|
||||
@@ -0,0 +1,79 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Test that MessageQueue uses the local node's IP for binding,
|
||||
not a remote master_addr. This validates the fix for cross-node
|
||||
data-parallel where each DP group leader must bind to its own IP.
|
||||
|
||||
The bug: multiproc_executor used `parallel_config.master_addr` as
|
||||
`connect_ip` for every DP group's MessageQueue. For DP groups whose
|
||||
leader is NOT on the master node, binding to master_addr fails with
|
||||
"Cannot assign requested address".
|
||||
|
||||
The fix: use `get_ip()` (local node IP) instead of `master_addr`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import zmq
|
||||
|
||||
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
||||
from vllm.utils.network_utils import get_ip
|
||||
|
||||
|
||||
def test_mq_bind_with_local_ip():
|
||||
"""MessageQueue with remote readers should successfully bind
|
||||
when connect_ip is the local node's IP."""
|
||||
# n_reader=2, n_local_reader=1 means 1 remote reader,
|
||||
# which triggers the remote ZMQ socket bind.
|
||||
mq = MessageQueue(
|
||||
n_reader=2,
|
||||
n_local_reader=1,
|
||||
connect_ip=get_ip(),
|
||||
)
|
||||
handle = mq.export_handle()
|
||||
assert handle.remote_subscribe_addr is not None
|
||||
# The bound address should contain our local IP
|
||||
local_ip = get_ip()
|
||||
assert (
|
||||
local_ip in handle.remote_subscribe_addr
|
||||
or f"[{local_ip}]" in handle.remote_subscribe_addr
|
||||
)
|
||||
del mq
|
||||
|
||||
|
||||
def test_mq_bind_with_non_local_ip_fails():
|
||||
"""MessageQueue should fail to bind when connect_ip is a
|
||||
non-local IP address (simulating the bug where master_addr
|
||||
from a different node was used)."""
|
||||
# Use a non-local IP that we definitely can't bind to.
|
||||
# 198.51.100.1 is from TEST-NET-2 (RFC 5737), never locally assigned.
|
||||
non_local_ip = "198.51.100.1"
|
||||
with pytest.raises(zmq.error.ZMQError, match="Cannot assign requested address"):
|
||||
MessageQueue(
|
||||
n_reader=2,
|
||||
n_local_reader=1,
|
||||
connect_ip=non_local_ip,
|
||||
)
|
||||
|
||||
|
||||
def test_mq_bind_defaults_to_local_ip():
|
||||
"""When connect_ip is None, MessageQueue should auto-detect
|
||||
the local IP and bind successfully."""
|
||||
mq = MessageQueue(
|
||||
n_reader=2,
|
||||
n_local_reader=1,
|
||||
connect_ip=None, # should fallback to get_ip()
|
||||
)
|
||||
handle = mq.export_handle()
|
||||
assert handle.remote_subscribe_addr is not None
|
||||
del mq
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mq_bind_with_local_ip()
|
||||
print("PASSED: test_mq_bind_with_local_ip")
|
||||
test_mq_bind_with_non_local_ip_fails()
|
||||
print("PASSED: test_mq_bind_with_non_local_ip_fails")
|
||||
test_mq_bind_defaults_to_local_ip()
|
||||
print("PASSED: test_mq_bind_defaults_to_local_ip")
|
||||
print("\nAll tests passed!")
|
||||
@@ -0,0 +1,64 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Make sure ray assigns GPU workers to the correct node.
|
||||
|
||||
Run:
|
||||
```sh
|
||||
cd $VLLM_PATH/tests
|
||||
|
||||
pytest distributed/test_multi_node_assignment.py
|
||||
```
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
from vllm import initialize_ray_cluster
|
||||
from vllm.config import ParallelConfig
|
||||
from vllm.utils.network_utils import get_ip
|
||||
from vllm.v1.executor.ray_utils import _wait_until_pg_removed
|
||||
|
||||
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not VLLM_MULTI_NODE, reason="Need at least 2 nodes to run the test."
|
||||
)
|
||||
def test_multi_node_assignment() -> None:
|
||||
# NOTE: important to keep this class definition here
|
||||
# to let ray use cloudpickle to serialize it.
|
||||
class Actor:
|
||||
def get_ip(self):
|
||||
return get_ip()
|
||||
|
||||
for _ in range(10):
|
||||
config = ParallelConfig(1, 2)
|
||||
initialize_ray_cluster(config)
|
||||
|
||||
current_ip = get_ip()
|
||||
workers = []
|
||||
for bundle_id, bundle in enumerate(config.placement_group.bundle_specs):
|
||||
if not bundle.get("GPU", 0):
|
||||
continue
|
||||
scheduling_strategy = PlacementGroupSchedulingStrategy(
|
||||
placement_group=config.placement_group,
|
||||
placement_group_capture_child_tasks=True,
|
||||
placement_group_bundle_index=bundle_id,
|
||||
)
|
||||
|
||||
worker = ray.remote(
|
||||
num_cpus=0,
|
||||
num_gpus=1,
|
||||
scheduling_strategy=scheduling_strategy,
|
||||
)(Actor).remote()
|
||||
worker_ip = ray.get(worker.get_ip.remote())
|
||||
assert worker_ip == current_ip
|
||||
workers.append(worker)
|
||||
|
||||
for worker in workers:
|
||||
ray.kill(worker)
|
||||
|
||||
_wait_until_pg_removed(config.placement_group)
|
||||
@@ -0,0 +1,444 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Integration tests for MultiprocExecutor at the executor level.
|
||||
This test directly tests the executor without going through the LLM interface,
|
||||
focusing on executor initialization, RPC calls, and distributed execution.
|
||||
"""
|
||||
|
||||
import multiprocessing
|
||||
import os
|
||||
import socket
|
||||
|
||||
from tests.utils import multi_gpu_test
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.executor.multiproc_executor import MultiprocExecutor
|
||||
|
||||
MODEL = "facebook/opt-125m"
|
||||
|
||||
|
||||
def create_vllm_config(
|
||||
tensor_parallel_size: int = 1,
|
||||
pipeline_parallel_size: int = 1,
|
||||
max_model_len: int = 256,
|
||||
gpu_memory_utilization: float = 0.3,
|
||||
distributed_executor_backend: str = "mp",
|
||||
nnodes: int = 1,
|
||||
node_rank: int = 0,
|
||||
master_port: int = 0,
|
||||
) -> VllmConfig:
|
||||
"""Create a VllmConfig for testing using EngineArgs."""
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
pipeline_parallel_size=pipeline_parallel_size,
|
||||
max_model_len=max_model_len,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
enforce_eager=True,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
|
||||
# Override distributed node settings if needed
|
||||
if nnodes > 1 or node_rank > 0:
|
||||
vllm_config.parallel_config.nnodes = nnodes
|
||||
vllm_config.parallel_config.node_rank = node_rank
|
||||
vllm_config.parallel_config.master_port = master_port
|
||||
if nnodes > 1:
|
||||
vllm_config.parallel_config.disable_custom_all_reduce = True
|
||||
|
||||
return vllm_config
|
||||
|
||||
|
||||
def create_test_scheduler_output(num_requests: int = 1) -> SchedulerOutput:
|
||||
"""Create a minimal SchedulerOutput for testing."""
|
||||
# This is a simplified version - in practice you'd need proper
|
||||
# SchedulerOutput construction based on the actual vLLM v1 API
|
||||
return SchedulerOutput(
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_resumed_reqs=[],
|
||||
scheduled_running_reqs=[],
|
||||
num_scheduled_tokens={},
|
||||
total_num_scheduled_tokens=0,
|
||||
)
|
||||
|
||||
|
||||
def test_multiproc_executor_initialization():
|
||||
"""Test that MultiprocExecutor can be initialized with proper config."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
# Create executor - this should initialize workers
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
# Verify executor properties
|
||||
assert executor.world_size == 1, "World size should be 1 for single GPU"
|
||||
assert executor.local_world_size == 1, "Local world size should be 1"
|
||||
assert hasattr(executor, "workers"), "Executor should have workers"
|
||||
assert len(executor.workers) == 1, "Should have 1 worker for single GPU"
|
||||
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_multiproc_executor_initialization_tensor_parallel():
|
||||
"""Test MultiprocExecutor initialization with tensor parallelism."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
# Create executor
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
# Verify executor properties
|
||||
assert executor.world_size == 2, "World size should be 2 for TP=2"
|
||||
assert executor.local_world_size == 2, "Local world size should be 2"
|
||||
assert len(executor.workers) == 2, "Should have 2 workers for TP=2"
|
||||
|
||||
# Verify output rank calculation
|
||||
output_rank = executor._get_output_rank()
|
||||
assert output_rank == 0, "Output rank should be 0 for TP=2, PP=1"
|
||||
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_multiproc_executor_collective_rpc():
|
||||
"""Test collective RPC calls to all workers."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
# Create executor
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
try:
|
||||
# Test check_health RPC - should work without errors
|
||||
executor.check_health()
|
||||
|
||||
# Test that RPC works correctly
|
||||
# Note: We're just testing that the RPC mechanism works,
|
||||
# not testing actual model execution here
|
||||
assert not executor.is_failed, "Executor should not be in failed state"
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
def test_multiproc_executor_failure_callback():
|
||||
"""Test failure callback registration and invocation."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
try:
|
||||
# Test callback registration
|
||||
callback_invoked = []
|
||||
|
||||
def test_callback():
|
||||
callback_invoked.append(True)
|
||||
|
||||
# Register callback
|
||||
executor.register_failure_callback(test_callback)
|
||||
|
||||
# Callback should not be invoked yet
|
||||
assert len(callback_invoked) == 0, "Callback should not be invoked immediately"
|
||||
|
||||
# Simulate failure
|
||||
executor.is_failed = True
|
||||
|
||||
# Register another callback - should be invoked immediately
|
||||
executor.register_failure_callback(test_callback)
|
||||
assert len(callback_invoked) == 1, (
|
||||
"Callback should be invoked when executor is failed"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_multiproc_executor_worker_monitor():
|
||||
"""Test that worker monitor is set up correctly."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
try:
|
||||
# Verify all worker processes are alive
|
||||
for worker in executor.workers:
|
||||
assert worker.proc.is_alive(), f"Worker rank {worker.rank} should be alive"
|
||||
|
||||
# Verify executor is not in failed state
|
||||
assert not executor.is_failed, "Executor should not be in failed state"
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
# After shutdown, workers should be terminated
|
||||
import time
|
||||
|
||||
time.sleep(0.5) # Give processes time to terminate
|
||||
for worker in executor.workers:
|
||||
assert not worker.proc.is_alive(), (
|
||||
f"Worker rank {worker.rank} should terminate after shutdown"
|
||||
)
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
def test_multiproc_executor_get_response_message_queues():
|
||||
"""Test message queue retrieval for different ranks."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
try:
|
||||
# Get all message queues
|
||||
all_queues = executor.get_response_mqs()
|
||||
assert len(all_queues) == 2, "Should have 2 message queues for 2 workers"
|
||||
|
||||
# Get message queue for specific rank
|
||||
rank0_queue = executor.get_response_mqs(unique_reply_rank=0)
|
||||
assert len(rank0_queue) == 1, "Should have 1 message queue for rank 0"
|
||||
|
||||
rank1_queue = executor.get_response_mqs(unique_reply_rank=1)
|
||||
assert len(rank1_queue) == 1, "Should have 1 message queue for rank 1"
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
def test_multiproc_executor_shutdown_cleanup():
|
||||
"""Test that shutdown properly cleans up resources."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
# Verify executor is set up
|
||||
assert hasattr(executor, "workers"), "Executor should have workers"
|
||||
assert len(executor.workers) > 0, "Should have at least one worker"
|
||||
|
||||
# Shutdown
|
||||
executor.shutdown()
|
||||
|
||||
# Verify cleanup
|
||||
import time
|
||||
|
||||
time.sleep(0.5) # Give processes time to terminate
|
||||
|
||||
for worker in executor.workers:
|
||||
assert not worker.proc.is_alive(), "Worker processes should be terminated"
|
||||
|
||||
# Verify shutdown event is set
|
||||
assert executor.shutdown_event.is_set(), "Shutdown event should be set"
|
||||
|
||||
# Multiple shutdowns should be safe (idempotent)
|
||||
executor.shutdown()
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_multiproc_executor_pipeline_parallel():
|
||||
"""Test MultiprocExecutor with pipeline parallelism."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2,
|
||||
pipeline_parallel_size=2,
|
||||
)
|
||||
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
try:
|
||||
# Verify executor properties
|
||||
assert executor.world_size == 4, "World size should be 4 for TP=2, PP=2"
|
||||
assert len(executor.workers) == 4, "Should have 4 workers"
|
||||
|
||||
# Verify output rank calculation
|
||||
# For TP=2, PP=2: output should be from the last PP stage (ranks 2-3)
|
||||
# Specifically rank 2 (first rank of last PP stage)
|
||||
output_rank = executor._get_output_rank()
|
||||
assert output_rank == 2, "Output rank should be 2 (first rank of last PP stage)"
|
||||
|
||||
# V2 model runner uses one extra batch to overlap async scheduling.
|
||||
expected_concurrent_batches = 2 + int(
|
||||
vllm_config.scheduler_config.async_scheduling
|
||||
and vllm_config.use_v2_model_runner
|
||||
)
|
||||
assert vllm_config.max_concurrent_batches == expected_concurrent_batches, (
|
||||
"Max concurrent batches should follow the configured PP/async "
|
||||
"scheduling policy"
|
||||
)
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
def test_multiproc_executor_properties():
|
||||
"""Test various executor properties and configurations."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=1,
|
||||
pipeline_parallel_size=1,
|
||||
)
|
||||
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
try:
|
||||
# Test supports_pp property
|
||||
assert MultiprocExecutor.supports_pp is True, (
|
||||
"MultiprocExecutor should support pipeline parallelism"
|
||||
)
|
||||
|
||||
# Test world_size calculation
|
||||
assert executor.world_size == (
|
||||
executor.parallel_config.tensor_parallel_size
|
||||
* executor.parallel_config.pipeline_parallel_size
|
||||
), "World size should equal TP * PP"
|
||||
|
||||
# Test local_world_size calculation
|
||||
assert executor.local_world_size == (
|
||||
executor.parallel_config.world_size // executor.parallel_config.nnodes
|
||||
), "Local world size should be world_size / nnodes"
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
def test_multiproc_executor_multi_node():
|
||||
"""
|
||||
Test MultiprocExecutor with multi-node configuration.
|
||||
This simulates 2 nodes with TP=4:
|
||||
- Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2
|
||||
- Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
|
||||
Total world_size = 4, nnodes = 2
|
||||
"""
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
port = s.getsockname()[1]
|
||||
# symm_mem does not work for simulating multi instance in single node
|
||||
os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
|
||||
|
||||
def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int):
|
||||
"""Run a single node's executor."""
|
||||
executor = None
|
||||
try:
|
||||
# Set CUDA_VISIBLE_DEVICES for this node
|
||||
if node_rank == 0:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
|
||||
else:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
|
||||
|
||||
# Create config for this node
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=4, # Total TP across all nodes
|
||||
pipeline_parallel_size=1,
|
||||
nnodes=2, # 2 nodes
|
||||
node_rank=node_rank,
|
||||
master_port=port, # same port
|
||||
)
|
||||
|
||||
# Create executor for this node
|
||||
executor = MultiprocExecutor(vllm_config=vllm_config)
|
||||
|
||||
# Verify node-specific properties
|
||||
assert executor.world_size == 4, (
|
||||
f"World size should be 4 on node {node_rank}"
|
||||
)
|
||||
assert executor.local_world_size == 2, (
|
||||
f"Local world size should be 2 on node {node_rank}"
|
||||
)
|
||||
assert len(executor.workers) == 2, (
|
||||
f"Should have 2 local workers on node {node_rank}"
|
||||
)
|
||||
|
||||
# Verify worker ranks are correct for this node
|
||||
expected_ranks = [node_rank * 2, node_rank * 2 + 1]
|
||||
actual_ranks = sorted([w.rank for w in executor.workers])
|
||||
assert actual_ranks == expected_ranks, (
|
||||
f"Node {node_rank} should have workers "
|
||||
f"with ranks {expected_ranks}, got {actual_ranks}"
|
||||
)
|
||||
# Verify all workers are alive
|
||||
for worker in executor.workers:
|
||||
assert worker.proc.is_alive(), (
|
||||
f"Worker rank {worker.rank} should be alive on node {node_rank}"
|
||||
)
|
||||
# executor.gen
|
||||
# Put success result in queue BEFORE shutdown to avoid hanging
|
||||
result_queue.put({"node": node_rank, "success": True})
|
||||
import time
|
||||
|
||||
time.sleep(2)
|
||||
executor.shutdown()
|
||||
except Exception as e:
|
||||
# Put failure result in queue
|
||||
result_queue.put({"node": node_rank, "success": False, "error": str(e)})
|
||||
raise e
|
||||
finally:
|
||||
if executor is not None:
|
||||
executor.shutdown()
|
||||
|
||||
# Create a queue to collect results from both processes
|
||||
result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue()
|
||||
|
||||
# Start both node processes
|
||||
processes = []
|
||||
for node_rank in range(2):
|
||||
p = multiprocessing.Process(
|
||||
target=run_node,
|
||||
args=(node_rank, result_queue, port),
|
||||
name=f"Node{node_rank}",
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
# Wait for both processes to complete
|
||||
all_completed = True
|
||||
for p in processes:
|
||||
p.join(timeout=60)
|
||||
if p.is_alive():
|
||||
p.terminate()
|
||||
p.join(timeout=20)
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
p.join()
|
||||
all_completed = False
|
||||
|
||||
# Check results from both nodes
|
||||
results: list[dict[str, int | bool]] = []
|
||||
while len(results) < 2:
|
||||
try:
|
||||
result = result_queue.get(timeout=1)
|
||||
results.append(result)
|
||||
except Exception:
|
||||
pass
|
||||
assert all_completed, "Not all processes completed successfully"
|
||||
assert len(results) == 2, f"Expected 2 results, got {len(results)}"
|
||||
assert results[0]["success"], f"Node 0 failed: {results[0]}"
|
||||
assert results[1]["success"], f"Node 1 failed: {results[1]}"
|
||||
@@ -0,0 +1,223 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
import typing
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.utils import ensure_current_vllm_config
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
|
||||
from vllm.distributed.device_communicators.pynccl import register_nccl_symmetric_ops
|
||||
from vllm.distributed.device_communicators.pynccl_allocator import (
|
||||
get_nccl_mem_pool,
|
||||
is_symmetric_memory_enabled,
|
||||
)
|
||||
from vllm.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
torch.manual_seed(42)
|
||||
random.seed(44)
|
||||
|
||||
test_size_elements = 4 * 1024 * 1024
|
||||
|
||||
|
||||
def nccl_symm_mem_allreduce_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
pynccl_comm = cuda_communicator.pynccl_comm
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory allreduce is disabled.")
|
||||
|
||||
register_nccl_symmetric_ops(pynccl_comm)
|
||||
input = torch.randint(1, 23, (test_size_elements,), dtype=dtype, device=device)
|
||||
input_clone = input.clone()
|
||||
output = torch.ops.vllm.all_reduce_symmetric_with_copy(input)
|
||||
assert output is not None
|
||||
|
||||
group = get_tp_group().device_group
|
||||
dist.all_reduce(input_clone, group=group)
|
||||
torch.testing.assert_close(output, input_clone, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCLSymmMemAllreduce is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
# Enable SymmMemCommunicator
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_allreduce_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def nccl_symm_mem_allgather_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12346",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory is disabled.")
|
||||
|
||||
per_rank_size = test_size_elements // world_size
|
||||
input_tensor = torch.randint(
|
||||
1, 23, (per_rank_size,), dtype=dtype, device=device
|
||||
)
|
||||
output = cuda_communicator.all_gatherv(input_tensor, dim=0)
|
||||
|
||||
group = get_tp_group().device_group
|
||||
expected = torch.empty(test_size_elements, dtype=dtype, device=device)
|
||||
dist.all_gather_into_tensor(expected, input_tensor, group=group)
|
||||
torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCL symmetric memory is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_allgather(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_allgather_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def nccl_symm_mem_reduce_scatter_worker(local_rank: int, world_size: int):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12347",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
with ensure_current_vllm_config():
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
if get_nccl_mem_pool() is None:
|
||||
pytest.skip(
|
||||
"NCCL allocator compilation failed (probably missing NCCL headers)."
|
||||
)
|
||||
if not is_symmetric_memory_enabled():
|
||||
pytest.skip("NCCL symmetric memory is disabled.")
|
||||
|
||||
per_rank_size = test_size_elements // world_size
|
||||
input_tensor = torch.randint(
|
||||
1, 23, (test_size_elements,), dtype=dtype, device=device
|
||||
)
|
||||
input_clone = input_tensor.clone()
|
||||
output = cuda_communicator.reduce_scatter(input_tensor, dim=0)
|
||||
|
||||
group = get_tp_group().device_group
|
||||
expected = torch.empty(per_rank_size, dtype=dtype, device=device)
|
||||
dist.reduce_scatter_tensor(expected, input_clone, group=group)
|
||||
torch.testing.assert_close(output, expected, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="NCCL symmetric memory is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("world_size", [2])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_nccl_symm_mem_reduce_scatter(monkeypatch: pytest.MonkeyPatch, world_size):
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
monkeypatch.setenv("VLLM_USE_NCCL_SYMM_MEM", "1")
|
||||
monkeypatch.setenv("NCCL_NVLS_ENABLE", "1")
|
||||
monkeypatch.setenv("NCCL_CUMEM_ENABLE", "1")
|
||||
|
||||
mp.spawn(nccl_symm_mem_reduce_scatter_worker, args=(world_size,), nprocs=world_size)
|
||||
cleanup_dist_env_and_memory()
|
||||
@@ -0,0 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm.distributed.parallel_state import _node_count
|
||||
from vllm.distributed.utils import StatelessProcessGroup
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
if __name__ == "__main__":
|
||||
dist.init_process_group(backend="gloo")
|
||||
|
||||
rank = dist.get_rank()
|
||||
world_size = dist.get_world_size()
|
||||
|
||||
if rank == 0:
|
||||
port = get_open_port()
|
||||
ip = get_ip()
|
||||
dist.broadcast_object_list([ip, port], src=0)
|
||||
else:
|
||||
recv = [None, None]
|
||||
dist.broadcast_object_list(recv, src=0)
|
||||
ip, port = recv
|
||||
|
||||
stateless_pg = StatelessProcessGroup.create(ip, port, rank, world_size)
|
||||
|
||||
for pg in [dist.group.WORLD, stateless_pg]:
|
||||
test_result = _node_count(pg)
|
||||
|
||||
# Expected node count based on environment variable)
|
||||
expected = int(os.environ.get("NUM_NODES", "1"))
|
||||
|
||||
assert test_result == expected, f"Expected {expected} nodes, got {test_result}"
|
||||
|
||||
if pg == dist.group.WORLD:
|
||||
print(
|
||||
f"Node count test passed! Got {test_result} nodes "
|
||||
f"when using torch distributed!"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"Node count test passed! Got {test_result} nodes "
|
||||
f"when using StatelessProcessGroup!"
|
||||
)
|
||||
@@ -0,0 +1,775 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for packed tensor broadcasting functionality.
|
||||
|
||||
Unit tests for packed_nccl_broadcast_producer and packed_nccl_broadcast_consumer.
|
||||
These utilities enable efficient batched tensor transfer over NCCL.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferUpdateInfo
|
||||
from vllm.distributed.weight_transfer.packed_tensor import (
|
||||
pack_tensors,
|
||||
packed_ipc_consumer,
|
||||
packed_ipc_producer,
|
||||
packed_nccl_broadcast_consumer,
|
||||
packed_nccl_broadcast_producer,
|
||||
unpack_tensor,
|
||||
)
|
||||
|
||||
|
||||
class MockCommunicationGroup:
|
||||
"""Mock communication group for testing producer broadcast operations."""
|
||||
|
||||
def __init__(self):
|
||||
self.broadcasted_tensors: list[torch.Tensor] = []
|
||||
self.broadcast_count = 0
|
||||
self.device = torch.device("cuda:0")
|
||||
|
||||
def broadcast(self, tensor, src):
|
||||
"""Mock broadcast that stores the tensor for later verification."""
|
||||
self.broadcasted_tensors.append(tensor.clone())
|
||||
self.broadcast_count += 1
|
||||
|
||||
|
||||
class MockConsumerCommunicationGroup:
|
||||
"""Mock communication group for consumer that returns pre-stored tensors."""
|
||||
|
||||
def __init__(self, tensors_to_return: list[torch.Tensor]):
|
||||
self.tensors_to_return = tensors_to_return
|
||||
self.current_index = 0
|
||||
self.device = torch.device("cuda:0")
|
||||
|
||||
def broadcast(self, tensor, src):
|
||||
"""Mock broadcast that fills the tensor with pre-stored data."""
|
||||
if self.current_index < len(self.tensors_to_return):
|
||||
tensor.copy_(self.tensors_to_return[self.current_index])
|
||||
self.current_index += 1
|
||||
|
||||
|
||||
def create_mock_model_params(
|
||||
num_layers: int = 3,
|
||||
dtype: torch.dtype = torch.float32,
|
||||
) -> list[tuple[str, torch.Tensor]]:
|
||||
"""Create mock model parameters for testing."""
|
||||
params = []
|
||||
for i in range(num_layers):
|
||||
params.append((f"layer{i}.weight", torch.randn(10, 20, dtype=dtype)))
|
||||
params.append((f"layer{i}.bias", torch.randn(10, dtype=dtype)))
|
||||
return params
|
||||
|
||||
|
||||
def create_state_dict_info(
|
||||
params: list[tuple[str, torch.Tensor]],
|
||||
) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
|
||||
"""Create state dict info (name -> (shape, dtype)) from params."""
|
||||
return {name: (tuple(tensor.shape), tensor.dtype) for name, tensor in params}
|
||||
|
||||
|
||||
# --- Unit Tests: NCCLWeightTransferUpdateInfo packed field ---
|
||||
|
||||
|
||||
class TestNCCLWeightTransferUpdateInfoPacked:
|
||||
"""Test NCCLWeightTransferUpdateInfo dataclass packed field."""
|
||||
|
||||
def test_packed_default_false(self):
|
||||
"""Test that packed defaults to False."""
|
||||
info = NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
)
|
||||
assert info.packed is False
|
||||
|
||||
def test_packed_can_be_set_true(self):
|
||||
"""Test that packed can be set to True."""
|
||||
info = NCCLWeightTransferUpdateInfo(
|
||||
names=["layer.weight"],
|
||||
dtype_names=["float32"],
|
||||
shapes=[[10, 10]],
|
||||
packed=True,
|
||||
)
|
||||
assert info.packed is True
|
||||
|
||||
|
||||
# --- Unit Tests: packed_nccl_broadcast_producer ---
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
class TestPackedBroadcastProducer:
|
||||
"""Test packed_nccl_broadcast_producer function."""
|
||||
|
||||
def test_producer_empty_iterator(self):
|
||||
"""Test producer handles empty iterator gracefully."""
|
||||
mock_group = MockCommunicationGroup()
|
||||
|
||||
packed_nccl_broadcast_producer(
|
||||
iterator=iter([]),
|
||||
group=mock_group,
|
||||
src=0,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=1000,
|
||||
)
|
||||
|
||||
# No broadcasts for empty iterator
|
||||
assert mock_group.broadcast_count == 0
|
||||
|
||||
|
||||
# --- Integration Tests: Producer-Consumer Roundtrip ---
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
class TestPackedBroadcastRoundtrip:
|
||||
"""Test producer-consumer roundtrip behavior."""
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_roundtrip_different_dtypes(self, dtype):
|
||||
"""Test roundtrip with different data types."""
|
||||
params = create_mock_model_params(num_layers=2, dtype=dtype)
|
||||
params_cuda = [(name, tensor.cuda()) for name, tensor in params]
|
||||
|
||||
buffer_size = 1000
|
||||
producer_group = MockCommunicationGroup()
|
||||
|
||||
packed_nccl_broadcast_producer(
|
||||
iterator=iter(params_cuda),
|
||||
group=producer_group,
|
||||
src=0,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=buffer_size,
|
||||
)
|
||||
|
||||
consumer_group = MockConsumerCommunicationGroup(
|
||||
producer_group.broadcasted_tensors
|
||||
)
|
||||
|
||||
state_dict_info = create_state_dict_info(params_cuda)
|
||||
unpacked_tensors = {}
|
||||
|
||||
def post_unpack_func(tensor_list):
|
||||
for name, tensor in tensor_list:
|
||||
unpacked_tensors[name] = tensor.clone()
|
||||
|
||||
packed_nccl_broadcast_consumer(
|
||||
iterator=iter(state_dict_info.items()),
|
||||
group=consumer_group,
|
||||
src=0,
|
||||
post_unpack_func=post_unpack_func,
|
||||
buffer_size_bytes=buffer_size,
|
||||
)
|
||||
|
||||
# Verify roundtrip preserves data
|
||||
for name, original_tensor in params_cuda:
|
||||
assert name in unpacked_tensors
|
||||
unpacked = unpacked_tensors[name]
|
||||
assert unpacked.dtype == dtype
|
||||
assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
|
||||
|
||||
def test_roundtrip_mixed_dtypes(self):
|
||||
"""Test roundtrip with mixed data types."""
|
||||
# Create params with mixed dtypes
|
||||
params = [
|
||||
("layer1.weight", torch.randn(10, 20, dtype=torch.float32).cuda()),
|
||||
("layer1.bias", torch.randn(10, dtype=torch.float16).cuda()),
|
||||
("layer2.weight", torch.randn(20, 30, dtype=torch.bfloat16).cuda()),
|
||||
]
|
||||
|
||||
buffer_size = 500
|
||||
producer_group = MockCommunicationGroup()
|
||||
|
||||
packed_nccl_broadcast_producer(
|
||||
iterator=iter(params),
|
||||
group=producer_group,
|
||||
src=0,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=buffer_size,
|
||||
)
|
||||
|
||||
consumer_group = MockConsumerCommunicationGroup(
|
||||
producer_group.broadcasted_tensors
|
||||
)
|
||||
|
||||
state_dict_info = create_state_dict_info(params)
|
||||
unpacked_tensors = {}
|
||||
|
||||
def post_unpack_func(tensor_list):
|
||||
for name, tensor in tensor_list:
|
||||
unpacked_tensors[name] = tensor.clone()
|
||||
|
||||
packed_nccl_broadcast_consumer(
|
||||
iterator=iter(state_dict_info.items()),
|
||||
group=consumer_group,
|
||||
src=0,
|
||||
post_unpack_func=post_unpack_func,
|
||||
buffer_size_bytes=buffer_size,
|
||||
)
|
||||
|
||||
# Verify all params roundtrip correctly with correct dtypes
|
||||
for name, original_tensor in params:
|
||||
assert name in unpacked_tensors
|
||||
unpacked = unpacked_tensors[name]
|
||||
assert unpacked.shape == original_tensor.shape
|
||||
assert unpacked.dtype == original_tensor.dtype
|
||||
assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
|
||||
|
||||
@pytest.mark.parametrize("target_size", [100, 100000])
|
||||
def test_roundtrip_different_batch_sizes(self, target_size):
|
||||
"""Test roundtrip with different target batch sizes."""
|
||||
params = create_mock_model_params(num_layers=5)
|
||||
params_cuda = [(name, tensor.cuda()) for name, tensor in params]
|
||||
|
||||
producer_group = MockCommunicationGroup()
|
||||
|
||||
packed_nccl_broadcast_producer(
|
||||
iterator=iter(params_cuda),
|
||||
group=producer_group,
|
||||
src=0,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=target_size,
|
||||
)
|
||||
|
||||
consumer_group = MockConsumerCommunicationGroup(
|
||||
producer_group.broadcasted_tensors
|
||||
)
|
||||
|
||||
state_dict_info = create_state_dict_info(params_cuda)
|
||||
unpacked_tensors = {}
|
||||
|
||||
def post_unpack_func(tensor_list):
|
||||
for name, tensor in tensor_list:
|
||||
unpacked_tensors[name] = tensor.clone()
|
||||
|
||||
packed_nccl_broadcast_consumer(
|
||||
iterator=iter(state_dict_info.items()),
|
||||
group=consumer_group,
|
||||
src=0,
|
||||
post_unpack_func=post_unpack_func,
|
||||
buffer_size_bytes=target_size,
|
||||
)
|
||||
|
||||
# Verify all params roundtrip correctly
|
||||
assert len(unpacked_tensors) == len(params)
|
||||
for name, original_tensor in params_cuda:
|
||||
assert name in unpacked_tensors
|
||||
assert torch.allclose(
|
||||
unpacked_tensors[name], original_tensor, rtol=1e-5, atol=1e-7
|
||||
)
|
||||
|
||||
def test_roundtrip_non_contiguous_tensors(self):
|
||||
"""Test roundtrip with non-contiguous tensors from the trainer."""
|
||||
# Create non-contiguous tensors (simulating trainer outputs)
|
||||
# Transposed tensors are non-contiguous
|
||||
weight1 = torch.randn(20, 10, dtype=torch.float32).cuda().T
|
||||
# Sliced tensors with step are non-contiguous
|
||||
weight2 = torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2]
|
||||
# Permuted tensors are non-contiguous
|
||||
weight3 = torch.randn(5, 10, 15, dtype=torch.bfloat16).cuda().permute(2, 0, 1)
|
||||
|
||||
params = [
|
||||
("layer1.weight", weight1),
|
||||
("layer2.weight", weight2),
|
||||
("layer3.weight", weight3),
|
||||
]
|
||||
|
||||
# Verify tensors are indeed non-contiguous
|
||||
for name, tensor in params:
|
||||
assert not tensor.is_contiguous(), f"{name} should be non-contiguous"
|
||||
|
||||
buffer_size = 500
|
||||
producer_group = MockCommunicationGroup()
|
||||
|
||||
packed_nccl_broadcast_producer(
|
||||
iterator=iter(params),
|
||||
group=producer_group,
|
||||
src=0,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=buffer_size,
|
||||
)
|
||||
|
||||
consumer_group = MockConsumerCommunicationGroup(
|
||||
producer_group.broadcasted_tensors
|
||||
)
|
||||
|
||||
state_dict_info = create_state_dict_info(params)
|
||||
unpacked_tensors = {}
|
||||
|
||||
def post_unpack_func(tensor_list):
|
||||
for name, tensor in tensor_list:
|
||||
unpacked_tensors[name] = tensor.clone()
|
||||
|
||||
packed_nccl_broadcast_consumer(
|
||||
iterator=iter(state_dict_info.items()),
|
||||
group=consumer_group,
|
||||
src=0,
|
||||
post_unpack_func=post_unpack_func,
|
||||
buffer_size_bytes=buffer_size,
|
||||
)
|
||||
|
||||
# Verify all non-contiguous params roundtrip correctly
|
||||
for name, original_tensor in params:
|
||||
assert name in unpacked_tensors
|
||||
unpacked = unpacked_tensors[name]
|
||||
assert unpacked.shape == original_tensor.shape
|
||||
assert unpacked.dtype == original_tensor.dtype
|
||||
assert torch.allclose(unpacked, original_tensor, rtol=1e-4, atol=1e-6)
|
||||
|
||||
|
||||
# --- Unit Tests: unpack_tensor ---
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
class TestUnpackTensor:
|
||||
"""Test the shared unpack_tensor function."""
|
||||
|
||||
def test_unpack_produces_independent_copies(self):
|
||||
"""Verify unpacked tensors don't share memory with packed buffer."""
|
||||
original = torch.randn(10, dtype=torch.float32).cuda()
|
||||
packed = original.contiguous().view(torch.uint8).view(-1)
|
||||
|
||||
result = unpack_tensor(
|
||||
packed,
|
||||
names=["w"],
|
||||
shapes=[[10]],
|
||||
dtypes=[torch.float32],
|
||||
tensor_sizes=[packed.numel()],
|
||||
)
|
||||
|
||||
# Mutate the packed buffer
|
||||
packed.zero_()
|
||||
|
||||
# Unpacked tensor should be unaffected
|
||||
assert torch.allclose(result[0][1], original)
|
||||
|
||||
|
||||
# --- Unit Tests: pack_tensors ---
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
class TestPackTensors:
|
||||
"""Test the shared pack_tensors function."""
|
||||
|
||||
def test_pack_basic(self):
|
||||
"""Test packing a few tensors into one buffer."""
|
||||
params = [
|
||||
("w1", torch.randn(10, 20, dtype=torch.float32).cuda()),
|
||||
("w2", torch.randn(5, dtype=torch.float16).cuda()),
|
||||
]
|
||||
|
||||
chunk = pack_tensors(
|
||||
iterator=iter(params),
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
)
|
||||
|
||||
assert chunk is not None
|
||||
assert len(chunk.names) == 2
|
||||
assert chunk.names == ["w1", "w2"]
|
||||
assert chunk.shapes == [[10, 20], [5]]
|
||||
assert chunk.dtypes == [torch.float32, torch.float16]
|
||||
assert chunk.packed_tensor.dtype == torch.uint8
|
||||
|
||||
def test_pack_respects_buffer_limit(self):
|
||||
"""Test that packing stops when buffer_size_bytes is exceeded."""
|
||||
params = [
|
||||
(f"w{i}", torch.randn(100, 100, dtype=torch.float32).cuda())
|
||||
for i in range(10)
|
||||
]
|
||||
|
||||
chunk = pack_tensors(
|
||||
iterator=iter(params),
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=50_000,
|
||||
)
|
||||
|
||||
assert chunk is not None
|
||||
assert len(chunk.names) < 10
|
||||
|
||||
def test_pack_empty_iterator(self):
|
||||
"""Test that an empty iterator returns None."""
|
||||
chunk = pack_tensors(
|
||||
iterator=iter([]),
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=1000,
|
||||
)
|
||||
assert chunk is None
|
||||
|
||||
def test_pack_single_tensor_larger_than_buffer_warns(self):
|
||||
"""Test that a tensor exceeding buffer_size_bytes emits a warning."""
|
||||
big = torch.randn(1000, 1000, dtype=torch.float32).cuda()
|
||||
params = [("big", big)]
|
||||
|
||||
import warnings
|
||||
|
||||
with warnings.catch_warnings(record=True) as w:
|
||||
warnings.simplefilter("always")
|
||||
chunk = pack_tensors(
|
||||
iterator=iter(params),
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=100,
|
||||
)
|
||||
assert chunk is not None
|
||||
assert len(chunk.names) == 1
|
||||
assert any("exceeds buffer_size_bytes" in str(wi.message) for wi in w)
|
||||
|
||||
def test_pack_unpack_roundtrip(self):
|
||||
"""Test pack then unpack produces identical tensors."""
|
||||
params = [
|
||||
("a", torch.randn(8, 16, dtype=torch.float32).cuda()),
|
||||
("b", torch.randn(4, dtype=torch.float16).cuda()),
|
||||
("c", torch.randn(3, 5, 7, dtype=torch.bfloat16).cuda()),
|
||||
]
|
||||
|
||||
chunk = pack_tensors(
|
||||
iterator=iter(params),
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
)
|
||||
|
||||
assert chunk is not None
|
||||
result = unpack_tensor(
|
||||
chunk.packed_tensor,
|
||||
chunk.names,
|
||||
chunk.shapes,
|
||||
chunk.dtypes,
|
||||
chunk.tensor_sizes,
|
||||
)
|
||||
|
||||
assert len(result) == len(params)
|
||||
for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
|
||||
assert orig_name == res_name
|
||||
assert res_tensor.shape == orig_tensor.shape
|
||||
assert res_tensor.dtype == orig_tensor.dtype
|
||||
assert torch.allclose(res_tensor, orig_tensor, rtol=1e-4, atol=1e-6)
|
||||
|
||||
def test_pack_multiple_chunks(self):
|
||||
"""Test consuming an iterator across multiple pack_tensors calls."""
|
||||
params = [
|
||||
(f"w{i}", torch.randn(50, 50, dtype=torch.float32).cuda()) for i in range(6)
|
||||
]
|
||||
it = iter(params)
|
||||
|
||||
all_names = []
|
||||
chunks = []
|
||||
while True:
|
||||
chunk = pack_tensors(it, lambda x: x[1], buffer_size_bytes=12_000)
|
||||
if chunk is None:
|
||||
break
|
||||
chunks.append(chunk)
|
||||
all_names.extend(chunk.names)
|
||||
|
||||
assert len(chunks) > 1
|
||||
assert all_names == [f"w{i}" for i in range(6)]
|
||||
|
||||
|
||||
# --- Unit Tests: packed_ipc_producer ---
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
class TestPackedIpcProducer:
|
||||
"""Test the packed_ipc_producer generator."""
|
||||
|
||||
def test_producer_yields_chunks(self):
|
||||
"""Test that the producer yields PackedIpcChunk objects."""
|
||||
params = [
|
||||
(f"w{i}", torch.randn(50, 50, dtype=torch.float32).cuda()) for i in range(6)
|
||||
]
|
||||
|
||||
chunks = list(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid="test-uuid",
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=12_000,
|
||||
)
|
||||
)
|
||||
|
||||
assert len(chunks) > 1
|
||||
|
||||
def test_producer_ipc_handle_has_uuid(self):
|
||||
"""Test that each chunk's ipc_handle is keyed by the given UUID."""
|
||||
params = [("w", torch.randn(10, dtype=torch.float32).cuda())]
|
||||
|
||||
chunks = list(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid="my-gpu-uuid",
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
)
|
||||
)
|
||||
|
||||
assert "my-gpu-uuid" in chunks[0].ipc_handle
|
||||
|
||||
def test_producer_dtype_names_are_strings(self):
|
||||
"""Test that dtype_names are string representations."""
|
||||
params = [
|
||||
("a", torch.randn(10, dtype=torch.float32).cuda()),
|
||||
("b", torch.randn(10, dtype=torch.float16).cuda()),
|
||||
]
|
||||
|
||||
chunks = list(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid="uuid",
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
)
|
||||
)
|
||||
|
||||
assert chunks[0].dtype_names == ["float32", "float16"]
|
||||
|
||||
def test_producer_empty_iterator(self):
|
||||
"""Test producer with empty iterator yields nothing."""
|
||||
chunks = list(
|
||||
packed_ipc_producer(
|
||||
iterator=iter([]),
|
||||
gpu_uuid="uuid",
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=1000,
|
||||
)
|
||||
)
|
||||
assert len(chunks) == 0
|
||||
|
||||
|
||||
# --- Integration Tests: IPC Producer-Consumer Roundtrip ---
|
||||
|
||||
|
||||
def _ipc_consumer_worker(cmd_q, ack_q, result_q, done_event, device_index):
|
||||
"""Worker that consumes chunks streamed one at a time from the parent.
|
||||
|
||||
CUDA IPC requires the consumer to be in a separate process from the
|
||||
producer. The producer reuses a single IPC buffer between chunks, so
|
||||
the parent must wait for our ack (sent after we copy the chunk to
|
||||
CPU) before advancing the producer.
|
||||
"""
|
||||
try:
|
||||
torch.accelerator.set_device_index(device_index)
|
||||
all_results = []
|
||||
while True:
|
||||
cd = cmd_q.get()
|
||||
if cd is None:
|
||||
break
|
||||
result = packed_ipc_consumer(
|
||||
ipc_handle=cd["ipc_handle"],
|
||||
names=cd["names"],
|
||||
shapes=cd["shapes"],
|
||||
dtype_names=cd["dtype_names"],
|
||||
tensor_sizes=cd["tensor_sizes"],
|
||||
device_index=device_index,
|
||||
)
|
||||
# .cpu() forces a GPU→CPU copy off the shared IPC buffer, so
|
||||
# the producer is free to overwrite it once we ack.
|
||||
all_results.extend([(name, tensor.cpu()) for name, tensor in result])
|
||||
del result
|
||||
ack_q.put("ack")
|
||||
result_q.put(("ok", all_results))
|
||||
except Exception as e:
|
||||
result_q.put(("error", str(e)))
|
||||
# Keep the process alive until the parent has finished reading from
|
||||
# the result queue — torch serializes CPU tensors via fd sharing,
|
||||
# which requires this process's resource-sharer server to be running.
|
||||
done_event.wait(timeout=60)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
class TestPackedIpcRoundtrip:
|
||||
"""Test IPC producer-consumer roundtrip using real CUDA IPC.
|
||||
|
||||
These tests spawn a child process for the consumer because
|
||||
rebuild_cuda_tensor requires a separate process from the one that
|
||||
called reduce_tensor.
|
||||
"""
|
||||
|
||||
def _get_gpu_uuid(self) -> str:
|
||||
device_index = torch.cuda.current_device()
|
||||
props = torch.cuda.get_device_properties(device_index)
|
||||
return str(props.uuid)
|
||||
|
||||
def _run_roundtrip(self, chunk_iter, device_index, timeout=30):
|
||||
"""Stream chunks through a child consumer one at a time.
|
||||
|
||||
``packed_ipc_producer`` reuses a single IPC buffer for every
|
||||
chunk, so the producer must not be advanced until the consumer
|
||||
has finished reading the current chunk. We enforce that with an
|
||||
ack queue: the consumer puts ``"ack"`` after it has copied the
|
||||
chunk to CPU, and only then do we pull the next chunk from the
|
||||
generator.
|
||||
|
||||
Returns ``(num_chunks, results)``.
|
||||
"""
|
||||
import multiprocessing as mp
|
||||
|
||||
ctx = mp.get_context("spawn")
|
||||
cmd_q = ctx.Queue()
|
||||
ack_q = ctx.Queue()
|
||||
result_q = ctx.Queue()
|
||||
done_event = ctx.Event()
|
||||
proc = ctx.Process(
|
||||
target=_ipc_consumer_worker,
|
||||
args=(cmd_q, ack_q, result_q, done_event, device_index),
|
||||
)
|
||||
proc.start()
|
||||
|
||||
num_chunks = 0
|
||||
try:
|
||||
for chunk in chunk_iter:
|
||||
cmd_q.put(
|
||||
{
|
||||
"ipc_handle": chunk.ipc_handle,
|
||||
"names": chunk.names,
|
||||
"shapes": chunk.shapes,
|
||||
"dtype_names": chunk.dtype_names,
|
||||
"tensor_sizes": chunk.tensor_sizes,
|
||||
}
|
||||
)
|
||||
if ack_q.get(timeout=timeout) != "ack":
|
||||
raise RuntimeError("Consumer did not ack chunk")
|
||||
num_chunks += 1
|
||||
cmd_q.put(None)
|
||||
status, payload = result_q.get(timeout=timeout)
|
||||
finally:
|
||||
done_event.set()
|
||||
proc.join(timeout=10)
|
||||
if proc.is_alive():
|
||||
proc.kill()
|
||||
|
||||
if status == "error":
|
||||
raise RuntimeError(f"Consumer process failed: {payload}")
|
||||
# Reclaim IPC-shared memory now that the child has released it
|
||||
torch.cuda.ipc_collect()
|
||||
return num_chunks, payload
|
||||
|
||||
def test_roundtrip_basic(self):
|
||||
"""Test basic IPC producer -> consumer roundtrip."""
|
||||
params = [
|
||||
("w1", torch.randn(10, 20, dtype=torch.float32).cuda()),
|
||||
("w2", torch.randn(5, dtype=torch.float16).cuda()),
|
||||
]
|
||||
gpu_uuid = self._get_gpu_uuid()
|
||||
device_index = torch.cuda.current_device()
|
||||
|
||||
num_chunks, result = self._run_roundtrip(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid=gpu_uuid,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
),
|
||||
device_index,
|
||||
)
|
||||
|
||||
assert num_chunks == 1
|
||||
assert len(result) == 2
|
||||
for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
|
||||
assert orig_name == res_name
|
||||
assert res_tensor.shape == orig_tensor.shape
|
||||
assert res_tensor.dtype == orig_tensor.dtype
|
||||
assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_roundtrip_dtypes(self, dtype):
|
||||
"""Test IPC roundtrip with different dtypes."""
|
||||
params = create_mock_model_params(num_layers=2, dtype=dtype)
|
||||
params_cuda = [(n, t.cuda()) for n, t in params]
|
||||
gpu_uuid = self._get_gpu_uuid()
|
||||
device_index = torch.cuda.current_device()
|
||||
|
||||
_, result = self._run_roundtrip(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params_cuda),
|
||||
gpu_uuid=gpu_uuid,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
),
|
||||
device_index,
|
||||
)
|
||||
|
||||
assert len(result) == len(params_cuda)
|
||||
for (orig_name, orig_tensor), (res_name, res_tensor) in zip(
|
||||
params_cuda, result
|
||||
):
|
||||
assert orig_name == res_name
|
||||
assert res_tensor.dtype == dtype
|
||||
assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
|
||||
|
||||
def test_roundtrip_multiple_chunks(self):
|
||||
"""Test IPC roundtrip across multiple chunks."""
|
||||
params = [
|
||||
(f"layer{i}.weight", torch.randn(100, 100, dtype=torch.float32).cuda())
|
||||
for i in range(8)
|
||||
]
|
||||
gpu_uuid = self._get_gpu_uuid()
|
||||
device_index = torch.cuda.current_device()
|
||||
|
||||
num_chunks, result = self._run_roundtrip(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid=gpu_uuid,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=50_000,
|
||||
),
|
||||
device_index,
|
||||
)
|
||||
|
||||
assert num_chunks > 1
|
||||
assert len(result) == len(params)
|
||||
for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
|
||||
assert orig_name == res_name
|
||||
assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-5, atol=1e-7)
|
||||
|
||||
def test_roundtrip_non_contiguous(self):
|
||||
"""Test IPC roundtrip with non-contiguous tensors."""
|
||||
params = [
|
||||
("transposed", torch.randn(20, 10, dtype=torch.float32).cuda().T),
|
||||
("sliced", torch.randn(40, 30, dtype=torch.float16).cuda()[::2, ::2]),
|
||||
]
|
||||
gpu_uuid = self._get_gpu_uuid()
|
||||
device_index = torch.cuda.current_device()
|
||||
|
||||
for _, t in params:
|
||||
assert not t.is_contiguous()
|
||||
|
||||
_, result = self._run_roundtrip(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid=gpu_uuid,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
),
|
||||
device_index,
|
||||
)
|
||||
|
||||
for (orig_name, orig_tensor), (res_name, res_tensor) in zip(params, result):
|
||||
assert orig_name == res_name
|
||||
assert res_tensor.shape == orig_tensor.shape
|
||||
assert res_tensor.dtype == orig_tensor.dtype
|
||||
assert torch.allclose(res_tensor, orig_tensor.cpu(), rtol=1e-4, atol=1e-6)
|
||||
|
||||
def test_consumer_wrong_uuid_raises(self):
|
||||
"""Test that consumer raises ValueError for unknown GPU UUID."""
|
||||
params = [("w", torch.randn(10, dtype=torch.float32).cuda())]
|
||||
gpu_uuid = self._get_gpu_uuid()
|
||||
|
||||
chunks = list(
|
||||
packed_ipc_producer(
|
||||
iterator=iter(params),
|
||||
gpu_uuid=gpu_uuid,
|
||||
post_iter_func=lambda x: x[1],
|
||||
buffer_size_bytes=10_000_000,
|
||||
)
|
||||
)
|
||||
|
||||
c = chunks[0]
|
||||
fake_handle = {"fake-uuid-12345": c.ipc_handle[gpu_uuid]}
|
||||
|
||||
with pytest.raises(ValueError, match="IPC handle not found"):
|
||||
packed_ipc_consumer(
|
||||
ipc_handle=fake_handle,
|
||||
names=c.names,
|
||||
shapes=c.shapes,
|
||||
dtype_names=c.dtype_names,
|
||||
tensor_sizes=c.tensor_sizes,
|
||||
device_index=torch.cuda.current_device(),
|
||||
)
|
||||
@@ -0,0 +1,434 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
WARNING: This test runs in both single-node (4 GPUs) and multi-node
|
||||
(2 node with 2 GPUs each) modes. If the test only uses 2 GPUs, it is
|
||||
important to set the distributed backend to "mp" to avoid Ray scheduling
|
||||
all workers in a node other than the head node, which can cause the test
|
||||
to fail.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal, NamedTuple
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config.model import _FLOAT16_NOT_SUPPORTED_MODELS, RunnerOption
|
||||
from vllm.logger import init_logger
|
||||
from vllm.transformers_utils.config import get_config
|
||||
|
||||
from ..models.registry import HF_EXAMPLE_MODELS
|
||||
from ..utils import compare_two_settings, create_new_process_for_each_test
|
||||
|
||||
logger = init_logger("test_pipeline_parallel")
|
||||
|
||||
VLLM_MULTI_NODE = os.getenv("VLLM_MULTI_NODE", "0") == "1"
|
||||
|
||||
|
||||
class ParallelSetup(NamedTuple):
|
||||
tp_size: int
|
||||
pp_size: int
|
||||
eager_mode: bool
|
||||
|
||||
|
||||
class PPTestOptions(NamedTuple):
|
||||
multi_node_only: bool
|
||||
load_format: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PPTestSettings:
|
||||
parallel_setups: list[ParallelSetup]
|
||||
distributed_backends: list[str]
|
||||
runner: RunnerOption
|
||||
test_options: PPTestOptions
|
||||
|
||||
@staticmethod
|
||||
def detailed(
|
||||
*,
|
||||
tp_base: int = 1,
|
||||
pp_base: int = 2,
|
||||
multi_node_only: bool = False,
|
||||
runner: RunnerOption = "auto",
|
||||
load_format: str | None = None,
|
||||
):
|
||||
return PPTestSettings(
|
||||
parallel_setups=[
|
||||
ParallelSetup(tp_size=tp_base, pp_size=pp_base, eager_mode=False),
|
||||
ParallelSetup(tp_size=tp_base, pp_size=2 * pp_base, eager_mode=False),
|
||||
ParallelSetup(tp_size=tp_base, pp_size=2 * pp_base, eager_mode=True),
|
||||
ParallelSetup(tp_size=2 * tp_base, pp_size=pp_base, eager_mode=False),
|
||||
ParallelSetup(tp_size=2 * tp_base, pp_size=pp_base, eager_mode=True),
|
||||
],
|
||||
distributed_backends=["mp", "ray"],
|
||||
runner=runner,
|
||||
test_options=PPTestOptions(
|
||||
multi_node_only=multi_node_only, load_format=load_format
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def fast(
|
||||
*,
|
||||
tp_base: int = 1,
|
||||
pp_base: int = 2,
|
||||
runner: RunnerOption = "auto",
|
||||
multi_node_only: bool = False,
|
||||
load_format: str | None = None,
|
||||
):
|
||||
return PPTestSettings(
|
||||
parallel_setups=[
|
||||
ParallelSetup(tp_size=tp_base, pp_size=pp_base, eager_mode=True),
|
||||
],
|
||||
distributed_backends=["mp"],
|
||||
runner=runner,
|
||||
test_options=PPTestOptions(
|
||||
multi_node_only=multi_node_only, load_format=load_format
|
||||
),
|
||||
)
|
||||
|
||||
def iter_params(self, model_id: str):
|
||||
opts = self.test_options
|
||||
|
||||
for parallel_setup in self.parallel_setups:
|
||||
for backend in self.distributed_backends:
|
||||
yield (model_id, parallel_setup, backend, self.runner, opts)
|
||||
|
||||
|
||||
# NOTE: You can adjust tp_base and/or pp_base locally to fit the model in GPU
|
||||
# The values displayed here are only a rough indicator of the size of the model
|
||||
|
||||
TEXT_GENERATION_MODELS = {
|
||||
# [Decoder-only]
|
||||
"Snowflake/snowflake-arctic-instruct": PPTestSettings.fast(load_format="dummy"),
|
||||
"bigscience/bloomz-1b1": PPTestSettings.fast(),
|
||||
"zai-org/chatglm3-6b": PPTestSettings.fast(),
|
||||
"CohereLabs/c4ai-command-r-v01": PPTestSettings.fast(load_format="dummy"),
|
||||
"databricks/dbrx-instruct": PPTestSettings.fast(load_format="dummy"),
|
||||
"Deci/DeciLM-7B-instruct": PPTestSettings.fast(),
|
||||
"deepseek-ai/deepseek-llm-7b-chat": PPTestSettings.fast(),
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat": PPTestSettings.fast(tp_base=2),
|
||||
"LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct": PPTestSettings.fast(),
|
||||
"tiiuae/falcon-7b": PPTestSettings.fast(),
|
||||
"google/gemma-1.1-2b-it": PPTestSettings.fast(),
|
||||
"google/gemma-2-9b": PPTestSettings.fast(),
|
||||
"openai-community/gpt2": PPTestSettings.fast(),
|
||||
"EleutherAI/gpt-j-6b": PPTestSettings.fast(),
|
||||
"EleutherAI/pythia-1.4b": PPTestSettings.fast(),
|
||||
"ibm/PowerLM-3b": PPTestSettings.fast(),
|
||||
"ibm/PowerMoE-3b": PPTestSettings.fast(),
|
||||
"internlm/internlm2-chat-7b": PPTestSettings.fast(),
|
||||
"ai21labs/Jamba-tiny-dev": PPTestSettings.fast(),
|
||||
"pfnet/plamo-2-1b": PPTestSettings.fast(),
|
||||
"pfnet/plamo-3-nict-2b-base": PPTestSettings.fast(),
|
||||
"meta-llama/Llama-3.2-1B-Instruct": PPTestSettings.detailed(),
|
||||
# Tests TransformersForCausalLM
|
||||
"hmellor/Ilama-3.2-1B": PPTestSettings.fast(),
|
||||
"openbmb/MiniCPM-2B-sft-bf16": PPTestSettings.fast(),
|
||||
"openbmb/MiniCPM3-4B": PPTestSettings.fast(),
|
||||
# Uses Llama
|
||||
# "mistralai/Mistral-7B-Instruct-v0.1": PPTestSettings.fast(),
|
||||
"state-spaces/mamba-130m-hf": PPTestSettings.fast(),
|
||||
"mistralai/Mixtral-8x7B-Instruct-v0.1": PPTestSettings.fast(load_format="dummy"),
|
||||
"mosaicml/mpt-7b": PPTestSettings.fast(),
|
||||
"nvidia/Minitron-8B-Base": PPTestSettings.fast(),
|
||||
"allenai/OLMo-1B-hf": PPTestSettings.fast(),
|
||||
"allenai/OLMo-2-0425-1B": PPTestSettings.fast(),
|
||||
"allenai/OLMoE-1B-7B-0924-Instruct": PPTestSettings.fast(),
|
||||
"facebook/opt-iml-max-1.3b": PPTestSettings.fast(),
|
||||
"OrionStarAI/Orion-14B-Chat": PPTestSettings.fast(),
|
||||
"microsoft/phi-2": PPTestSettings.fast(),
|
||||
"microsoft/Phi-3-small-8k-instruct": PPTestSettings.fast(),
|
||||
"microsoft/Phi-3.5-MoE-instruct": PPTestSettings.detailed(
|
||||
multi_node_only=True, load_format="dummy"
|
||||
),
|
||||
"Qwen/Qwen2.5-0.5B-Instruct": PPTestSettings.fast(),
|
||||
"Qwen/Qwen1.5-MoE-A2.7B-Chat": PPTestSettings.fast(),
|
||||
"stabilityai/stablelm-3b-4e1t": PPTestSettings.fast(),
|
||||
"bigcode/starcoder2-3b": PPTestSettings.fast(),
|
||||
"upstage/solar-pro-preview-instruct": PPTestSettings.fast(load_format="dummy"),
|
||||
# [Encoder-only]
|
||||
# TODO: Implement PP
|
||||
# "facebook/bart-base": PPTestSettings.fast(),
|
||||
}
|
||||
|
||||
EMBEDDING_MODELS = { # type: ignore[var-annotated]
|
||||
# [Text-only]
|
||||
"intfloat/e5-mistral-7b-instruct": PPTestSettings.fast(runner="pooling"),
|
||||
"BAAI/bge-multilingual-gemma2": PPTestSettings.fast(runner="pooling"),
|
||||
"Qwen/Qwen2.5-Math-RM-72B": PPTestSettings.fast(
|
||||
load_format="dummy", runner="pooling"
|
||||
),
|
||||
}
|
||||
|
||||
MULTIMODAL_MODELS = {
|
||||
# [Decoder-only]
|
||||
"Salesforce/blip2-opt-6.7b": PPTestSettings.fast(),
|
||||
"facebook/chameleon-7b": PPTestSettings.fast(),
|
||||
"zai-org/glm-4v-9b": PPTestSettings.fast(),
|
||||
"OpenGVLab/InternVL3-1B": PPTestSettings.fast(),
|
||||
"llava-hf/llava-1.5-7b-hf": PPTestSettings.fast(),
|
||||
"llava-hf/llava-v1.6-mistral-7b-hf": PPTestSettings.fast(),
|
||||
"llava-hf/LLaVA-NeXT-Video-7B-hf": PPTestSettings.fast(),
|
||||
"llava-hf/llava-onevision-qwen2-0.5b-ov-hf": PPTestSettings.fast(),
|
||||
"openbmb/MiniCPM-Llama3-V-2_5": PPTestSettings.fast(),
|
||||
"allenai/Molmo-7B-D-0924": PPTestSettings.fast(),
|
||||
"AIDC-AI/Ovis2-1B": PPTestSettings.fast(),
|
||||
"AIDC-AI/Ovis2.5-2B": PPTestSettings.fast(),
|
||||
"microsoft/Phi-3.5-vision-instruct": PPTestSettings.fast(),
|
||||
"mistralai/Pixtral-12B-2409": PPTestSettings.fast(load_format="dummy"),
|
||||
"Qwen/Qwen2-Audio-7B-Instruct": PPTestSettings.fast(),
|
||||
"Qwen/Qwen2-VL-2B-Instruct": PPTestSettings.fast(),
|
||||
"fixie-ai/ultravox-v0_5-llama-3_2-1b": PPTestSettings.fast(),
|
||||
}
|
||||
|
||||
# NOTE: You can update this on your local machine to run specific tests
|
||||
TEST_MODELS = [
|
||||
# [LANGUAGE GENERATION]
|
||||
"microsoft/Phi-3.5-MoE-instruct",
|
||||
"meta-llama/Llama-3.2-1B-Instruct",
|
||||
"hmellor/Ilama-3.2-1B",
|
||||
"ibm/PowerLM-3b",
|
||||
"deepseek-ai/DeepSeek-V2-Lite-Chat",
|
||||
# [LANGUAGE EMBEDDING]
|
||||
"intfloat/e5-mistral-7b-instruct",
|
||||
"BAAI/bge-multilingual-gemma2",
|
||||
# [MULTIMODAL GENERATION]
|
||||
"OpenGVLab/InternVL3-1B",
|
||||
"microsoft/Phi-3.5-vision-instruct",
|
||||
"fixie-ai/ultravox-v0_5-llama-3_2-1b",
|
||||
# [LANGUAGE GENERATION - HYBRID ARCH]
|
||||
"ai21labs/Jamba-tiny-dev",
|
||||
]
|
||||
|
||||
|
||||
def _compare_tp(
|
||||
model_id: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: PPTestOptions,
|
||||
num_gpus_available: int,
|
||||
*,
|
||||
method: Literal["generate", "encode"],
|
||||
is_multimodal: bool,
|
||||
):
|
||||
(
|
||||
tp_size,
|
||||
pp_size,
|
||||
eager_mode,
|
||||
) = parallel_setup
|
||||
|
||||
multi_node_only, load_format = test_options
|
||||
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model_id)
|
||||
model_info.check_transformers_version(on_fail="skip")
|
||||
|
||||
trust_remote_code = model_info.trust_remote_code
|
||||
tokenizer_mode = model_info.tokenizer_mode
|
||||
hf_overrides = model_info.hf_overrides
|
||||
hf_config = get_config(model_id, trust_remote_code)
|
||||
require_embed_inputs = model_info.require_embed_inputs
|
||||
max_num_seqs = model_info.max_num_seqs
|
||||
enable_prefix_caching = model_info.enable_prefix_caching
|
||||
|
||||
dtype = "float16"
|
||||
if hf_config.model_type in _FLOAT16_NOT_SUPPORTED_MODELS:
|
||||
dtype = "bfloat16"
|
||||
|
||||
if load_format == "dummy":
|
||||
# Avoid OOM
|
||||
text_overrides = {
|
||||
"num_hidden_layers": 4,
|
||||
"hidden_size": 512,
|
||||
"intermediate_size": 800,
|
||||
"num_attention_heads": 4,
|
||||
"num_key_value_heads": 1,
|
||||
}
|
||||
|
||||
if is_multimodal:
|
||||
hf_overrides.update({"text_config": text_overrides})
|
||||
else:
|
||||
hf_overrides.update(text_overrides)
|
||||
else:
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
if num_gpus_available < tp_size * pp_size:
|
||||
pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
|
||||
if VLLM_MULTI_NODE and distributed_backend == "mp":
|
||||
pytest.skip(
|
||||
"Skipping multi-node pipeline parallel test for "
|
||||
"multiprocessing distributed backend"
|
||||
)
|
||||
if multi_node_only and not VLLM_MULTI_NODE:
|
||||
pytest.skip("Not in multi-node setting")
|
||||
|
||||
common_args = [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
dtype,
|
||||
"--max-model-len",
|
||||
"2048",
|
||||
"--max-num-seqs",
|
||||
"8",
|
||||
]
|
||||
if eager_mode:
|
||||
common_args.append("--enforce-eager")
|
||||
if runner != "auto":
|
||||
common_args.extend(["--runner", runner])
|
||||
if trust_remote_code:
|
||||
common_args.append("--trust-remote-code")
|
||||
if tokenizer_mode:
|
||||
common_args.extend(["--tokenizer-mode", tokenizer_mode])
|
||||
if load_format:
|
||||
common_args.extend(["--load-format", load_format])
|
||||
if hf_overrides:
|
||||
common_args.extend(["--hf-overrides", json.dumps(hf_overrides)])
|
||||
if not enable_prefix_caching:
|
||||
common_args.append("--no-enable-prefix-caching")
|
||||
if require_embed_inputs:
|
||||
common_args.extend(
|
||||
[
|
||||
"--skip-tokenizer-init",
|
||||
"--enable-prompt-embeds",
|
||||
"--enable-mm-embeds",
|
||||
]
|
||||
)
|
||||
if max_num_seqs:
|
||||
common_args.extend(["--max-num-seqs", f"{max_num_seqs}"])
|
||||
|
||||
if distributed_backend == "ray":
|
||||
# Test Ray Compiled Graph for all the tests
|
||||
pp_env = {
|
||||
"VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL": "1",
|
||||
}
|
||||
elif distributed_backend == "mp":
|
||||
pp_env = None
|
||||
else:
|
||||
pp_env = None
|
||||
|
||||
tp_env = None
|
||||
|
||||
pp_args = [
|
||||
*common_args,
|
||||
"--pipeline-parallel-size",
|
||||
str(pp_size),
|
||||
"--tensor-parallel-size",
|
||||
str(tp_size),
|
||||
"--distributed-executor-backend",
|
||||
distributed_backend,
|
||||
]
|
||||
|
||||
# compare without pipeline parallelism
|
||||
# NOTE: use mp backend for TP
|
||||
# PP tests might involve multiple nodes, and ray might
|
||||
# schedule all workers in a node other than the head node,
|
||||
# which can cause the test to fail.
|
||||
tp_args = [
|
||||
*common_args,
|
||||
"--tensor-parallel-size",
|
||||
str(tp_size),
|
||||
"--distributed-executor-backend",
|
||||
"mp",
|
||||
]
|
||||
|
||||
compare_two_settings(
|
||||
model_id,
|
||||
pp_args,
|
||||
tp_args,
|
||||
pp_env,
|
||||
tp_env,
|
||||
method=method,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_id", "parallel_setup", "distributed_backend", "runner", "test_options"),
|
||||
[
|
||||
params
|
||||
for model_id, settings in TEXT_GENERATION_MODELS.items()
|
||||
for params in settings.iter_params(model_id)
|
||||
if model_id in TEST_MODELS
|
||||
],
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_tp_language_generation(
|
||||
model_id: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: PPTestOptions,
|
||||
num_gpus_available,
|
||||
):
|
||||
_compare_tp(
|
||||
model_id,
|
||||
parallel_setup,
|
||||
distributed_backend,
|
||||
runner,
|
||||
test_options,
|
||||
num_gpus_available,
|
||||
method="generate",
|
||||
is_multimodal=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_id", "parallel_setup", "distributed_backend", "runner", "test_options"),
|
||||
[
|
||||
params
|
||||
for model_id, settings in EMBEDDING_MODELS.items()
|
||||
for params in settings.iter_params(model_id)
|
||||
if model_id in TEST_MODELS
|
||||
],
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_tp_language_embedding(
|
||||
model_id: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: PPTestOptions,
|
||||
num_gpus_available,
|
||||
):
|
||||
_compare_tp(
|
||||
model_id,
|
||||
parallel_setup,
|
||||
distributed_backend,
|
||||
runner,
|
||||
test_options,
|
||||
num_gpus_available,
|
||||
method="encode",
|
||||
is_multimodal=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model_id", "parallel_setup", "distributed_backend", "runner", "test_options"),
|
||||
[
|
||||
params
|
||||
for model_id, settings in MULTIMODAL_MODELS.items()
|
||||
for params in settings.iter_params(model_id)
|
||||
if model_id in TEST_MODELS
|
||||
],
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_tp_multimodal_generation(
|
||||
model_id: str,
|
||||
parallel_setup: ParallelSetup,
|
||||
distributed_backend: str,
|
||||
runner: RunnerOption,
|
||||
test_options: PPTestOptions,
|
||||
num_gpus_available,
|
||||
):
|
||||
_compare_tp(
|
||||
model_id,
|
||||
parallel_setup,
|
||||
distributed_backend,
|
||||
runner,
|
||||
test_options,
|
||||
num_gpus_available,
|
||||
method="generate",
|
||||
is_multimodal=True,
|
||||
)
|
||||
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.distributed.utils import get_pp_indices
|
||||
|
||||
|
||||
def test_custom_layer_partition(monkeypatch: pytest.MonkeyPatch):
|
||||
with monkeypatch.context() as m:
|
||||
|
||||
def _verify(partition_str, num_layers, pp_size, goldens):
|
||||
bak = os.environ.get("VLLM_PP_LAYER_PARTITION", None)
|
||||
m.setenv("VLLM_PP_LAYER_PARTITION", partition_str)
|
||||
for pp_rank, golden in enumerate(goldens):
|
||||
assert get_pp_indices(num_layers, pp_rank, pp_size) == golden
|
||||
if bak is not None:
|
||||
m.setenv("VLLM_PP_LAYER_PARTITION", bak)
|
||||
|
||||
# Even partition
|
||||
_verify("5,5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
|
||||
# Balanced partition
|
||||
_verify("4,6,6,4", 20, 4, [(0, 4), (4, 10), (10, 16), (16, 20)])
|
||||
# Put reminder somewhere
|
||||
_verify("5,6,5,6", 22, 4, [(0, 5), (5, 11), (11, 16), (16, 22)])
|
||||
# Invalid partition strings
|
||||
with pytest.raises(ValueError):
|
||||
_verify("5,5,5,5,", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
|
||||
with pytest.raises(ValueError):
|
||||
_verify("5,5,5,a", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
|
||||
# Wrong number of partitions
|
||||
with pytest.raises(ValueError):
|
||||
_verify("5,5,5", 20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
|
||||
# Wrong number of layers
|
||||
with pytest.raises(ValueError):
|
||||
_verify("5,5,5,5", 21, 4, [(0, 5), (5, 10), (10, 15), (15, 20)])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_hidden_layers,pp_size,pp_rank,indices",
|
||||
[
|
||||
# pp_size 2
|
||||
(2, 2, 0, (0, 1)),
|
||||
(2, 2, 1, (1, 2)),
|
||||
(3, 2, 0, (0, 2)),
|
||||
(3, 2, 1, (2, 3)),
|
||||
# pp_size 3
|
||||
(3, 3, 0, (0, 1)),
|
||||
(3, 3, 1, (1, 2)),
|
||||
(3, 3, 2, (2, 3)),
|
||||
(4, 3, 0, (0, 1)),
|
||||
(4, 3, 1, (1, 3)),
|
||||
(4, 3, 2, (3, 4)),
|
||||
(5, 3, 0, (0, 2)),
|
||||
(5, 3, 1, (2, 4)),
|
||||
(5, 3, 2, (4, 5)),
|
||||
],
|
||||
)
|
||||
def test_uneven_auto_partition(
|
||||
num_hidden_layers: int,
|
||||
pp_size: int,
|
||||
pp_rank: int,
|
||||
indices: tuple[int, int],
|
||||
):
|
||||
assert indices == get_pp_indices(num_hidden_layers, pp_rank, pp_size)
|
||||
@@ -0,0 +1,41 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import compare_two_settings, create_new_process_for_each_test
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"PP_SIZE, MODEL_NAME",
|
||||
[
|
||||
(2, "JackFram/llama-160m"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"ATTN_BACKEND",
|
||||
[None] if current_platform.is_rocm() else ["FLASH_ATTN"],
|
||||
)
|
||||
@create_new_process_for_each_test()
|
||||
def test_pp_cudagraph(
|
||||
PP_SIZE: int,
|
||||
MODEL_NAME: str,
|
||||
ATTN_BACKEND: str | None,
|
||||
):
|
||||
cudagraph_args = [
|
||||
# use half precision for speed and memory savings in CI environment
|
||||
"--dtype",
|
||||
"float16",
|
||||
"--pipeline-parallel-size",
|
||||
str(PP_SIZE),
|
||||
"--distributed-executor-backend",
|
||||
"mp",
|
||||
]
|
||||
# On ROCm, defer to the platform attention selector instead of forcing a backend.
|
||||
if ATTN_BACKEND is not None:
|
||||
cudagraph_args.append(f"--attention-backend={ATTN_BACKEND}")
|
||||
|
||||
eager_args = cudagraph_args + ["--enforce-eager"]
|
||||
|
||||
compare_two_settings(MODEL_NAME, eager_args, cudagraph_args)
|
||||
@@ -0,0 +1,474 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
|
||||
import multiprocess as mp
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed
|
||||
|
||||
import vllm.envs as envs
|
||||
from tests.utils import ensure_current_vllm_config
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
|
||||
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
|
||||
from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary
|
||||
from vllm.distributed.parallel_state import (
|
||||
ensure_model_parallel_initialized,
|
||||
get_tp_group,
|
||||
get_world_group,
|
||||
graph_capture,
|
||||
init_distributed_environment,
|
||||
)
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
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"] = "12345"
|
||||
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):
|
||||
# `multiprocessing.Process` cannot accept environment variables directly
|
||||
# so we need to pass the environment variables as arguments
|
||||
# and update the environment variables in the function
|
||||
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
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn():
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
tensor = torch.ones(16, 1024, 1024, dtype=torch.float32).cuda(pynccl_comm.rank)
|
||||
tensor = pynccl_comm.all_reduce(tensor)
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(tensor == pynccl_comm.world_size).cpu().item()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl():
|
||||
distributed_run(worker_fn, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def multiple_allreduce_worker_fn():
|
||||
device = torch.device(f"cuda:{torch.distributed.get_rank()}")
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
# Eager-init path: parent PG has bound_device_id + a CPU backend,
|
||||
# so split_group is supported.
|
||||
group = torch.distributed.split_group(
|
||||
split_ranks=[[0, 1], [2, 3]], backend="cpu:gloo,cuda:nccl"
|
||||
)
|
||||
else:
|
||||
groups = [
|
||||
torch.distributed.new_group(ranks=[0, 1], backend="gloo"),
|
||||
torch.distributed.new_group(ranks=[2, 3], backend="gloo"),
|
||||
]
|
||||
group = groups[0] if torch.distributed.get_rank() in [0, 1] else groups[1]
|
||||
pynccl_comm = PyNcclCommunicator(group=group, device=device)
|
||||
tensor = torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
|
||||
# two groups can communicate independently
|
||||
if torch.distributed.get_rank() in [0, 1]:
|
||||
tensor = pynccl_comm.all_reduce(tensor)
|
||||
tensor = pynccl_comm.all_reduce(tensor)
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(tensor == 4).cpu().item()
|
||||
else:
|
||||
tensor = pynccl_comm.all_reduce(tensor)
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(tensor == 2).cpu().item()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_multiple_allreduce():
|
||||
# this tests pynccl for multiple tp groups, in a standalone way
|
||||
# i.e. call `pynccl_comm.all_reduce` directly
|
||||
distributed_run(multiple_allreduce_worker_fn, 4)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def multiple_allreduce_with_vllm_worker_fn():
|
||||
device = torch.device(f"cuda:{torch.distributed.get_rank()}")
|
||||
with ensure_current_vllm_config():
|
||||
ensure_model_parallel_initialized(2, 2)
|
||||
tensor = torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
|
||||
with graph_capture(device=device):
|
||||
# two tp groups can communicate independently
|
||||
if torch.distributed.get_rank() in [0, 1]:
|
||||
tensor = tensor_model_parallel_all_reduce(tensor)
|
||||
tensor = tensor_model_parallel_all_reduce(tensor)
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(tensor == 4).cpu().item()
|
||||
else:
|
||||
tensor = tensor_model_parallel_all_reduce(tensor)
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(tensor == 2).cpu().item()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_multiple_allreduce_with_vllm():
|
||||
# this tests pynccl for multiple tp groups, together with vllm
|
||||
# i.e. call `tensor_model_parallel_all_reduce`
|
||||
distributed_run(multiple_allreduce_with_vllm_worker_fn, 4)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn_with_cudagraph():
|
||||
with torch.no_grad():
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
# run something in the default stream to initialize torch engine
|
||||
a = torch.ones((4, 4), device=f"cuda:{pynccl_comm.rank}")
|
||||
torch.accelerator.synchronize()
|
||||
with torch.cuda.graph(graph):
|
||||
a_out = pynccl_comm.all_reduce(a)
|
||||
torch.accelerator.synchronize()
|
||||
graph.replay()
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(a_out == pynccl_comm.world_size).cpu().item()
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def all_gather_worker_fn():
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
|
||||
rank = pynccl_comm.rank
|
||||
world_size = pynccl_comm.world_size
|
||||
device = f"cuda:{pynccl_comm.rank}"
|
||||
|
||||
num_elems = 1000
|
||||
tensor = (
|
||||
torch.arange(num_elems, dtype=torch.float32, device=device) + rank * num_elems
|
||||
)
|
||||
result = torch.zeros(num_elems * world_size, dtype=torch.float32, device=device)
|
||||
|
||||
expected = torch.cat(
|
||||
[
|
||||
torch.arange(num_elems, dtype=torch.float32) + r * num_elems
|
||||
for r in range(world_size)
|
||||
]
|
||||
).to(device)
|
||||
|
||||
pynccl_comm.all_gather(result, tensor)
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_all_gather():
|
||||
distributed_run(all_gather_worker_fn, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def cuda_communicator_all_gather_dim_worker_fn():
|
||||
with ensure_current_vllm_config():
|
||||
ensure_model_parallel_initialized(2, 1)
|
||||
|
||||
tp_group = get_tp_group()
|
||||
comm = tp_group.device_communicator
|
||||
assert comm is not None
|
||||
|
||||
rank = tp_group.rank_in_group
|
||||
world_size = tp_group.world_size
|
||||
device = tp_group.device
|
||||
|
||||
shape = (2, 3, 4)
|
||||
num_elems = 1
|
||||
for size in shape:
|
||||
num_elems *= size
|
||||
|
||||
for dim in (1, -1):
|
||||
tensor = (
|
||||
torch.arange(num_elems, dtype=torch.float32, device=device).reshape(shape)
|
||||
+ rank * num_elems
|
||||
)
|
||||
expected = torch.cat(
|
||||
[
|
||||
torch.arange(num_elems, dtype=torch.float32, device=device).reshape(
|
||||
shape
|
||||
)
|
||||
+ r * num_elems
|
||||
for r in range(world_size)
|
||||
],
|
||||
dim=dim,
|
||||
)
|
||||
|
||||
result = comm.all_gather(tensor, dim=dim)
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_cuda_communicator_all_gather_dim_not_zero():
|
||||
distributed_run(cuda_communicator_all_gather_dim_worker_fn, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def all_gatherv_worker_fn():
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
|
||||
rank = pynccl_comm.rank
|
||||
world_size = pynccl_comm.world_size
|
||||
device = f"cuda:{pynccl_comm.rank}"
|
||||
|
||||
assert world_size <= 8
|
||||
sizes = [81, 20, 57, 52, 81, 5, 49, 49][:world_size]
|
||||
num_elems = sizes[rank]
|
||||
tensor = torch.arange(num_elems, dtype=torch.float32, device=device) + rank * 100
|
||||
result = torch.zeros(sum(sizes), dtype=torch.float32, device=device)
|
||||
|
||||
expected = torch.cat(
|
||||
[
|
||||
torch.arange(sizes[r], dtype=torch.float32) + r * 100
|
||||
for r in range(world_size)
|
||||
]
|
||||
).to(device)
|
||||
|
||||
pynccl_comm.all_gatherv(result, tensor, sizes=sizes)
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_all_gatherv():
|
||||
distributed_run(all_gatherv_worker_fn, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def reduce_scatter_worker_fn():
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
|
||||
rank = pynccl_comm.rank
|
||||
world_size = pynccl_comm.world_size
|
||||
device = f"cuda:{pynccl_comm.rank}"
|
||||
|
||||
num_elems = 1000
|
||||
tensor = (
|
||||
torch.arange(num_elems, dtype=torch.float32, device=device) + rank * num_elems
|
||||
)
|
||||
assert num_elems % world_size == 0
|
||||
result = torch.zeros(num_elems // world_size, dtype=torch.float32, device=device)
|
||||
|
||||
# Calculate expected result for this rank's chunk
|
||||
scattered_size = num_elems // world_size
|
||||
all_tensors = [
|
||||
torch.arange(num_elems, dtype=torch.float32) + r * num_elems
|
||||
for r in range(world_size)
|
||||
]
|
||||
expected = sum(
|
||||
tensor[rank * scattered_size : (rank + 1) * scattered_size]
|
||||
for tensor in all_tensors
|
||||
).to(device)
|
||||
|
||||
pynccl_comm.reduce_scatter(result, tensor)
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_reduce_scatter():
|
||||
distributed_run(reduce_scatter_worker_fn, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def reduce_scatterv_worker_fn():
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
|
||||
rank = pynccl_comm.rank
|
||||
world_size = pynccl_comm.world_size
|
||||
device = f"cuda:{pynccl_comm.rank}"
|
||||
|
||||
assert world_size <= 8
|
||||
sizes = [81, 20, 57, 52, 81, 5, 49, 49][:world_size]
|
||||
num_elems = sum(sizes)
|
||||
tensor = torch.arange(num_elems, dtype=torch.float32, device=device) + rank * 100
|
||||
result = torch.zeros(sizes[rank], dtype=torch.float32, device=device)
|
||||
|
||||
# Calculate expected result for this rank's chunk
|
||||
all_tensors = [
|
||||
torch.arange(num_elems, dtype=torch.float32) + r * 100
|
||||
for r in range(world_size)
|
||||
]
|
||||
sizes_cumsum = np.cumsum(sizes)
|
||||
start = 0 if rank == 0 else sizes_cumsum[rank - 1]
|
||||
end = sizes_cumsum[rank]
|
||||
expected = sum(tensor[start:end] for tensor in all_tensors).to(device)
|
||||
|
||||
pynccl_comm.reduce_scatterv(result, tensor, sizes=sizes)
|
||||
torch.accelerator.synchronize()
|
||||
torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-8)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_reduce_scatterv():
|
||||
distributed_run(reduce_scatterv_worker_fn, 2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_with_cudagraph():
|
||||
distributed_run(worker_fn_with_cudagraph, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def send_recv_worker_fn():
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
if pynccl_comm.rank == 0:
|
||||
tensor = torch.ones(16, 1024, 1024, dtype=torch.float32).cuda(pynccl_comm.rank)
|
||||
else:
|
||||
tensor = torch.empty(16, 1024, 1024, dtype=torch.float32).cuda(pynccl_comm.rank)
|
||||
|
||||
if pynccl_comm.rank == 0:
|
||||
pynccl_comm.send(tensor, dst=(pynccl_comm.rank + 1) % pynccl_comm.world_size)
|
||||
else:
|
||||
pynccl_comm.recv(tensor, src=(pynccl_comm.rank - 1) % pynccl_comm.world_size)
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(tensor == 1).cpu().item()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_send_recv():
|
||||
distributed_run(send_recv_worker_fn, 2)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def multiple_send_recv_worker_fn():
|
||||
device = torch.device(f"cuda:{torch.distributed.get_rank()}")
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
group = torch.distributed.split_group(
|
||||
split_ranks=[[0, 2], [1, 3]], backend="cpu:gloo,cuda:nccl"
|
||||
)
|
||||
else:
|
||||
groups = [
|
||||
torch.distributed.new_group(ranks=[0, 2], backend="gloo"),
|
||||
torch.distributed.new_group(ranks=[1, 3], backend="gloo"),
|
||||
]
|
||||
group = groups[0] if torch.distributed.get_rank() in [0, 2] else groups[1]
|
||||
pynccl_comm = PyNcclCommunicator(group=group, device=device)
|
||||
if torch.distributed.get_rank() == 0:
|
||||
tensor = torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
|
||||
elif torch.distributed.get_rank() == 1:
|
||||
tensor = 2 * torch.ones(16, 1024, 1024, dtype=torch.float32, device=device)
|
||||
else:
|
||||
tensor = torch.empty(16, 1024, 1024, dtype=torch.float32, device=device)
|
||||
if torch.distributed.get_rank() in [0, 1]:
|
||||
pynccl_comm.send(tensor, dst=(pynccl_comm.rank + 1) % pynccl_comm.world_size)
|
||||
else:
|
||||
pynccl_comm.recv(tensor, src=(pynccl_comm.rank - 1) % pynccl_comm.world_size)
|
||||
torch.accelerator.synchronize()
|
||||
if torch.distributed.get_rank() in [0, 2]:
|
||||
assert torch.all(tensor == 1).cpu().item()
|
||||
else:
|
||||
assert torch.all(tensor == 2).cpu().item()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_multiple_send_recv():
|
||||
distributed_run(multiple_send_recv_worker_fn, 4)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs to run the test."
|
||||
)
|
||||
def test_pynccl_broadcast():
|
||||
distributed_run(broadcast_worker_fn, 4)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def broadcast_worker_fn():
|
||||
# Test broadcast for every root rank.
|
||||
# Essentially this is an all-gather operation.
|
||||
pynccl_comm = PyNcclCommunicator(
|
||||
get_world_group().cpu_group, device=get_world_group().device
|
||||
)
|
||||
recv_tensors = [
|
||||
torch.empty(16, 1024, 1024, dtype=torch.float32, device=pynccl_comm.device)
|
||||
for i in range(pynccl_comm.world_size)
|
||||
]
|
||||
recv_tensors[pynccl_comm.rank] = (
|
||||
torch.ones(16, 1024, 1024, dtype=torch.float32, device=pynccl_comm.device)
|
||||
* pynccl_comm.rank
|
||||
)
|
||||
|
||||
for i in range(pynccl_comm.world_size):
|
||||
pynccl_comm.broadcast(recv_tensors[i], src=i)
|
||||
# the broadcast op might be launched in a different stream
|
||||
# need to synchronize to make sure the tensor is ready
|
||||
torch.accelerator.synchronize()
|
||||
assert torch.all(recv_tensors[i] == i).cpu().item()
|
||||
|
||||
|
||||
def test_ncclGetUniqueId():
|
||||
lib = NCCLLibrary()
|
||||
unique_id = lib.ncclGetUniqueId()
|
||||
# `list(unique_id.internal)` is something like this:
|
||||
# [34, -16, 23, 83, 109, -19, 59, 95, 2, 0, -86, 55, 10, -128, 0, 29, 0,
|
||||
# 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
# 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
# 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
# 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
|
||||
# 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
||||
# as long as the function doesn't raise an exception, we're good
|
||||
assert unique_id is not None
|
||||
@@ -0,0 +1,483 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import multiprocessing
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import (
|
||||
KB,
|
||||
MB,
|
||||
QuickAllReduce,
|
||||
QuickReduceRegime,
|
||||
)
|
||||
from vllm.distributed.parallel_state import get_tp_group, graph_capture
|
||||
from vllm.envs import disable_envs_cache
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import (
|
||||
ensure_model_parallel_initialized,
|
||||
init_test_distributed_environment,
|
||||
multi_process_parallel,
|
||||
set_random_seed,
|
||||
)
|
||||
|
||||
|
||||
def on_gfx942() -> bool:
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_gfx942 as rocm_on_gfx942
|
||||
|
||||
return rocm_on_gfx942()
|
||||
return False
|
||||
|
||||
|
||||
set_random_seed(42)
|
||||
_test_size_rng = random.Random(44)
|
||||
# Size over 8MB is sufficient for custom quick allreduce.
|
||||
test_sizes = [
|
||||
_test_size_rng.randint(8 * 1024 * 1024, 10 * 1024 * 1024) for _ in range(8)
|
||||
]
|
||||
for i, v in enumerate(test_sizes):
|
||||
test_sizes[i] -= v % 8
|
||||
|
||||
|
||||
def _assert_quickreduce(fa, inp):
|
||||
assert fa is not None
|
||||
assert not fa.disabled
|
||||
assert fa.should_quick_allreduce(inp)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def envs_cache_disabled():
|
||||
disable_envs_cache()
|
||||
yield
|
||||
disable_envs_cache()
|
||||
|
||||
|
||||
def _make_quick_allreduce_for_test(
|
||||
min_size_mb: int | None = None,
|
||||
quantization_min_size: int | None = None,
|
||||
) -> QuickAllReduce:
|
||||
quick_reduce = QuickAllReduce.__new__(QuickAllReduce)
|
||||
quick_reduce.disabled = False
|
||||
quick_reduce.qr_max_size = 16 * MB
|
||||
quick_reduce.qr_min_size = min_size_mb * MB if min_size_mb is not None else None
|
||||
quick_reduce.qr_quant_level = QuickReduceRegime.INT4
|
||||
quick_reduce.qr_quantization_min_size = quantization_min_size
|
||||
quick_reduce.use_fp16_kernels = False
|
||||
quick_reduce.world_size = 2
|
||||
return quick_reduce
|
||||
|
||||
|
||||
def test_should_quick_allreduce_uses_builtin_min_size_when_unset():
|
||||
quick_reduce = _make_quick_allreduce_for_test(min_size_mb=None)
|
||||
|
||||
below_builtin_min = torch.empty(MB // 4, dtype=torch.float16)
|
||||
at_builtin_min = torch.empty(MB // 2, dtype=torch.float16)
|
||||
|
||||
assert not quick_reduce.should_quick_allreduce(below_builtin_min)
|
||||
assert quick_reduce.should_quick_allreduce(at_builtin_min)
|
||||
|
||||
|
||||
def test_should_quick_allreduce_uses_min_size_override():
|
||||
quick_reduce = _make_quick_allreduce_for_test(min_size_mb=0)
|
||||
|
||||
below_builtin_min = torch.empty(8, dtype=torch.float16)
|
||||
|
||||
assert quick_reduce.should_quick_allreduce(below_builtin_min)
|
||||
|
||||
|
||||
def test_quick_allreduce_min_size_env_unset(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", raising=False)
|
||||
|
||||
assert QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB) is None
|
||||
|
||||
|
||||
def test_quick_allreduce_min_size_env_converts_mb_to_bytes(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "4")
|
||||
|
||||
assert QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB) == 4 * MB
|
||||
|
||||
|
||||
def test_quick_allreduce_min_size_env_rejects_negative(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "-1")
|
||||
|
||||
with pytest.raises(ValueError, match="must be non-negative"):
|
||||
QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB)
|
||||
|
||||
|
||||
def test_quick_allreduce_min_size_env_allows_equal_to_max(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "16")
|
||||
|
||||
assert QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB) == 16 * MB
|
||||
|
||||
|
||||
def test_quick_allreduce_min_size_env_rejects_larger_than_max(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB", "17")
|
||||
|
||||
with pytest.raises(ValueError, match="effective QuickReduce max size"):
|
||||
QuickAllReduce._get_qr_min_size(qr_max_size=16 * MB)
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_env_unset(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", raising=False)
|
||||
|
||||
assert QuickAllReduce._get_qr_quantization_min_size() is None
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_env_converts_kb_to_bytes(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", "2048")
|
||||
|
||||
assert QuickAllReduce._get_qr_quantization_min_size() == 2048 * KB
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_env_rejects_negative(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
envs_cache_disabled,
|
||||
):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB", "-1")
|
||||
|
||||
with pytest.raises(ValueError, match="must be non-negative"):
|
||||
QuickAllReduce._get_qr_quantization_min_size()
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_unset_uses_configured_codec():
|
||||
quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=None)
|
||||
inp = torch.empty(8, dtype=torch.float16)
|
||||
|
||||
assert quick_reduce._get_qr_quant_level(inp) == QuickReduceRegime.INT4.value
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_uses_fp_below_threshold():
|
||||
quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2048)
|
||||
inp = torch.empty(1024 // 2, dtype=torch.float16)
|
||||
|
||||
assert quick_reduce._get_qr_quant_level(inp) == QuickReduceRegime.FP.value
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_uses_configured_codec_at_threshold():
|
||||
quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2048)
|
||||
inp = torch.empty(2048 // 2, dtype=torch.float16)
|
||||
|
||||
assert quick_reduce._get_qr_quant_level(inp) == QuickReduceRegime.INT4.value
|
||||
|
||||
|
||||
def test_quick_allreduce_quantization_min_size_does_not_change_eligibility():
|
||||
quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2 * MB)
|
||||
|
||||
below_builtin_min = torch.empty(MB // 4, dtype=torch.float16)
|
||||
at_builtin_min = torch.empty(MB // 2, dtype=torch.float16)
|
||||
|
||||
assert not quick_reduce.should_quick_allreduce(below_builtin_min)
|
||||
assert quick_reduce.should_quick_allreduce(at_builtin_min)
|
||||
|
||||
|
||||
def test_quick_allreduce_passes_dynamic_quant_level(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
quick_reduce = _make_quick_allreduce_for_test(quantization_min_size=2 * KB)
|
||||
quick_reduce._ptr = object()
|
||||
inp = torch.empty(KB // 2, dtype=torch.float16)
|
||||
called_quant_level = None
|
||||
|
||||
def fake_qr_all_reduce(
|
||||
fa,
|
||||
inp,
|
||||
out,
|
||||
quant_level,
|
||||
cast_bf2half,
|
||||
):
|
||||
nonlocal called_quant_level
|
||||
called_quant_level = quant_level
|
||||
|
||||
monkeypatch.setattr(ops, "qr_all_reduce", fake_qr_all_reduce)
|
||||
|
||||
quick_reduce.quick_all_reduce(inp)
|
||||
|
||||
assert called_quant_level == QuickReduceRegime.FP.value
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_quickreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
ensure_model_parallel_initialized(tp_size, pp_size)
|
||||
group = get_tp_group().device_group
|
||||
fa = get_tp_group().device_communicator.qr_comm
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
torch.distributed.all_reduce(data, group=group)
|
||||
torch.accelerator.synchronize()
|
||||
del data
|
||||
|
||||
# we use the first group to communicate once
|
||||
# and the second group to communicate twice
|
||||
# and so on
|
||||
# this is used to demonstrate that each group can
|
||||
# communicate independently
|
||||
num_communication = rank // tp_size + 1
|
||||
|
||||
for sz in test_sizes:
|
||||
for dtype in [torch.float16, torch.bfloat16]:
|
||||
with graph_capture(device=device) as graph_capture_context:
|
||||
device_idx = torch.accelerator.current_device_index()
|
||||
inp1 = torch.randint(1, 23, (sz,), dtype=dtype, device=device_idx)
|
||||
inp2 = torch.randint(-23, 1, (sz,), dtype=dtype, device=device_idx)
|
||||
_assert_quickreduce(fa, inp1)
|
||||
_assert_quickreduce(fa, inp2)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=graph_capture_context.stream):
|
||||
for _ in range(num_communication):
|
||||
out1 = tensor_model_parallel_all_reduce(inp1)
|
||||
dist.all_reduce(inp1, group=group)
|
||||
out2 = tensor_model_parallel_all_reduce(inp2)
|
||||
dist.all_reduce(inp2, group=group)
|
||||
graph.replay()
|
||||
torch.testing.assert_close(out1, inp1, atol=2.5, rtol=0.1)
|
||||
torch.testing.assert_close(out2, inp2, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_quickreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
|
||||
# Size over 8MB is sufficient for custom quick allreduce.
|
||||
sz = 16 * 1024 * 1024
|
||||
fa = get_tp_group().device_communicator.qr_comm
|
||||
inp = torch.tensor(
|
||||
[1.0 * ((i) % 23) for i in range(sz)], dtype=torch.float16, device=device
|
||||
)
|
||||
_assert_quickreduce(fa, inp)
|
||||
out = fa.quick_all_reduce(inp)
|
||||
torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
|
||||
|
||||
inp = torch.tensor(
|
||||
[1.0 * ((i) % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
|
||||
)
|
||||
_assert_quickreduce(fa, inp)
|
||||
out = fa.quick_all_reduce(inp)
|
||||
torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def bf16_cast_quickreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
):
|
||||
with monkeypatch.context() as m:
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("HIP_VISIBLE_DEVICES", raising=False)
|
||||
m.delenv("ROCR_VISIBLE_DEVICES", raising=False)
|
||||
m.setenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1")
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
|
||||
sz = 16 * 1024 * 1024
|
||||
fa = get_tp_group().device_communicator.qr_comm
|
||||
inp = torch.tensor(
|
||||
[1.0 * (i % 23) for i in range(sz)], dtype=torch.bfloat16, device=device
|
||||
)
|
||||
_assert_quickreduce(fa, inp)
|
||||
assert fa.use_fp16_kernels
|
||||
out = fa.quick_all_reduce(inp)
|
||||
torch.testing.assert_close(out, inp * tp_size, atol=2.5, rtol=0.1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
|
||||
)
|
||||
@pytest.mark.parametrize("quant_mode", ["FP", "INT8", "INT6", "INT4", "INT3"])
|
||||
@pytest.mark.parametrize("tp_size", [2])
|
||||
@pytest.mark.parametrize("pipeline_parallel_size", [1, 2])
|
||||
@pytest.mark.parametrize("test_target", [graph_quickreduce, eager_quickreduce])
|
||||
def test_custom_quick_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pipeline_parallel_size,
|
||||
test_target,
|
||||
quant_mode,
|
||||
):
|
||||
world_size = tp_size * pipeline_parallel_size
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
if test_target is graph_quickreduce and on_gfx942():
|
||||
pytest.xfail(
|
||||
"CUDA graph capture with quick reduce hits "
|
||||
"hipErrorStreamCaptureInvalidated on gfx942"
|
||||
)
|
||||
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
|
||||
|
||||
multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
|
||||
)
|
||||
def test_custom_quick_allreduce_bf16_cast(monkeypatch: pytest.MonkeyPatch):
|
||||
if torch.accelerator.device_count() < 2:
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "FP")
|
||||
multi_process_parallel(monkeypatch, 2, 1, bf16_cast_quickreduce)
|
||||
|
||||
|
||||
def qr_variable_input(rank, world_size):
|
||||
"""
|
||||
When the tensor parallelism is set to 4 or 8, frequent changes
|
||||
in the input shape can cause QuickReduce to hang (this issue
|
||||
has been observed with the gpt_oss model).
|
||||
"""
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
qr_max_size = None # MB
|
||||
_ptr = ops.init_custom_qr(rank, world_size, qr_max_size)
|
||||
ranks = []
|
||||
for i in range(world_size):
|
||||
ranks.append(i)
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
dist.init_process_group(
|
||||
backend="cpu:gloo,cuda:nccl",
|
||||
init_method="tcp://127.0.0.1:29500",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
device_id=device,
|
||||
)
|
||||
else:
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method="tcp://127.0.0.1:29500",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
)
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
cpu_group = torch.distributed.split_group(
|
||||
split_ranks=[ranks], backend="cpu:gloo,cuda:nccl"
|
||||
)
|
||||
else:
|
||||
cpu_group = torch.distributed.new_group(ranks, backend="nccl")
|
||||
|
||||
handle = ops.qr_get_handle(_ptr)
|
||||
world_size = dist.get_world_size(group=cpu_group)
|
||||
handles = [None] * world_size
|
||||
dist.all_gather_object(handles, handle, group=cpu_group)
|
||||
ops.qr_open_handles(_ptr, handles)
|
||||
|
||||
num = 1
|
||||
s1 = 1024
|
||||
while num < 50000: # 50000 is sufficient to identify issues.
|
||||
dtype = torch.float16
|
||||
device_idx = torch.accelerator.current_device_index()
|
||||
if num % 2 == 0:
|
||||
s2 = 1024
|
||||
inp1 = torch.zeros((s1, s2), dtype=dtype, device=device_idx)
|
||||
else:
|
||||
s2 = 2048
|
||||
inp1 = torch.ones((s1, s2), dtype=dtype, device=device_idx)
|
||||
result = torch.empty_like(inp1)
|
||||
# FP = 0 INT8 = 1 INT6 = 2 INT4 = 3 INT3 = 4
|
||||
ops.qr_all_reduce(_ptr, inp1, result, 3, cast_bf2half=True)
|
||||
try:
|
||||
if inp1[0, 0] == 0:
|
||||
assert torch.all(result == 0)
|
||||
else:
|
||||
assert torch.all(result == world_size)
|
||||
except AssertionError:
|
||||
print("Assertion failed! Allreduce results are incorrect.")
|
||||
raise
|
||||
num += 1
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_rocm(), reason="only test quick allreduce for rocm"
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [4, 8])
|
||||
@pytest.mark.parametrize("pipeline_parallel_size", [1])
|
||||
def test_custom_quick_allreduce_variable_input(tp_size, pipeline_parallel_size):
|
||||
world_size = tp_size * pipeline_parallel_size
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
|
||||
multiprocessing.set_start_method("spawn", force=True)
|
||||
# 60s is enough
|
||||
timeout = 60
|
||||
processes = []
|
||||
for rank in range(tp_size):
|
||||
p = multiprocessing.Process(target=qr_variable_input, args=(rank, tp_size))
|
||||
p.start()
|
||||
processes.append((rank, p))
|
||||
for rank, p in processes:
|
||||
p.join(timeout=timeout)
|
||||
if p.is_alive():
|
||||
for r, proc in processes:
|
||||
if proc.is_alive():
|
||||
proc.terminate()
|
||||
proc.join()
|
||||
raise RuntimeError(f"QuickReduce hang detected after {timeout} seconds!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_custom_quick_allreduce_variable_input(tp_size=4, pipeline_parallel_size=1)
|
||||
@@ -0,0 +1,389 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Integration tests for RayExecutorV2 at the executor level.
|
||||
Validates executor initialization, placement group support, RPC calls,
|
||||
and distributed execution with various TP/PP configurations.
|
||||
"""
|
||||
|
||||
import gc
|
||||
import threading
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
|
||||
from vllm import LLM
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.v1.executor import ray_executor_v2
|
||||
from vllm.v1.executor.ray_executor_v2 import RayExecutorV2
|
||||
|
||||
pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend")
|
||||
|
||||
MODEL = "facebook/opt-125m"
|
||||
|
||||
|
||||
def create_vllm_config(
|
||||
tensor_parallel_size: int = 1,
|
||||
pipeline_parallel_size: int = 1,
|
||||
max_model_len: int = 256,
|
||||
gpu_memory_utilization: float = 0.3,
|
||||
placement_group=None,
|
||||
) -> VllmConfig:
|
||||
engine_args = EngineArgs(
|
||||
model=MODEL,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
pipeline_parallel_size=pipeline_parallel_size,
|
||||
max_model_len=max_model_len,
|
||||
gpu_memory_utilization=gpu_memory_utilization,
|
||||
distributed_executor_backend="ray",
|
||||
enforce_eager=True,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
|
||||
if placement_group is not None:
|
||||
vllm_config.parallel_config.placement_group = placement_group
|
||||
|
||||
return vllm_config
|
||||
|
||||
|
||||
def ensure_ray_initialized():
|
||||
if not ray.is_initialized():
|
||||
ray.init(ignore_reinit_error=True)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def create_placement_group(request):
|
||||
ensure_ray_initialized()
|
||||
num_gpus = request.param
|
||||
bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]
|
||||
pg = ray.util.placement_group(bundles, strategy="PACK")
|
||||
ray.get(pg.ready())
|
||||
yield pg
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor(request):
|
||||
"""Create a RayExecutorV2 and shut it down after the test."""
|
||||
executor = RayExecutorV2(vllm_config=request.param)
|
||||
yield executor
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
def assert_executor(executor, tp_size, pp_size):
|
||||
"""Common assertions for executor initialization tests."""
|
||||
world_size = tp_size * pp_size
|
||||
expected_output_rank = (pp_size - 1) * tp_size
|
||||
|
||||
assert executor.world_size == world_size
|
||||
assert len(executor.ray_worker_handles) == world_size
|
||||
assert len(executor.response_mqs) == world_size
|
||||
assert executor._get_output_rank() == expected_output_rank
|
||||
|
||||
if pp_size > 1:
|
||||
expected_concurrent_batches = pp_size + int(
|
||||
executor.vllm_config.scheduler_config.async_scheduling
|
||||
and executor.vllm_config.use_v2_model_runner
|
||||
)
|
||||
assert (
|
||||
executor.vllm_config.max_concurrent_batches == expected_concurrent_batches
|
||||
)
|
||||
|
||||
executor.check_health()
|
||||
assert not executor.is_failed
|
||||
|
||||
ranks = sorted(h.rank for h in executor.ray_worker_handles)
|
||||
assert ranks == list(range(world_size))
|
||||
|
||||
for handle in executor.ray_worker_handles:
|
||||
assert handle.node_id is not None
|
||||
|
||||
|
||||
def test_select_tcpstore_port_seeds_disjoint_windows(monkeypatch):
|
||||
"""Co-located DP engines scan distinct, adjacent port windows, so two
|
||||
engines on a node cannot pick the same TCPStore port."""
|
||||
requested = []
|
||||
|
||||
def fake_get_open_port(start_port, max_attempts):
|
||||
requested.append((start_port, max_attempts))
|
||||
return start_port
|
||||
|
||||
monkeypatch.setattr(ray_executor_v2, "_get_open_port", fake_get_open_port)
|
||||
|
||||
ports = [
|
||||
RayExecutorV2._select_tcpstore_port(rank, master_port=29500)
|
||||
for rank in range(4)
|
||||
]
|
||||
|
||||
assert requested == [(29600, 32), (29632, 32), (29664, 32), (29696, 32)]
|
||||
assert len(set(ports)) == 4
|
||||
|
||||
|
||||
def test_select_tcpstore_port_non_dp_uses_random(monkeypatch):
|
||||
"""A non-DP engine has no local rank and uses a random port."""
|
||||
monkeypatch.setattr(ray_executor_v2, "get_open_port", lambda: 54321)
|
||||
assert RayExecutorV2._select_tcpstore_port(None, master_port=29500) == 54321
|
||||
|
||||
|
||||
def test_select_tcpstore_port_full_window_uses_random(monkeypatch):
|
||||
"""A fully occupied window falls back to a random port."""
|
||||
|
||||
def raise_full(start_port, max_attempts):
|
||||
raise RuntimeError("no open port")
|
||||
|
||||
monkeypatch.setattr(ray_executor_v2, "_get_open_port", raise_full)
|
||||
monkeypatch.setattr(ray_executor_v2, "get_open_port", lambda: 54321)
|
||||
assert RayExecutorV2._select_tcpstore_port(0, master_port=29500) == 54321
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp_size, pp_size", [(1, 1), (2, 1), (4, 1), (2, 2)])
|
||||
def test_ray_v2_executor(tp_size, pp_size):
|
||||
"""Validate RayExecutorV2 with various TP/PP configs."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=tp_size,
|
||||
pipeline_parallel_size=pp_size,
|
||||
)
|
||||
executor = RayExecutorV2(vllm_config=vllm_config)
|
||||
try:
|
||||
assert_executor(executor, tp_size, pp_size)
|
||||
finally:
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tp_size, pp_size, create_placement_group",
|
||||
[(2, 1, 2), (4, 1, 4), (2, 2, 4)],
|
||||
indirect=["create_placement_group"],
|
||||
)
|
||||
def test_ray_v2_executor_pg(tp_size, pp_size, create_placement_group):
|
||||
"""Validate RayExecutorV2 with various TP/PP configs using external PG."""
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=tp_size,
|
||||
pipeline_parallel_size=pp_size,
|
||||
placement_group=create_placement_group,
|
||||
)
|
||||
executor = RayExecutorV2(vllm_config=vllm_config)
|
||||
try:
|
||||
assert_executor(executor, tp_size, pp_size)
|
||||
finally:
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"executor",
|
||||
[create_vllm_config(tensor_parallel_size=2)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_ray_v2_executor_failure_callback(executor):
|
||||
"""Validate failure callback registration."""
|
||||
callback_invoked = False
|
||||
|
||||
def test_callback():
|
||||
nonlocal callback_invoked
|
||||
callback_invoked = True
|
||||
|
||||
executor.register_failure_callback(test_callback)
|
||||
assert not callback_invoked
|
||||
|
||||
executor.is_failed = True
|
||||
executor.register_failure_callback(test_callback)
|
||||
assert callback_invoked
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"executor",
|
||||
[create_vllm_config(tensor_parallel_size=2)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_ray_v2_executor_collective_rpc(executor):
|
||||
"""Validate collective RPC calls through MessageQueue."""
|
||||
executor.check_health()
|
||||
assert not executor.is_failed
|
||||
assert executor.rpc_broadcast_mq is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"executor",
|
||||
[create_vllm_config(tensor_parallel_size=2)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_ray_v2_executor_driver_node_rank_0(executor):
|
||||
"""Validate that driver node workers get the lowest ranks."""
|
||||
driver_node = ray.get_runtime_context().get_node_id()
|
||||
|
||||
for handle in executor.ray_worker_handles:
|
||||
assert handle.node_id == driver_node
|
||||
|
||||
rank0_handle = next(h for h in executor.ray_worker_handles if h.rank == 0)
|
||||
assert rank0_handle.node_id == driver_node
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"executor",
|
||||
[create_vllm_config(tensor_parallel_size=2)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_ray_v2_executor_worker_death(executor):
|
||||
"""Validate executor detects worker death via ray.wait()."""
|
||||
callback_event = threading.Event()
|
||||
|
||||
def on_failure():
|
||||
callback_event.set()
|
||||
|
||||
executor.register_failure_callback(on_failure)
|
||||
assert not executor.is_failed
|
||||
|
||||
# Kill one worker actor externally
|
||||
victim = executor.ray_worker_handles[1].actor
|
||||
ray.kill(victim, no_restart=True)
|
||||
|
||||
# Monitor thread should detect the death and invoke callback
|
||||
assert callback_event.wait(timeout=30)
|
||||
assert executor.is_failed
|
||||
assert executor.shutting_down
|
||||
|
||||
|
||||
def test_ray_v2_executor_shutdown():
|
||||
"""Validate graceful shutdown: ray.kill() terminates all worker actors."""
|
||||
executor = RayExecutorV2(vllm_config=create_vllm_config(tensor_parallel_size=2))
|
||||
assert executor.rpc_broadcast_mq is not None
|
||||
assert len(executor.response_mqs) == executor.world_size
|
||||
|
||||
actors = [h.actor for h in executor.ray_worker_handles]
|
||||
executor.shutdown()
|
||||
|
||||
for actor in actors:
|
||||
with pytest.raises(ray.exceptions.RayActorError):
|
||||
ray.get(actor.wait_for_init.remote(), timeout=5)
|
||||
|
||||
assert executor.rpc_broadcast_mq is None
|
||||
assert len(executor.response_mqs) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"executor",
|
||||
[create_vllm_config(tensor_parallel_size=2)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_ray_v2_run_refs_stored_for_monitoring(executor):
|
||||
"""Validate worker handles store run_ref for monitoring."""
|
||||
for handle in executor.ray_worker_handles:
|
||||
assert handle.run_ref is not None
|
||||
ready, _ = ray.wait([handle.run_ref], timeout=0)
|
||||
assert len(ready) == 0, "run_ref should be pending"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)])
|
||||
def test_ray_v2_single_node_generation(tp_size, pp_size):
|
||||
"""End-to-end LLM generation with RayExecutorV2."""
|
||||
|
||||
llm = LLM(
|
||||
model=MODEL,
|
||||
tensor_parallel_size=tp_size,
|
||||
pipeline_parallel_size=pp_size,
|
||||
distributed_executor_backend="ray",
|
||||
enforce_eager=True,
|
||||
max_model_len=256,
|
||||
gpu_memory_utilization=0.3,
|
||||
)
|
||||
try:
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
outputs = llm.generate(prompts)
|
||||
|
||||
assert len(outputs) == len(prompts)
|
||||
for output in outputs:
|
||||
assert len(output.outputs) > 0
|
||||
assert len(output.outputs[0].text) > 0
|
||||
finally:
|
||||
llm.llm_engine.model_executor.shutdown()
|
||||
del llm
|
||||
gc.collect()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bundle_indices, expected_bundle_ids, create_placement_group",
|
||||
[("2,3", [2, 3], 4), ("3,2", [3, 2], 4)],
|
||||
indirect=["create_placement_group"],
|
||||
)
|
||||
def test_ray_v2_bundle_indices_env(
|
||||
bundle_indices, expected_bundle_ids, create_placement_group, monkeypatch
|
||||
):
|
||||
"""Validate explicit VLLM_RAY_BUNDLE_INDICES bundle placement."""
|
||||
monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", bundle_indices)
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2,
|
||||
placement_group=create_placement_group,
|
||||
)
|
||||
executor = RayExecutorV2(vllm_config=vllm_config)
|
||||
try:
|
||||
actual = [
|
||||
h.bundle_id_idx
|
||||
for h in sorted(executor.ray_worker_handles, key=lambda h: h.rank)
|
||||
]
|
||||
assert actual == expected_bundle_ids
|
||||
assert_executor(executor, tp_size=2, pp_size=1)
|
||||
finally:
|
||||
executor.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bundle_indices, expected_error, create_placement_group",
|
||||
[
|
||||
("1,1", "cannot have duplicate values,", 4),
|
||||
("0,1,2", "must have the same size", 4),
|
||||
],
|
||||
indirect=["create_placement_group"],
|
||||
)
|
||||
def test_ray_v2_invalid_bundle_indices(
|
||||
bundle_indices, expected_error, create_placement_group, monkeypatch
|
||||
):
|
||||
"""Validate invalid bundle indices are rejected."""
|
||||
monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", bundle_indices)
|
||||
vllm_config = create_vllm_config(
|
||||
tensor_parallel_size=2, placement_group=create_placement_group
|
||||
)
|
||||
with pytest.raises(AssertionError, match=expected_error):
|
||||
RayExecutorV2(vllm_config=vllm_config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp_size, pp_size", [(2, 1), (2, 2)])
|
||||
def test_ray_v2_single_node_generation_with_pg(tp_size, pp_size):
|
||||
"""E2E LLM generation with a user-provided placement group."""
|
||||
ensure_ray_initialized()
|
||||
bundles = [{"GPU": 1, "CPU": 1} for _ in range(tp_size * pp_size)]
|
||||
pg = ray.util.placement_group(bundles, strategy="PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
try:
|
||||
with patch.object(ray.util, "get_current_placement_group", return_value=pg):
|
||||
llm = LLM(
|
||||
model=MODEL,
|
||||
tensor_parallel_size=tp_size,
|
||||
pipeline_parallel_size=pp_size,
|
||||
distributed_executor_backend="ray",
|
||||
enforce_eager=True,
|
||||
max_model_len=256,
|
||||
gpu_memory_utilization=0.3,
|
||||
)
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
outputs = llm.generate(prompts)
|
||||
|
||||
assert len(outputs) == len(prompts)
|
||||
for output in outputs:
|
||||
assert len(output.outputs) > 0
|
||||
assert len(output.outputs[0].text) > 0
|
||||
finally:
|
||||
llm.llm_engine.model_executor.shutdown()
|
||||
del llm
|
||||
gc.collect()
|
||||
@@ -0,0 +1,209 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Orchestration-level integration tests for RayExecutorV2.
|
||||
"""
|
||||
|
||||
import gc
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
|
||||
pytestmark = pytest.mark.usefixtures("enable_ray_v2_backend")
|
||||
|
||||
MODEL = "facebook/opt-125m"
|
||||
|
||||
|
||||
def _get_env_var(worker, name):
|
||||
return os.environ.get(name)
|
||||
|
||||
|
||||
def _ray_init():
|
||||
"""Start Ray with the project root on workers' PYTHONPATH.
|
||||
|
||||
Without this, workers cannot unpickle actor classes defined in the
|
||||
``tests`` package, causing FunctionActorManager to fall back to
|
||||
TemporaryActor which drops async method signatures."""
|
||||
project_root = str(pathlib.Path(__file__).resolve().parents[2])
|
||||
ray.init(
|
||||
ignore_reinit_error=True,
|
||||
runtime_env={"env_vars": {"PYTHONPATH": project_root}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_init():
|
||||
_ray_init()
|
||||
|
||||
|
||||
class _AsyncLLMActor:
|
||||
def start(self, pg, bundle_indices=None, ray_runtime_env=None):
|
||||
os.environ["VLLM_USE_RAY_V2_EXECUTOR_BACKEND"] = "1"
|
||||
# Needed so collective_rpc can pickle _get_env_var over the
|
||||
# AsyncLLM -> EngineCore ZMQ boundary.
|
||||
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
|
||||
if bundle_indices is not None:
|
||||
os.environ["VLLM_RAY_BUNDLE_INDICES"] = bundle_indices
|
||||
else:
|
||||
os.environ.pop("VLLM_RAY_BUNDLE_INDICES", None)
|
||||
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
from vllm.v1.executor.abstract import Executor
|
||||
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=MODEL,
|
||||
tensor_parallel_size=2,
|
||||
distributed_executor_backend="ray",
|
||||
enforce_eager=True,
|
||||
max_model_len=256,
|
||||
gpu_memory_utilization=0.8,
|
||||
)
|
||||
vllm_config = engine_args.create_engine_config()
|
||||
vllm_config.parallel_config.placement_group = pg
|
||||
if ray_runtime_env is not None:
|
||||
vllm_config.parallel_config.ray_runtime_env = ray_runtime_env
|
||||
|
||||
executor_class = Executor.get_class(vllm_config)
|
||||
self.engine = AsyncLLM(
|
||||
vllm_config=vllm_config,
|
||||
executor_class=executor_class,
|
||||
log_stats=False,
|
||||
log_requests=False,
|
||||
)
|
||||
|
||||
async def generate(self, prompt):
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
params = SamplingParams(max_tokens=16)
|
||||
result = None
|
||||
async for output in self.engine.generate(
|
||||
prompt, params, request_id="test_request_id"
|
||||
):
|
||||
result = output
|
||||
assert result is not None
|
||||
return result.outputs[0].text
|
||||
|
||||
async def generate_and_get_worker_envs(self, prompt, env_names):
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
params = SamplingParams(max_tokens=16)
|
||||
result = None
|
||||
async for output in self.engine.generate(
|
||||
prompt, params, request_id="test_request_id"
|
||||
):
|
||||
result = output
|
||||
assert result is not None
|
||||
text = result.outputs[0].text
|
||||
|
||||
env_results = {}
|
||||
for name in env_names:
|
||||
vals = await self.engine.collective_rpc(
|
||||
_get_env_var, timeout=10, args=(name,)
|
||||
)
|
||||
env_results[name] = vals
|
||||
return text, env_results
|
||||
|
||||
def shutdown(self):
|
||||
if engine := getattr(self, "engine", None):
|
||||
engine.shutdown()
|
||||
del self.engine
|
||||
gc.collect()
|
||||
|
||||
|
||||
AsyncLLMActor = ray.remote(num_cpus=0, max_concurrency=1)(_AsyncLLMActor)
|
||||
|
||||
|
||||
def test_multi_replicas(ray_init):
|
||||
pg1 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
|
||||
pg2 = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
|
||||
ray.get([pg1.ready(), pg2.ready()])
|
||||
|
||||
actor1 = AsyncLLMActor.remote()
|
||||
actor2 = AsyncLLMActor.remote()
|
||||
|
||||
ray.get(actor1.start.remote(pg1))
|
||||
ray.get(actor2.start.remote(pg2))
|
||||
|
||||
out1, out2 = ray.get(
|
||||
[
|
||||
actor1.generate.remote("Hello world"),
|
||||
actor2.generate.remote("Hello world"),
|
||||
]
|
||||
)
|
||||
assert len(out1) > 0
|
||||
assert len(out2) > 0
|
||||
|
||||
|
||||
def test_multi_replicas_with_bundle_indices(ray_init):
|
||||
pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 4, strategy="PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
actor1 = AsyncLLMActor.remote()
|
||||
actor2 = AsyncLLMActor.remote()
|
||||
|
||||
ray.get(actor1.start.remote(pg, bundle_indices="2,1"))
|
||||
ray.get(actor2.start.remote(pg, bundle_indices="0,3"))
|
||||
|
||||
out1, out2 = ray.get(
|
||||
[
|
||||
actor1.generate.remote("Hello world"),
|
||||
actor2.generate.remote("Hello world"),
|
||||
]
|
||||
)
|
||||
assert len(out1) > 0
|
||||
assert len(out2) > 0
|
||||
|
||||
|
||||
def test_env_var_and_runtime_env_propagation():
|
||||
"""
|
||||
Verify env vars (NCCL_, HF_) and parallel_config.ray_runtime_env
|
||||
propagate to RayWorkerProc actors.
|
||||
"""
|
||||
sentinel_vars = {
|
||||
"NCCL_DEBUG": "INFO",
|
||||
"HF_TOKEN": "test_sentinel_token",
|
||||
}
|
||||
for k, v in sentinel_vars.items():
|
||||
os.environ[k] = v
|
||||
|
||||
try:
|
||||
# Called directly (not via the ray_init fixture) because sentinel
|
||||
# env vars must be in os.environ before ray.init() so that Ray
|
||||
# worker processes inherit them.
|
||||
_ray_init()
|
||||
|
||||
pg = ray.util.placement_group([{"GPU": 1, "CPU": 1}] * 2, strategy="PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
# Include the project root so that RayWorkerProc actors can
|
||||
# unpickle _get_env_var.
|
||||
project_root = str(pathlib.Path(__file__).resolve().parents[2])
|
||||
ray_runtime_env = {
|
||||
"env_vars": {
|
||||
"RAY_RUNTIME_ENV_TEST": "ray_runtime_env",
|
||||
"PYTHONPATH": project_root,
|
||||
},
|
||||
}
|
||||
|
||||
actor = AsyncLLMActor.remote()
|
||||
ray.get(actor.start.remote(pg, ray_runtime_env=ray_runtime_env))
|
||||
|
||||
all_env_names = list(sentinel_vars) + ["RAY_RUNTIME_ENV_TEST"]
|
||||
text, env_results = ray.get(
|
||||
actor.generate_and_get_worker_envs.remote("Hello world", all_env_names)
|
||||
)
|
||||
assert len(text) > 0
|
||||
|
||||
for name, expected in sentinel_vars.items():
|
||||
for val in env_results[name]:
|
||||
assert val == expected
|
||||
|
||||
for val in env_results["RAY_RUNTIME_ENV_TEST"]:
|
||||
assert val == "ray_runtime_env"
|
||||
|
||||
finally:
|
||||
for k in sentinel_vars:
|
||||
os.environ.pop(k, None)
|
||||
@@ -0,0 +1,134 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm._aiter_ops import is_aiter_found, rocm_aiter_ops
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce # noqa
|
||||
from vllm.distributed.parallel_state import get_tp_group, graph_capture
|
||||
from vllm.envs import disable_envs_cache
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
from ..utils import (
|
||||
assert_rocm_custom_allreduce_backend_state,
|
||||
ensure_model_parallel_initialized,
|
||||
init_test_distributed_environment,
|
||||
multi_gpu_test,
|
||||
multi_process_parallel,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="ROCm-only AITER custom allreduce tests",
|
||||
)
|
||||
|
||||
test_cases = [
|
||||
((2, 7168), torch.float16),
|
||||
((2, 7168), torch.bfloat16),
|
||||
((128, 8192), torch.float16),
|
||||
((128, 8192), torch.bfloat16),
|
||||
]
|
||||
|
||||
|
||||
def _configure_aiter_custom_ar_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising=False)
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
|
||||
monkeypatch.setenv("VLLM_ROCM_USE_AITER_CUSTOM_AR", "1")
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", "NONE")
|
||||
disable_envs_cache()
|
||||
rocm_aiter_ops.refresh_env_variables()
|
||||
|
||||
|
||||
def _assert_aiter_handles_input(inp: torch.Tensor) -> None:
|
||||
aiter_ar_comm = get_tp_group().device_communicator.aiter_ar_comm
|
||||
assert aiter_ar_comm is not None
|
||||
assert aiter_ar_comm.should_custom_ar(inp), (
|
||||
f"AITER CustomAllreduce does not support input shape {inp.shape}."
|
||||
)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def graph_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
) -> None:
|
||||
with monkeypatch.context() as m:
|
||||
_configure_aiter_custom_ar_env(m)
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
ensure_model_parallel_initialized(tp_size, pp_size)
|
||||
assert_rocm_custom_allreduce_backend_state(True, "NONE")
|
||||
group = get_tp_group().device_group
|
||||
|
||||
# A small all_reduce for warmup.
|
||||
# this is needed because device communicators might be created lazily
|
||||
# (e.g. NCCL). This will ensure that the communicator is initialized
|
||||
# before any communication happens, so that this group can be used for
|
||||
# graph capture immediately.
|
||||
data = torch.zeros(1)
|
||||
data = data.to(device=device)
|
||||
dist.all_reduce(data, group=group)
|
||||
torch.accelerator.synchronize()
|
||||
del data
|
||||
|
||||
for shape, dtype in test_cases:
|
||||
with graph_capture(device=device) as graph_capture_context:
|
||||
inp = torch.ones(shape, dtype=dtype, device=device)
|
||||
_assert_aiter_handles_input(inp)
|
||||
expected = inp * tp_size
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph, stream=graph_capture_context.stream):
|
||||
out = tensor_model_parallel_all_reduce(inp)
|
||||
|
||||
graph.replay()
|
||||
torch.testing.assert_close(out, expected)
|
||||
|
||||
|
||||
@ray.remote(num_gpus=1, max_calls=1)
|
||||
def eager_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pp_size,
|
||||
rank,
|
||||
distributed_init_port,
|
||||
) -> None:
|
||||
with monkeypatch.context() as m:
|
||||
_configure_aiter_custom_ar_env(m)
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
init_test_distributed_environment(tp_size, pp_size, rank, distributed_init_port)
|
||||
ensure_model_parallel_initialized(tp_size, pp_size)
|
||||
assert_rocm_custom_allreduce_backend_state(True, "NONE")
|
||||
|
||||
for shape, dtype in test_cases:
|
||||
inp = torch.ones(shape, dtype=dtype, device=device)
|
||||
_assert_aiter_handles_input(inp)
|
||||
expected = inp * tp_size
|
||||
out = tensor_model_parallel_all_reduce(inp)
|
||||
torch.testing.assert_close(out, expected)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_aiter_found(), reason="AITER is not installed")
|
||||
@multi_gpu_test(num_gpus=2)
|
||||
@pytest.mark.parametrize("tp_size", [2])
|
||||
@pytest.mark.parametrize("pipeline_parallel_size", [1])
|
||||
@pytest.mark.parametrize("test_target", [eager_allreduce, graph_allreduce])
|
||||
def test_rocm_aiter_custom_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tp_size,
|
||||
pipeline_parallel_size,
|
||||
test_target,
|
||||
):
|
||||
multi_process_parallel(monkeypatch, tp_size, pipeline_parallel_size, test_target)
|
||||
@@ -0,0 +1,750 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import contextlib
|
||||
import importlib
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import queue
|
||||
import traceback
|
||||
from functools import lru_cache
|
||||
from types import SimpleNamespace
|
||||
from typing import Literal
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_rocm(),
|
||||
reason="ROCm-only quick-reduce tests",
|
||||
)
|
||||
|
||||
MB = 1024 * 1024
|
||||
WORLD_SIZE = 2
|
||||
QUANT_LEVELS = ["FP", "INT8", "INT6", "INT4"]
|
||||
|
||||
|
||||
def _log(message: str) -> None:
|
||||
print(f"[rocm_quick_reduce] {message}", flush=True)
|
||||
|
||||
|
||||
def _reload_envs():
|
||||
return importlib.reload(envs)
|
||||
|
||||
|
||||
def _make_quick_allreduce(
|
||||
*,
|
||||
disabled: bool = False,
|
||||
world_size: int = 2,
|
||||
quant_level: str = "FP",
|
||||
use_fp16_kernels: bool = False,
|
||||
qr_max_size: int = 64 * MB,
|
||||
):
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import (
|
||||
QuickAllReduce,
|
||||
QuickReduceRegime,
|
||||
)
|
||||
|
||||
qar = QuickAllReduce.__new__(QuickAllReduce)
|
||||
qar.disabled = disabled
|
||||
qar.world_size = world_size
|
||||
qar.use_fp16_kernels = use_fp16_kernels
|
||||
qar.qr_quant_level = QuickReduceRegime[quant_level]
|
||||
qar.qr_max_size = qr_max_size
|
||||
return qar
|
||||
|
||||
|
||||
def _quick_allreduce_worker(
|
||||
rank: int,
|
||||
port: int,
|
||||
quant_level: str,
|
||||
dtype_name: str,
|
||||
cast_bf16: bool,
|
||||
):
|
||||
os.environ["VLLM_ROCM_QUICK_REDUCE_QUANTIZATION"] = quant_level
|
||||
os.environ["VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16"] = "1" if cast_bf16 else "0"
|
||||
_log(
|
||||
f"worker start: rank={rank} quant={quant_level} "
|
||||
f"dtype={dtype_name} cast_bf16={cast_bf16}"
|
||||
)
|
||||
|
||||
device = torch.device(f"cuda:{rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
dist.init_process_group(
|
||||
backend="gloo",
|
||||
init_method=f"tcp://127.0.0.1:{port}",
|
||||
rank=rank,
|
||||
world_size=WORLD_SIZE,
|
||||
)
|
||||
|
||||
qar = None
|
||||
try:
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import (
|
||||
QuickAllReduce,
|
||||
)
|
||||
|
||||
qar = QuickAllReduce(group=dist.GroupMember.WORLD, device=rank)
|
||||
assert not qar.disabled
|
||||
|
||||
num_elements = 8 * MB if dtype_name == "float16" else 4 * MB
|
||||
|
||||
dtype = getattr(torch, dtype_name)
|
||||
inp = torch.ones(num_elements, dtype=dtype, device=device)
|
||||
|
||||
assert qar.should_quick_allreduce(inp)
|
||||
if cast_bf16:
|
||||
assert qar.use_fp16_kernels
|
||||
|
||||
out = qar.quick_all_reduce(inp)
|
||||
assert torch.allclose(out, inp * WORLD_SIZE, atol=2.5, rtol=0.1)
|
||||
_log(
|
||||
f"worker complete: rank={rank} quant={quant_level} "
|
||||
f"dtype={dtype_name} num_elements={num_elements} "
|
||||
f"use_fp16_kernels={qar.use_fp16_kernels}"
|
||||
)
|
||||
finally:
|
||||
if qar is not None:
|
||||
qar.close()
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _run_two_gpu_quick_allreduce_test(
|
||||
*,
|
||||
quant_level: str,
|
||||
dtype_name: str,
|
||||
cast_bf16: bool,
|
||||
):
|
||||
_log(
|
||||
f"launch 2-GPU case: quant={quant_level} "
|
||||
f"dtype={dtype_name} cast_bf16={cast_bf16}"
|
||||
)
|
||||
ctx = mp.get_context("spawn")
|
||||
port = get_open_port()
|
||||
procs = []
|
||||
|
||||
for rank in range(WORLD_SIZE):
|
||||
proc = ctx.Process(
|
||||
target=_quick_allreduce_worker,
|
||||
args=(rank, port, quant_level, dtype_name, cast_bf16),
|
||||
)
|
||||
proc.start()
|
||||
procs.append(proc)
|
||||
|
||||
for proc in procs:
|
||||
proc.join(timeout=60)
|
||||
assert proc.exitcode == 0, f"worker exited with code {proc.exitcode}"
|
||||
_log(
|
||||
f"finished 2-GPU case: quant={quant_level} "
|
||||
f"dtype={dtype_name} cast_bf16={cast_bf16}"
|
||||
)
|
||||
|
||||
|
||||
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
E2E_PREFILL_TOKENS = 1024
|
||||
E2E_MAX_MODEL_LEN = 1536
|
||||
E2E_GPU_MEMORY_UTILIZATION = 0.3
|
||||
E2E_KV_CACHE_MEMORY_BYTES = 2 << 30
|
||||
|
||||
_BACKGROUND_LINE = (
|
||||
"Background filler: this archived operations memo repeats a routine status "
|
||||
"line so the distributed test uses a realistically long prefill."
|
||||
)
|
||||
_BACKGROUND_BLOCK = " ".join([_BACKGROUND_LINE] * 48)
|
||||
|
||||
|
||||
def _build_prompt(*, fact_block: str, question: str) -> str:
|
||||
return (
|
||||
"Read the archived operations memo below. Most of the memo is filler. "
|
||||
"Use only the fact block near the end when answering.\n"
|
||||
f"{_BACKGROUND_BLOCK}\n"
|
||||
"Fact block:\n"
|
||||
f"{fact_block}\n"
|
||||
f"Question: {question}\n"
|
||||
"Answer in one short sentence."
|
||||
)
|
||||
|
||||
|
||||
E2E_PROMPTS = [
|
||||
_build_prompt(
|
||||
fact_block=(
|
||||
"- Festival city: Oslo\n- Mascot animal: otter\n- Welcome drink: tea"
|
||||
),
|
||||
question="Which city hosts the festival, and what animal is the mascot?",
|
||||
),
|
||||
_build_prompt(
|
||||
fact_block=(
|
||||
"- Meeting day: Tuesday\n"
|
||||
"- Planned snack: apricot cake\n"
|
||||
"- Backup room: Cedar"
|
||||
),
|
||||
question="What day is the meeting, and what snack is planned?",
|
||||
),
|
||||
]
|
||||
RECORDED_RESPONSE_TEXTS = (
|
||||
" The city hosting the festival is Oslo, and the mascot is an otter.",
|
||||
" The meeting is on Tuesday and the snack planned is apricot cake.",
|
||||
)
|
||||
REQUIRED_WORDS = (("oslo", "otter"), ("tuesday", "apricot"))
|
||||
|
||||
|
||||
def _log_prompt_summaries() -> None:
|
||||
for i, prompt in enumerate(E2E_PROMPTS):
|
||||
prompt_lines = prompt.splitlines()
|
||||
fact_block = [line for line in prompt_lines if line.startswith("- ")]
|
||||
fact_summary = "; ".join(line.removeprefix("- ") for line in fact_block)
|
||||
_log(f"prompt {i} facts: {fact_summary}")
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _get_model_path() -> str:
|
||||
try:
|
||||
path = snapshot_download(repo_id=MODEL_NAME, local_files_only=True)
|
||||
_log(f"using cached model snapshot: {path}")
|
||||
return path
|
||||
except Exception:
|
||||
path = snapshot_download(repo_id=MODEL_NAME)
|
||||
_log(f"downloaded model snapshot: {path}")
|
||||
return path
|
||||
|
||||
|
||||
def _get_hidden_size(model_config) -> int:
|
||||
hidden_size = getattr(model_config, "hidden_size", None)
|
||||
if hidden_size is None and hasattr(model_config, "text_config"):
|
||||
hidden_size = getattr(model_config.text_config, "hidden_size", None)
|
||||
assert isinstance(hidden_size, int)
|
||||
return hidden_size
|
||||
|
||||
|
||||
def _check_tp_allreduce_uses_quick_reduce(
|
||||
self,
|
||||
num_tokens: int,
|
||||
dtype_name: str = "float16",
|
||||
) -> dict[str, int | bool]:
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
|
||||
assert self.device is not None
|
||||
qr_comm = get_tp_group().device_communicator.qr_comm
|
||||
assert qr_comm is not None
|
||||
assert not qr_comm.disabled
|
||||
|
||||
hidden_size = _get_hidden_size(self.model_runner.model.config)
|
||||
dtype = getattr(torch, dtype_name)
|
||||
sample = torch.full(
|
||||
(num_tokens, hidden_size),
|
||||
fill_value=float(self.rank + 1),
|
||||
dtype=dtype,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
assert qr_comm.should_quick_allreduce(sample)
|
||||
|
||||
expected = sample.clone()
|
||||
reduced = tensor_model_parallel_all_reduce(sample)
|
||||
dist.all_reduce(expected, group=get_tp_group().device_group)
|
||||
torch.testing.assert_close(reduced, expected, atol=2.5, rtol=0.1)
|
||||
|
||||
stats = {
|
||||
"rank": self.rank,
|
||||
"hidden_size": hidden_size,
|
||||
"num_tokens": num_tokens,
|
||||
"use_fp16_kernels": qr_comm.use_fp16_kernels,
|
||||
}
|
||||
_log(
|
||||
"worker quick-reduce check: "
|
||||
f"rank={self.rank} hidden_size={hidden_size} "
|
||||
f"num_tokens={num_tokens} use_fp16_kernels={qr_comm.use_fp16_kernels}"
|
||||
)
|
||||
return stats
|
||||
|
||||
|
||||
def _check_quick_reduce_disabled(self) -> int:
|
||||
from vllm.distributed.parallel_state import get_tp_group
|
||||
|
||||
qr_comm = get_tp_group().device_communicator.qr_comm
|
||||
assert qr_comm is not None
|
||||
assert qr_comm.disabled
|
||||
_log(f"worker confirmed quick reduce is disabled: rank={self.rank}")
|
||||
return self.rank
|
||||
|
||||
|
||||
def _collect_generations(outputs) -> list[tuple[tuple[int, ...], str]]:
|
||||
return [
|
||||
(tuple(output.outputs[0].token_ids), output.outputs[0].text)
|
||||
for output in outputs
|
||||
]
|
||||
|
||||
|
||||
def _shutdown_llm(llm: LLM | None) -> None:
|
||||
if llm is None:
|
||||
cleanup_dist_env_and_memory()
|
||||
return
|
||||
|
||||
with contextlib.suppress(Exception):
|
||||
llm.llm_engine.engine_core.shutdown()
|
||||
|
||||
del llm
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def _log_generations(
|
||||
label: str,
|
||||
generations: list[tuple[tuple[int, ...], str]],
|
||||
) -> None:
|
||||
for i, (token_ids, text) in enumerate(generations):
|
||||
_log(f"{label} prompt {i} token ids: {list(token_ids)}")
|
||||
_log(f"{label} prompt {i} text: {text!r}")
|
||||
|
||||
|
||||
def _assert_required_words(
|
||||
label: str,
|
||||
generations: list[tuple[tuple[int, ...], str]],
|
||||
) -> None:
|
||||
for i, (_, text) in enumerate(generations):
|
||||
lowered = text.lower()
|
||||
missing = [word for word in REQUIRED_WORDS[i] if word not in lowered]
|
||||
assert not missing, (
|
||||
f"{label} prompt {i} is missing required words {missing}. "
|
||||
f"Observed text: {text!r}"
|
||||
)
|
||||
|
||||
|
||||
def _collect_soft_mismatches(
|
||||
baseline_generations: list[tuple[tuple[int, ...], str]],
|
||||
quick_reduce_generations: list[tuple[tuple[int, ...], str]],
|
||||
) -> list[str]:
|
||||
mismatches = []
|
||||
|
||||
for i, (_, text) in enumerate(baseline_generations):
|
||||
expected = RECORDED_RESPONSE_TEXTS[i]
|
||||
if text != expected:
|
||||
mismatches.append(
|
||||
f"baseline prompt {i} drifted from the recorded response.\n"
|
||||
f"expected={expected!r}\nactual={text!r}"
|
||||
)
|
||||
|
||||
for i, (_, text) in enumerate(quick_reduce_generations):
|
||||
expected = RECORDED_RESPONSE_TEXTS[i]
|
||||
if text != expected:
|
||||
mismatches.append(
|
||||
f"quick-reduce prompt {i} drifted from the recorded response.\n"
|
||||
f"expected={expected!r}\nactual={text!r}"
|
||||
)
|
||||
|
||||
for i, ((_, baseline_text), (_, quick_reduce_text)) in enumerate(
|
||||
zip(baseline_generations, quick_reduce_generations)
|
||||
):
|
||||
if baseline_text != quick_reduce_text:
|
||||
mismatches.append(
|
||||
f"baseline and quick-reduce responses differ for prompt {i}.\n"
|
||||
f"baseline={baseline_text!r}\nquick_reduce={quick_reduce_text!r}"
|
||||
)
|
||||
|
||||
return mismatches
|
||||
|
||||
|
||||
def _run_generation(
|
||||
*,
|
||||
backend: Literal["mp", "ray"],
|
||||
quant_mode: str,
|
||||
expect_quick_reduce: bool,
|
||||
) -> list[tuple[tuple[int, ...], str]]:
|
||||
llm = None
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
m.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_mode)
|
||||
model_path = _get_model_path()
|
||||
_log(
|
||||
f"starting generation: backend={backend} quant={quant_mode} "
|
||||
f"gpu_memory_utilization={E2E_GPU_MEMORY_UTILIZATION} "
|
||||
f"kv_cache_bytes={E2E_KV_CACHE_MEMORY_BYTES} model={model_path}"
|
||||
)
|
||||
|
||||
try:
|
||||
llm = LLM(
|
||||
model=model_path,
|
||||
tokenizer=model_path,
|
||||
tensor_parallel_size=2,
|
||||
distributed_executor_backend=backend,
|
||||
dtype="half",
|
||||
enforce_eager=True,
|
||||
max_model_len=E2E_MAX_MODEL_LEN,
|
||||
max_num_seqs=len(E2E_PROMPTS),
|
||||
gpu_memory_utilization=E2E_GPU_MEMORY_UTILIZATION,
|
||||
kv_cache_memory_bytes=E2E_KV_CACHE_MEMORY_BYTES,
|
||||
seed=0,
|
||||
)
|
||||
|
||||
if not expect_quick_reduce:
|
||||
assert llm.collective_rpc(_check_quick_reduce_disabled) == [0, 1]
|
||||
|
||||
if expect_quick_reduce:
|
||||
worker_stats = llm.collective_rpc(
|
||||
_check_tp_allreduce_uses_quick_reduce,
|
||||
args=(E2E_PREFILL_TOKENS,),
|
||||
)
|
||||
assert [stat["rank"] for stat in worker_stats] == [0, 1]
|
||||
worker_summary = "; ".join(
|
||||
"rank={rank} hidden_size={hidden_size} num_tokens={num_tokens} "
|
||||
"use_fp16_kernels={use_fp16_kernels}".format(**stat)
|
||||
for stat in worker_stats
|
||||
)
|
||||
_log(f"{backend} quick-reduce worker checks: {worker_summary}")
|
||||
|
||||
outputs = llm.generate(
|
||||
E2E_PROMPTS,
|
||||
SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=20,
|
||||
stop=["\nAnswer:", " Answer:"],
|
||||
),
|
||||
use_tqdm=False,
|
||||
)
|
||||
generations = _collect_generations(outputs)
|
||||
assert all(text.strip() for _, text in generations)
|
||||
_log_generations(f"{backend} {quant_mode}", generations)
|
||||
return generations
|
||||
finally:
|
||||
_shutdown_llm(llm)
|
||||
|
||||
|
||||
def _run_quick_reduce_llm_e2e_in_subprocess(
|
||||
*,
|
||||
backend: Literal["mp", "ray"],
|
||||
) -> str | None:
|
||||
_log(f"running LLM e2e: backend={backend}")
|
||||
_log_prompt_summaries()
|
||||
baseline_outputs = _run_generation(
|
||||
backend=backend,
|
||||
quant_mode="NONE",
|
||||
expect_quick_reduce=False,
|
||||
)
|
||||
quick_reduce_outputs = _run_generation(
|
||||
backend=backend,
|
||||
quant_mode="FP",
|
||||
expect_quick_reduce=True,
|
||||
)
|
||||
|
||||
_assert_required_words("baseline", baseline_outputs)
|
||||
_assert_required_words("quick-reduce", quick_reduce_outputs)
|
||||
|
||||
mismatches = _collect_soft_mismatches(baseline_outputs, quick_reduce_outputs)
|
||||
if mismatches:
|
||||
details = "\n\n".join(mismatches)
|
||||
_log(f"soft response mismatch:\n{details}")
|
||||
return details
|
||||
|
||||
_log(f"LLM e2e backend={backend} matched the recorded responses exactly")
|
||||
return None
|
||||
|
||||
|
||||
def _quick_reduce_llm_e2e_worker(
|
||||
result_queue: mp.Queue,
|
||||
backend: Literal["mp", "ray"],
|
||||
) -> None:
|
||||
try:
|
||||
xfail_reason = _run_quick_reduce_llm_e2e_in_subprocess(backend=backend)
|
||||
except Exception:
|
||||
result_queue.put({"status": "error", "reason": traceback.format_exc()})
|
||||
raise
|
||||
else:
|
||||
if xfail_reason is not None:
|
||||
result_queue.put({"status": "xfail", "reason": xfail_reason})
|
||||
else:
|
||||
result_queue.put({"status": "ok"})
|
||||
|
||||
|
||||
def run_quick_reduce_llm_e2e(
|
||||
*,
|
||||
backend: Literal["mp", "ray"],
|
||||
) -> None:
|
||||
ctx = mp.get_context("spawn")
|
||||
result_queue = ctx.Queue()
|
||||
proc = ctx.Process(
|
||||
target=_quick_reduce_llm_e2e_worker,
|
||||
args=(result_queue, backend),
|
||||
)
|
||||
proc.start()
|
||||
proc.join(timeout=600)
|
||||
|
||||
try:
|
||||
result = result_queue.get(timeout=5)
|
||||
except queue.Empty as exc:
|
||||
if proc.exitcode != 0:
|
||||
raise AssertionError(
|
||||
f"quick-reduce llm e2e subprocess failed for backend={backend} "
|
||||
f"with exit code {proc.exitcode} and produced no result"
|
||||
) from exc
|
||||
raise AssertionError(
|
||||
f"quick-reduce llm e2e subprocess produced no result for backend={backend}"
|
||||
) from exc
|
||||
|
||||
if result["status"] == "xfail":
|
||||
pytest.xfail(result["reason"])
|
||||
if result["status"] == "error":
|
||||
raise AssertionError(
|
||||
f"quick-reduce llm e2e subprocess failed for backend={backend}:\n"
|
||||
f"{result['reason']}"
|
||||
)
|
||||
|
||||
assert proc.exitcode == 0, (
|
||||
f"quick-reduce llm e2e subprocess failed for backend={backend} "
|
||||
f"with exit code {proc.exitcode}"
|
||||
)
|
||||
|
||||
|
||||
def test_quick_reduce_regime_values():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
|
||||
|
||||
assert QuickReduceRegime.FP.value == 0
|
||||
assert QuickReduceRegime.INT8.value == 1
|
||||
assert QuickReduceRegime.INT6.value == 2
|
||||
assert QuickReduceRegime.INT4.value == 3
|
||||
assert QuickReduceRegime.NONE.value == 4
|
||||
|
||||
|
||||
def test_quick_reduce_regime_names():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickReduceRegime
|
||||
|
||||
assert set(QuickReduceRegime.__members__) == {"FP", "INT8", "INT6", "INT4", "NONE"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quant_level", QUANT_LEVELS + ["NONE"])
|
||||
def test_quick_reduce_quantization_env_var(monkeypatch, quant_level):
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", quant_level)
|
||||
|
||||
reloaded_envs = _reload_envs()
|
||||
assert quant_level == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION
|
||||
|
||||
|
||||
def test_quick_reduce_quantization_default(monkeypatch):
|
||||
monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_QUANTIZATION", raising=False)
|
||||
|
||||
reloaded_envs = _reload_envs()
|
||||
assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_QUANTIZATION == "NONE"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cast_bf16", [True, False])
|
||||
def test_quick_reduce_cast_bf16_to_fp16_env_var(monkeypatch, cast_bf16):
|
||||
monkeypatch.setenv(
|
||||
"VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", "1" if cast_bf16 else "0"
|
||||
)
|
||||
|
||||
reloaded_envs = _reload_envs()
|
||||
assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is cast_bf16
|
||||
|
||||
|
||||
def test_quick_reduce_cast_bf16_to_fp16_default(monkeypatch):
|
||||
monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16", raising=False)
|
||||
|
||||
reloaded_envs = _reload_envs()
|
||||
assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16 is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_mb", [128, 512, 2048, None])
|
||||
def test_quick_reduce_max_size_env_var(monkeypatch, max_mb):
|
||||
if max_mb is None:
|
||||
monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", str(max_mb))
|
||||
|
||||
reloaded_envs = _reload_envs()
|
||||
assert max_mb == reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB
|
||||
|
||||
|
||||
def test_quick_reduce_max_size_default(monkeypatch):
|
||||
monkeypatch.delenv("VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB", raising=False)
|
||||
|
||||
reloaded_envs = _reload_envs()
|
||||
assert reloaded_envs.VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("gcn_arch_name", "expected"),
|
||||
[
|
||||
("gfx942", True),
|
||||
("gfx950", True),
|
||||
("gfx90a", False),
|
||||
("", False),
|
||||
],
|
||||
)
|
||||
def test_quick_allreduce_rocm_arch_available(gcn_arch_name, expected):
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
|
||||
|
||||
qar = QuickAllReduce.__new__(QuickAllReduce)
|
||||
qar.disabled = True
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.distributed.device_communicators.quick_all_reduce.current_platform."
|
||||
"is_rocm",
|
||||
return_value=True,
|
||||
),
|
||||
patch(
|
||||
"torch.cuda.get_device_properties",
|
||||
return_value=SimpleNamespace(gcnArchName=gcn_arch_name),
|
||||
),
|
||||
):
|
||||
assert qar._rocm_arch_available() is expected
|
||||
|
||||
|
||||
def test_quick_allreduce_rocm_arch_available_handles_probe_failure():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
|
||||
|
||||
qar = QuickAllReduce.__new__(QuickAllReduce)
|
||||
qar.disabled = True
|
||||
|
||||
with (
|
||||
patch(
|
||||
"vllm.distributed.device_communicators.quick_all_reduce.current_platform."
|
||||
"is_rocm",
|
||||
return_value=True,
|
||||
),
|
||||
patch("torch.cuda.get_device_properties", side_effect=RuntimeError),
|
||||
):
|
||||
assert qar._rocm_arch_available() is False
|
||||
|
||||
|
||||
def test_quick_allreduce_rejects_disabled():
|
||||
qar = _make_quick_allreduce(disabled=True)
|
||||
|
||||
inp = torch.zeros(1024, dtype=torch.float16)
|
||||
assert qar.should_quick_allreduce(inp) is False
|
||||
|
||||
|
||||
def test_quick_allreduce_rejects_unsupported_dtype():
|
||||
qar = _make_quick_allreduce()
|
||||
|
||||
inp = torch.zeros(1024 * 1024, dtype=torch.float32)
|
||||
assert qar.should_quick_allreduce(inp) is False
|
||||
|
||||
|
||||
def test_quick_allreduce_rejects_non_aligned_input():
|
||||
qar = _make_quick_allreduce()
|
||||
|
||||
inp = torch.zeros(5, dtype=torch.float16)
|
||||
assert qar.should_quick_allreduce(inp) is False
|
||||
|
||||
|
||||
def test_quick_allreduce_rejects_non_contiguous_input():
|
||||
qar = _make_quick_allreduce()
|
||||
|
||||
inp = torch.zeros((1024, 1024), dtype=torch.float16)[:, ::2]
|
||||
assert qar.should_quick_allreduce(inp) is False
|
||||
|
||||
|
||||
def test_quick_allreduce_rejects_input_smaller_than_threshold():
|
||||
qar = _make_quick_allreduce()
|
||||
|
||||
inp = torch.zeros((MB // 2) - 8, dtype=torch.float16)
|
||||
assert qar.should_quick_allreduce(inp) is False
|
||||
|
||||
|
||||
def test_quick_allreduce_accepts_input_at_threshold():
|
||||
qar = _make_quick_allreduce()
|
||||
|
||||
inp = torch.zeros(MB // 2, dtype=torch.float16)
|
||||
assert qar.should_quick_allreduce(inp) is True
|
||||
|
||||
|
||||
def test_quick_allreduce_rejects_input_larger_than_max_size():
|
||||
qar = _make_quick_allreduce(qr_max_size=1 * MB)
|
||||
|
||||
inp = torch.zeros(MB, dtype=torch.float16)
|
||||
assert qar.should_quick_allreduce(inp) is False
|
||||
|
||||
|
||||
def test_quick_allreduce_bf16_uses_fp16_threshold_when_cast_enabled():
|
||||
inp = torch.zeros(MB // 2, dtype=torch.bfloat16)
|
||||
|
||||
without_cast = _make_quick_allreduce(use_fp16_kernels=False)
|
||||
with_cast = _make_quick_allreduce(use_fp16_kernels=True)
|
||||
|
||||
assert without_cast.should_quick_allreduce(inp) is False
|
||||
assert with_cast.should_quick_allreduce(inp) is True
|
||||
|
||||
|
||||
def test_quick_allreduce_supported_world_sizes():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
|
||||
|
||||
assert QuickAllReduce._SUPPORTED_WORLD_SIZES == [2, 4, 8]
|
||||
|
||||
|
||||
def test_quick_allreduce_supported_dtypes():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
|
||||
|
||||
assert [torch.float16, torch.bfloat16] == QuickAllReduce._SUPPORTED_DTYPES
|
||||
|
||||
|
||||
def test_quick_allreduce_min_size_table():
|
||||
from vllm.distributed.device_communicators.quick_all_reduce import QuickAllReduce
|
||||
|
||||
for dtype in [torch.float16, torch.bfloat16]:
|
||||
for world_size in QuickAllReduce._SUPPORTED_WORLD_SIZES:
|
||||
min_sizes = QuickAllReduce._QR_MIN_SIZE[(dtype, world_size)]
|
||||
assert len(min_sizes) == 4
|
||||
assert all(size > 0 for size in min_sizes)
|
||||
|
||||
|
||||
def test_qr_max_size():
|
||||
from vllm import _custom_ops as ops
|
||||
|
||||
max_size = ops.qr_max_size()
|
||||
assert isinstance(max_size, int)
|
||||
assert max_size > 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.device_count() < WORLD_SIZE,
|
||||
reason="requires 2 ROCm GPUs",
|
||||
)
|
||||
@pytest.mark.parametrize("quant_level", QUANT_LEVELS)
|
||||
def test_quick_allreduce_two_gpu_correctness(quant_level):
|
||||
_log(f"two-GPU correctness case: quant={quant_level}")
|
||||
_run_two_gpu_quick_allreduce_test(
|
||||
quant_level=quant_level,
|
||||
dtype_name="float16",
|
||||
cast_bf16=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.device_count() < WORLD_SIZE,
|
||||
reason="requires 2 ROCm GPUs",
|
||||
)
|
||||
def test_quick_allreduce_bf16_cast_mode():
|
||||
_log("BF16 cast case")
|
||||
_run_two_gpu_quick_allreduce_test(
|
||||
quant_level="FP",
|
||||
dtype_name="bfloat16",
|
||||
cast_bf16=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.device_count() < WORLD_SIZE,
|
||||
reason="requires 2 ROCm GPUs",
|
||||
)
|
||||
def test_quick_allreduce_llm_e2e():
|
||||
_log("LLM e2e case: backend=mp")
|
||||
run_quick_reduce_llm_e2e(backend="mp")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.device_count() < WORLD_SIZE,
|
||||
reason="requires 2 ROCm GPUs",
|
||||
)
|
||||
def test_quick_allreduce_llm_e2e_ray():
|
||||
_log("LLM e2e case: backend=ray")
|
||||
run_quick_reduce_llm_e2e(backend="ray")
|
||||
@@ -0,0 +1,49 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import os
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm.distributed.parallel_state import in_the_same_node_as
|
||||
from vllm.distributed.utils import StatelessProcessGroup
|
||||
from vllm.utils.network_utils import get_ip, get_open_port
|
||||
|
||||
|
||||
def _run_test(pg):
|
||||
test_result = all(in_the_same_node_as(pg, source_rank=0))
|
||||
|
||||
expected = os.environ.get("VLLM_TEST_SAME_HOST", "1") == "1"
|
||||
assert test_result == expected, f"Expected {expected}, got {test_result}"
|
||||
if pg == dist.group.WORLD:
|
||||
print("Same node test passed! when using torch distributed!")
|
||||
else:
|
||||
print("Same node test passed! when using StatelessProcessGroup!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
dist.init_process_group(backend="gloo")
|
||||
|
||||
rank = dist.get_rank()
|
||||
if rank == 0:
|
||||
port = get_open_port()
|
||||
ip = get_ip()
|
||||
dist.broadcast_object_list([ip, port], src=0)
|
||||
else:
|
||||
recv = [None, None]
|
||||
dist.broadcast_object_list(recv, src=0)
|
||||
ip, port = recv
|
||||
|
||||
stateless_pg = StatelessProcessGroup.create(ip, port, rank, dist.get_world_size())
|
||||
|
||||
for pg in [dist.group.WORLD, stateless_pg]:
|
||||
if os.environ.get("VLLM_TEST_WITH_DEFAULT_DEVICE_SET", "0") == "1":
|
||||
default_devices = ["cpu"]
|
||||
if torch.cuda.is_available():
|
||||
default_devices.append("cuda")
|
||||
for device in default_devices:
|
||||
torch.set_default_device(device)
|
||||
_run_test(pg)
|
||||
else:
|
||||
_run_test(pg)
|
||||
@@ -0,0 +1,394 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from unittest import mock
|
||||
|
||||
import multiprocess as mp
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch.distributed as dist
|
||||
|
||||
from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
|
||||
from vllm.distributed.utils import StatelessProcessGroup
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
|
||||
def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]:
|
||||
np.random.seed(seed)
|
||||
sizes = np.random.randint(1, 10_000, n)
|
||||
# on average, each array will have 5k elements
|
||||
# with int64, each array will have 40kb
|
||||
return [np.random.randint(1, 100, i) for i in sizes]
|
||||
|
||||
|
||||
def distributed_run(fn, world_size, timeout=60):
|
||||
"""Run a function in multiple processes with proper error handling.
|
||||
|
||||
Args:
|
||||
fn: Function to run in each process
|
||||
world_size: Number of processes to spawn
|
||||
timeout: Maximum time in seconds to wait for processes (default: 60)
|
||||
"""
|
||||
number_of_processes = world_size
|
||||
processes = []
|
||||
for i in range(number_of_processes):
|
||||
env = {}
|
||||
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"] = "12345"
|
||||
p = mp.Process(target=fn, args=(env,))
|
||||
processes.append(p)
|
||||
p.start()
|
||||
|
||||
# Monitor processes and fail fast if any process fails
|
||||
start_time = time.time()
|
||||
failed_processes = []
|
||||
|
||||
# Wait for all processes, checking for failures
|
||||
while time.time() - start_time < timeout:
|
||||
all_done = True
|
||||
for i, p in enumerate(processes):
|
||||
if p.is_alive():
|
||||
all_done = False
|
||||
elif p.exitcode != 0:
|
||||
# Process failed
|
||||
failed_processes.append((i, p.exitcode))
|
||||
break
|
||||
|
||||
if failed_processes or all_done:
|
||||
break
|
||||
time.sleep(0.1) # Check every 100ms
|
||||
|
||||
# Check for timeout if no failures detected yet
|
||||
for i, p in enumerate(processes):
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
p.join()
|
||||
|
||||
# Report failures
|
||||
if failed_processes:
|
||||
error_msg = "Distributed test failed:\n"
|
||||
for rank, status in failed_processes:
|
||||
error_msg += f" Rank {rank}: Exit code {status}\n"
|
||||
raise AssertionError(error_msg)
|
||||
|
||||
|
||||
def worker_fn_wrapper(fn):
|
||||
# `mp.Process` cannot accept environment variables directly
|
||||
# so we need to pass the environment variables as arguments
|
||||
# and update the environment variables in the function
|
||||
def wrapped_fn(env):
|
||||
update_environment_variables(env)
|
||||
dist.init_process_group(backend="gloo")
|
||||
fn()
|
||||
|
||||
return wrapped_fn
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn():
|
||||
rank = dist.get_rank()
|
||||
if rank == 0:
|
||||
port = get_open_port()
|
||||
ip = "127.0.0.1"
|
||||
dist.broadcast_object_list([ip, port], src=0)
|
||||
else:
|
||||
recv = [None, None]
|
||||
dist.broadcast_object_list(recv, src=0)
|
||||
ip, port = recv # type: ignore
|
||||
|
||||
stateless_pg = StatelessProcessGroup.create(ip, port, rank, dist.get_world_size())
|
||||
|
||||
for pg in [dist.group.WORLD, stateless_pg]:
|
||||
writer_rank = 2
|
||||
broadcaster = MessageQueue.create_from_process_group(
|
||||
pg, 40 * 1024, 2, writer_rank
|
||||
)
|
||||
if rank == writer_rank:
|
||||
seed = random.randint(0, 1000)
|
||||
dist.broadcast_object_list([seed], writer_rank)
|
||||
else:
|
||||
recv = [None]
|
||||
dist.broadcast_object_list(recv, writer_rank)
|
||||
seed = recv[0] # type: ignore
|
||||
|
||||
if pg == dist.group.WORLD:
|
||||
dist.barrier()
|
||||
else:
|
||||
pg.barrier()
|
||||
|
||||
# in case we find a race condition
|
||||
# print the seed so that we can reproduce the error
|
||||
print(f"Rank {rank} got seed {seed}")
|
||||
# test broadcasting with about 400MB of data
|
||||
N = 10_000
|
||||
if rank == writer_rank:
|
||||
arrs = get_arrays(N, seed)
|
||||
for x in arrs:
|
||||
broadcaster.broadcast_object(x)
|
||||
time.sleep(random.random() / 1000)
|
||||
else:
|
||||
arrs = get_arrays(N, seed)
|
||||
for x in arrs:
|
||||
y = broadcaster.broadcast_object(None)
|
||||
assert np.array_equal(x, y)
|
||||
time.sleep(random.random() / 1000)
|
||||
|
||||
if pg == dist.group.WORLD:
|
||||
dist.barrier()
|
||||
print(f"torch distributed passed the test! Rank {rank}")
|
||||
else:
|
||||
pg.barrier()
|
||||
print(f"StatelessProcessGroup passed the test! Rank {rank}")
|
||||
|
||||
|
||||
def test_shm_broadcast():
|
||||
distributed_run(worker_fn, 4)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn_test_shutdown_busy():
|
||||
rank = dist.get_rank()
|
||||
writer_rank = 2
|
||||
message_queue = MessageQueue.create_from_process_group(
|
||||
dist.group.WORLD, 40 * 1024, 2, writer_rank
|
||||
)
|
||||
|
||||
if not message_queue._is_writer:
|
||||
# Put into busy mode
|
||||
message_queue._spin_condition.busy_loop_s = 9999
|
||||
|
||||
shutdown_event = threading.Event()
|
||||
|
||||
def shutdown_thread(mq, shutdown_event):
|
||||
shutdown_event.wait()
|
||||
mq.shutdown()
|
||||
|
||||
threading.Thread(
|
||||
target=shutdown_thread, args=(message_queue, shutdown_event)
|
||||
).start()
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
message_queue.dequeue(timeout=0.01)
|
||||
|
||||
shutdown_event.set()
|
||||
|
||||
with pytest.raises(RuntimeError, match="cancelled"):
|
||||
message_queue.dequeue(timeout=1)
|
||||
|
||||
assert message_queue.shutting_down
|
||||
|
||||
print(f"torch distributed passed the test! Rank {rank}")
|
||||
dist.barrier()
|
||||
|
||||
|
||||
def test_message_queue_shutdown_busy(caplog_vllm):
|
||||
distributed_run(worker_fn_test_shutdown_busy, 4)
|
||||
print(caplog_vllm.text)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn_test_shutdown_idle():
|
||||
rank = dist.get_rank()
|
||||
writer_rank = 2
|
||||
message_queue = MessageQueue.create_from_process_group(
|
||||
dist.group.WORLD, 40 * 1024, 2, writer_rank
|
||||
)
|
||||
|
||||
if not message_queue._is_writer:
|
||||
# Put into idle mode
|
||||
message_queue._spin_condition.last_read = 0
|
||||
|
||||
shutdown_event = threading.Event()
|
||||
|
||||
def shutdown_thread(mq, shutdown_event):
|
||||
shutdown_event.wait()
|
||||
mq.shutdown()
|
||||
|
||||
threading.Thread(
|
||||
target=shutdown_thread, args=(message_queue, shutdown_event)
|
||||
).start()
|
||||
|
||||
with pytest.raises(TimeoutError):
|
||||
message_queue.dequeue(timeout=0.01)
|
||||
|
||||
shutdown_event.set()
|
||||
|
||||
with pytest.raises(RuntimeError, match="cancelled"):
|
||||
message_queue.dequeue(timeout=1)
|
||||
|
||||
assert message_queue.shutting_down
|
||||
|
||||
print(f"torch distributed passed the test! Rank {rank}")
|
||||
dist.barrier()
|
||||
|
||||
|
||||
def test_message_queue_shutdown_idle():
|
||||
distributed_run(worker_fn_test_shutdown_idle, 4)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn_test_idle_to_busy():
|
||||
rank = dist.get_rank()
|
||||
writer_rank = 2
|
||||
message_queue = MessageQueue.create_from_process_group(
|
||||
dist.group.WORLD, 40 * 1024, 2, writer_rank
|
||||
)
|
||||
|
||||
message1 = "hello world"
|
||||
message2 = np.random.randint(1, 100, 100)
|
||||
with mock.patch.object(
|
||||
message_queue._spin_condition, "wait", wraps=message_queue._spin_condition.wait
|
||||
) as wrapped_wait:
|
||||
if not message_queue._is_writer:
|
||||
# Put into idle mode
|
||||
message_queue._spin_condition.last_read = 0
|
||||
|
||||
# no messages, so expect a TimeoutError
|
||||
with pytest.raises(TimeoutError):
|
||||
message_queue.dequeue(timeout=0.01)
|
||||
# wait should only be called once while idle
|
||||
assert wrapped_wait.call_count == 1
|
||||
|
||||
# sync with the writer and wait for message1
|
||||
dist.barrier()
|
||||
recv_message = message_queue.dequeue(timeout=5)
|
||||
assert recv_message == message1
|
||||
# second call to wait, with a message read, this puts in a busy spin
|
||||
assert wrapped_wait.call_count == 2
|
||||
|
||||
# sync with the writer and wait for message2
|
||||
dist.barrier()
|
||||
recv_message = message_queue.dequeue(timeout=1)
|
||||
assert np.array_equal(recv_message, message2)
|
||||
# in busy mode, we expect wait to have been called multiple times
|
||||
assert wrapped_wait.call_count > 3
|
||||
else:
|
||||
# writer writes two messages in sync with the reader
|
||||
dist.barrier()
|
||||
# sleep delays the send to ensure reader enters the read loop
|
||||
time.sleep(0.1)
|
||||
message_queue.enqueue(message1)
|
||||
|
||||
dist.barrier()
|
||||
time.sleep(0.1)
|
||||
message_queue.enqueue(message2)
|
||||
|
||||
message_queue.shutdown()
|
||||
assert message_queue.shutting_down
|
||||
print(f"torch distributed passed the test! Rank {rank}")
|
||||
|
||||
|
||||
def test_message_queue_idle_wake():
|
||||
distributed_run(worker_fn_test_idle_to_busy, 4)
|
||||
|
||||
|
||||
@worker_fn_wrapper
|
||||
def worker_fn_test_busy_to_idle():
|
||||
rank = dist.get_rank()
|
||||
writer_rank = 2
|
||||
message_queue = MessageQueue.create_from_process_group(
|
||||
dist.group.WORLD, 40 * 1024, 2, writer_rank
|
||||
)
|
||||
|
||||
message1 = 12345
|
||||
message2 = list(range(3))
|
||||
with mock.patch.object(
|
||||
message_queue._spin_condition, "wait", wraps=message_queue._spin_condition.wait
|
||||
) as wrapped_wait:
|
||||
if not message_queue._is_writer:
|
||||
# Put into busy mode
|
||||
message_queue._spin_condition.busy_loop_s = 9999
|
||||
|
||||
# sync with the writer and wait for message1
|
||||
dist.barrier()
|
||||
recv_message = message_queue.dequeue(timeout=1)
|
||||
assert recv_message == message1
|
||||
# in busy mode, we expect wait to have been called many times
|
||||
assert wrapped_wait.call_count > 1
|
||||
|
||||
# simulate busy loop ending
|
||||
message_queue._spin_condition.busy_loop_s = 0
|
||||
# ensure we enter idle mode, then record call count
|
||||
with pytest.raises(TimeoutError):
|
||||
message_queue.dequeue(timeout=0.01)
|
||||
call_count = wrapped_wait.call_count
|
||||
|
||||
# sync with the writer and wait for message2
|
||||
dist.barrier()
|
||||
recv_message = message_queue.dequeue(timeout=1)
|
||||
assert recv_message == message2
|
||||
|
||||
# call to wait after idle should only happen once
|
||||
assert wrapped_wait.call_count == call_count + 1
|
||||
else:
|
||||
# writer writes two messages in sync with the reader
|
||||
dist.barrier()
|
||||
# sleep delays the send to ensure reader enters the read loop
|
||||
time.sleep(0.1)
|
||||
message_queue.enqueue(message1)
|
||||
|
||||
dist.barrier()
|
||||
time.sleep(0.1)
|
||||
message_queue.enqueue(message2)
|
||||
|
||||
message_queue.shutdown()
|
||||
assert message_queue.shutting_down
|
||||
print(f"torch distributed passed the test! Rank {rank}")
|
||||
|
||||
|
||||
def test_message_queue_busy_to_idle():
|
||||
distributed_run(worker_fn_test_busy_to_idle, 4)
|
||||
|
||||
|
||||
def test_warning_logs(caplog_vllm):
|
||||
"""
|
||||
Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals
|
||||
when indefinite=False, and are not emitted when indefinite=True.
|
||||
"""
|
||||
|
||||
# Patch the warning log interval to every 1 ms during reads
|
||||
with mock.patch(
|
||||
"vllm.distributed.device_communicators.shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL",
|
||||
new=0.001, # 1 ms
|
||||
):
|
||||
writer = MessageQueue(
|
||||
n_reader=1,
|
||||
n_local_reader=1,
|
||||
max_chunk_bytes=1024 * 1024, # 1MB chunks
|
||||
max_chunks=10,
|
||||
)
|
||||
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
|
||||
writer.wait_until_ready()
|
||||
reader.wait_until_ready()
|
||||
|
||||
# We should have at least one warning log here
|
||||
# "0 seconds" expected due to rounding of 1ms test interval
|
||||
with pytest.raises(TimeoutError):
|
||||
reader.dequeue(timeout=0.01, indefinite=False)
|
||||
assert any(
|
||||
"No available shared memory broadcast block found in 0 seconds"
|
||||
in record.message
|
||||
for record in caplog_vllm.records
|
||||
)
|
||||
caplog_vllm.clear()
|
||||
|
||||
# We should have no warnings this time
|
||||
with pytest.raises(TimeoutError):
|
||||
reader.dequeue(timeout=0.01, indefinite=True)
|
||||
assert all(
|
||||
"No available shared memory broadcast block found in 0 seconds"
|
||||
not in record.message
|
||||
for record in caplog_vllm.records
|
||||
)
|
||||
|
||||
# Clean up when done
|
||||
writer.shutdown()
|
||||
reader.shutdown()
|
||||
@@ -0,0 +1,243 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import traceback
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm.distributed.device_communicators.shm_object_storage import (
|
||||
SingleWriterShmRingBuffer,
|
||||
)
|
||||
|
||||
|
||||
class TestSingleWriterShmRingBuffer(unittest.TestCase):
|
||||
"""Test suite for the ring buffer implementation"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures"""
|
||||
self.buffer_size = 4096
|
||||
self.ring_buffer = None
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after tests"""
|
||||
if self.ring_buffer:
|
||||
self.ring_buffer.close()
|
||||
|
||||
def test_buffer_opening(self):
|
||||
"""Test opening an existing buffer"""
|
||||
# First create a buffer
|
||||
self.ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=self.buffer_size, create=True
|
||||
)
|
||||
|
||||
# Then open it with another instance
|
||||
reader_buffer = SingleWriterShmRingBuffer(*self.ring_buffer.handle())
|
||||
self.assertFalse(reader_buffer.is_writer)
|
||||
self.assertEqual(
|
||||
reader_buffer.shared_memory.name, self.ring_buffer.shared_memory.name
|
||||
)
|
||||
|
||||
def test_buffer_access(self):
|
||||
"""Test accessing allocated buffers"""
|
||||
self.ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=self.buffer_size, create=True
|
||||
)
|
||||
|
||||
size = 100
|
||||
address, monotonic_id = self.ring_buffer.allocate_buf(size)
|
||||
|
||||
# Write some test data
|
||||
test_data = b"Hello, World!" * 7 # 91 bytes
|
||||
with self.ring_buffer.access_buf(address) as (data_buf, metadata):
|
||||
data_buf[0 : len(test_data)] = test_data
|
||||
|
||||
# Read it back
|
||||
with self.ring_buffer.access_buf(address) as (data_buf2, metadata2):
|
||||
read_data = bytes(data_buf2[0 : len(test_data)])
|
||||
read_id = metadata2[0]
|
||||
|
||||
self.assertEqual(read_data, test_data)
|
||||
self.assertEqual(read_id, monotonic_id)
|
||||
|
||||
def test_memory_error_on_full_buffer(self):
|
||||
"""Test that MemoryError is raised when buffer is full"""
|
||||
small_buffer_size = 200
|
||||
self.ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=small_buffer_size, create=True
|
||||
)
|
||||
|
||||
# Fill up the buffer
|
||||
self.ring_buffer.allocate_buf(100)
|
||||
self.ring_buffer.allocate_buf(80) # Total: 196 bytes used
|
||||
|
||||
# This should fail
|
||||
with self.assertRaises(MemoryError):
|
||||
self.ring_buffer.allocate_buf(1) # Would exceed buffer capacity
|
||||
|
||||
def test_allocation_and_free(self):
|
||||
"""Test allocation and freeing of buffers"""
|
||||
small_buffer_size = 200
|
||||
self.ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=small_buffer_size, create=True
|
||||
)
|
||||
|
||||
size = 80
|
||||
# Write some data
|
||||
test_data = b"Repeated test data"
|
||||
for i in range(5):
|
||||
address, monotonic_id = self.ring_buffer.allocate_buf(size)
|
||||
with self.ring_buffer.access_buf(address) as (data_buf, metadata):
|
||||
data_buf[0:4] = (0).to_bytes(4, "little") # 0 for not in-use
|
||||
data_buf[4 : len(test_data) + 4] = test_data
|
||||
print(self.ring_buffer.metadata)
|
||||
freed_ids = self.ring_buffer.free_buf(lambda *args: True)
|
||||
print(f" Freed IDs: {freed_ids}")
|
||||
self.assertEqual(freed_ids[0], i)
|
||||
|
||||
def test_clear_buffer(self):
|
||||
"""Test clearing the buffer"""
|
||||
self.ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=self.buffer_size, create=True
|
||||
)
|
||||
|
||||
# Allocate some buffers
|
||||
for _ in range(3):
|
||||
self.ring_buffer.allocate_buf(100)
|
||||
|
||||
# Clear the buffer
|
||||
self.ring_buffer.clear()
|
||||
|
||||
# Check that metadata is empty and IDs reset
|
||||
self.assertEqual(len(self.ring_buffer.metadata), 0)
|
||||
self.assertEqual(self.ring_buffer.monotonic_id_start, 0)
|
||||
self.assertEqual(self.ring_buffer.monotonic_id_end, 0)
|
||||
self.assertEqual(self.ring_buffer.data_buffer_start, 0)
|
||||
self.assertEqual(self.ring_buffer.data_buffer_end, 0)
|
||||
|
||||
def test_allocation_cycles(self):
|
||||
buffer_size = 100
|
||||
ring = SingleWriterShmRingBuffer(data_buffer_size=buffer_size, create=True)
|
||||
|
||||
# tracking allocations for assertions
|
||||
allocated_bitmap = np.zeros(
|
||||
(buffer_size,), dtype=np.bool_
|
||||
) # addr -> is_allocated
|
||||
allocation_map = dict() # monotonic_id -> (addr, size)
|
||||
|
||||
def count_allocated(bitmap) -> int:
|
||||
return np.sum(bitmap).item()
|
||||
|
||||
def is_free_fn(a, b) -> bool:
|
||||
return True
|
||||
|
||||
def mark_allocated_with_assertion(id, addr, size):
|
||||
addr = addr % buffer_size
|
||||
self.assertEqual(count_allocated(allocated_bitmap[addr : addr + size]), 0)
|
||||
|
||||
allocated_bitmap[addr : addr + size] = True
|
||||
allocation_map[id] = (addr, size)
|
||||
|
||||
def mark_freed_with_assertion(id):
|
||||
self.assertTrue(id in allocation_map)
|
||||
|
||||
addr, size = allocation_map.pop(id)
|
||||
addr = addr % buffer_size
|
||||
self.assertEqual(
|
||||
count_allocated(allocated_bitmap[addr : addr + size]), size
|
||||
)
|
||||
|
||||
allocated_bitmap[addr : addr + size] = False
|
||||
|
||||
def ring_free(free_size=None):
|
||||
freed_ids = ring.free_buf(is_free_fn, free_size)
|
||||
for freed_id in freed_ids:
|
||||
mark_freed_with_assertion(freed_id)
|
||||
|
||||
def ring_allocate(allocate_size):
|
||||
allocate_size_with_md = allocate_size + ring.MD_SIZE
|
||||
try:
|
||||
addr, monotonic_id = ring.allocate_buf(allocate_size)
|
||||
mark_allocated_with_assertion(monotonic_id, addr, allocate_size_with_md)
|
||||
except MemoryError:
|
||||
# free 2x size for enough space if wrapping happened
|
||||
ring_free(allocate_size_with_md * 2)
|
||||
|
||||
# retry allocating
|
||||
addr, monotonic_id = ring.allocate_buf(allocate_size)
|
||||
mark_allocated_with_assertion(monotonic_id, addr, allocate_size_with_md)
|
||||
|
||||
# 1. allocation & free cycles
|
||||
for _ in range(33):
|
||||
# will consume 2 + 8 = 10 bytes per allocation
|
||||
ring_allocate(2)
|
||||
|
||||
# 2. free all allocations
|
||||
ring_free()
|
||||
|
||||
# 3. try allocate the largest possible buffer
|
||||
ring_allocate(buffer_size - ring.MD_SIZE)
|
||||
|
||||
|
||||
def main():
|
||||
"""Main function demonstrating usage and running tests"""
|
||||
print("=== SingleWriterShmRingBuffer Test Suite ===\n")
|
||||
|
||||
# Run unit tests
|
||||
print("Running unit tests...")
|
||||
unittest.main(argv=[""], exit=False, verbosity=2)
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("=== Manual Demo ===\n")
|
||||
|
||||
# Manual demonstration
|
||||
try:
|
||||
print("Creating ring buffer...")
|
||||
writer_buffer = SingleWriterShmRingBuffer(data_buffer_size=2048, create=True)
|
||||
reader_buffer = SingleWriterShmRingBuffer(*writer_buffer.handle())
|
||||
|
||||
print(f"Buffer created with name: {writer_buffer.shared_memory.name}")
|
||||
|
||||
# Allocate some buffers
|
||||
print("\nAllocating buffers...")
|
||||
address_array = []
|
||||
for i in range(3):
|
||||
size = 100 + i * 50
|
||||
try:
|
||||
writer_buffer.free_buf(lambda *args: True)
|
||||
address, monotonic_id = writer_buffer.allocate_buf(size)
|
||||
address_array.append((address, size, monotonic_id))
|
||||
|
||||
# Write some test data
|
||||
with writer_buffer.access_buf(address) as (data_buf, metadata):
|
||||
test_message = f"Test message {i}".encode()
|
||||
data_buf[0 : len(test_message)] = test_message
|
||||
|
||||
except MemoryError as e:
|
||||
print(f" Failed to allocate {size} bytes: {e}")
|
||||
|
||||
print("\nBuffer state:")
|
||||
print(f" Data buffer start: {writer_buffer.data_buffer_start}")
|
||||
print(f" Data buffer end: {writer_buffer.data_buffer_end}")
|
||||
print(f" Monotonic ID start: {writer_buffer.monotonic_id_start}")
|
||||
print(f" Monotonic ID end: {writer_buffer.monotonic_id_end}")
|
||||
print(f" Metadata entries: {len(writer_buffer.metadata)}")
|
||||
|
||||
# Try to read back the data
|
||||
print("\nReading back data...")
|
||||
for address, size, monotonic_id in address_array:
|
||||
with reader_buffer.access_buf(address) as (data_buf, metadata):
|
||||
# Find null terminator or read first 50 chars
|
||||
data_bytes = bytes(data_buf[0:size])
|
||||
message = data_bytes.decode()
|
||||
print(f" ID {monotonic_id}: '{message}'")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Demo error: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
print("\n=== Demo Complete ===")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,325 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import multiprocessing
|
||||
import random
|
||||
import time
|
||||
import traceback
|
||||
import unittest
|
||||
from multiprocessing import Lock
|
||||
|
||||
import torch
|
||||
|
||||
# Assuming these are imported from your module
|
||||
from vllm.distributed.device_communicators.shm_object_storage import (
|
||||
MsgpackSerde,
|
||||
SingleWriterShmObjectStorage,
|
||||
SingleWriterShmRingBuffer,
|
||||
)
|
||||
from vllm.multimodal.inputs import (
|
||||
MultiModalFieldElem,
|
||||
MultiModalKwargsItem,
|
||||
MultiModalSharedField,
|
||||
)
|
||||
|
||||
|
||||
def _dummy_elem(size: int):
|
||||
return MultiModalFieldElem(
|
||||
data=torch.empty((size,), dtype=torch.int8),
|
||||
field=MultiModalSharedField(batch_size=1),
|
||||
)
|
||||
|
||||
|
||||
def _dummy_item(size_by_key: dict[str, int]):
|
||||
return MultiModalKwargsItem(
|
||||
{key: _dummy_elem(size) for key, size in size_by_key.items()}
|
||||
)
|
||||
|
||||
|
||||
class TestSingleWriterShmObjectStorage(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test fixtures before each test method."""
|
||||
ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=1024 * 100,
|
||||
create=True, # 10 MB buffer
|
||||
)
|
||||
self.storage = SingleWriterShmObjectStorage(
|
||||
max_object_size=1024 * 10, # 10KB max object
|
||||
n_readers=2,
|
||||
ring_buffer=ring_buffer,
|
||||
serde_class=MsgpackSerde,
|
||||
reader_lock=Lock(),
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after each test."""
|
||||
if self.storage:
|
||||
self.storage.close()
|
||||
|
||||
def test_minimal_put_get_cycle(self):
|
||||
"""Test basic put and get operations."""
|
||||
key = "test_key"
|
||||
value = _dummy_item({"field1": 10, "field2": 20})
|
||||
|
||||
# Put operation
|
||||
address, monotonic_id = self.storage.put(key, value)
|
||||
|
||||
# Verify key is in index
|
||||
self.assertIn(key, self.storage.key_index)
|
||||
self.assertEqual(self.storage.key_index[key], (address, monotonic_id))
|
||||
self.assertEqual(self.storage.id_index[monotonic_id], key)
|
||||
|
||||
# Get operation
|
||||
result = self.storage.get(address, monotonic_id)
|
||||
|
||||
# Verify result
|
||||
self.assertEqual(result, value)
|
||||
|
||||
def test_put_same_key_twice(self):
|
||||
"""Test behavior when putting the same key multiple times."""
|
||||
key = "duplicate_key"
|
||||
value1 = "first value"
|
||||
value2 = "second value"
|
||||
|
||||
# First put
|
||||
address1, id1 = self.storage.put(key, value1)
|
||||
retrieved1 = self.storage.get(address1, id1)
|
||||
self.assertEqual(retrieved1, value1)
|
||||
|
||||
# should raise an error on second put
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.storage.put(key, value2)
|
||||
|
||||
self.assertIn("already exists in the storage", str(context.exception))
|
||||
|
||||
def test_large_object_rejection(self):
|
||||
"""Test that objects exceeding max_object_size are rejected."""
|
||||
# Create an object larger than max_object_size
|
||||
large_data = "x" * (self.storage.max_object_size + 100)
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.storage.put("large_key", large_data)
|
||||
|
||||
self.assertIn("exceeds max object size", str(context.exception))
|
||||
|
||||
def test_buffer_overflow_and_cleanup(self):
|
||||
"""Test behavior when buffer fills up and needs cleanup."""
|
||||
# Fill up the buffer with many small objects
|
||||
stored_items = []
|
||||
|
||||
try:
|
||||
for i in range(1000): # Try to store many items
|
||||
key = f"item_{i}"
|
||||
value = f"data_{i}" * 100 # Make it reasonably sized
|
||||
address, monotonic_id = self.storage.put(key, value)
|
||||
stored_items.append((key, value, address, monotonic_id))
|
||||
except MemoryError:
|
||||
print(f"Buffer filled after {len(stored_items)} items")
|
||||
|
||||
# Verify that some items are still accessible
|
||||
accessible_count = 0
|
||||
for key, original_value, address, monotonic_id in stored_items:
|
||||
for i in range(self.storage.n_readers):
|
||||
retrieved = self.storage.get(address, monotonic_id)
|
||||
if retrieved == original_value:
|
||||
accessible_count += 1
|
||||
|
||||
self.assertEqual(accessible_count, len(stored_items))
|
||||
|
||||
try:
|
||||
for i in range(len(stored_items), 1000): # Try to store many items
|
||||
key = f"item_{i}"
|
||||
value = f"data_{i}" * 100 # Make it reasonably sized
|
||||
address, monotonic_id = self.storage.put(key, value)
|
||||
stored_items.append((key, value, address, monotonic_id))
|
||||
except MemoryError:
|
||||
print(f"Buffer filled after {len(stored_items)} items")
|
||||
|
||||
# Verify that some items are still accessibles
|
||||
for key, original_value, address, monotonic_id in stored_items:
|
||||
try:
|
||||
for i in range(self.storage.n_readers):
|
||||
retrieved = self.storage.get(address, monotonic_id)
|
||||
if retrieved == original_value:
|
||||
accessible_count += 1
|
||||
except ValueError as e:
|
||||
print(f"Error retrieving {key}: {e}")
|
||||
|
||||
# some items from the first batch may still be accessible
|
||||
self.assertGreaterEqual(accessible_count, len(stored_items))
|
||||
|
||||
def test_blocking_unread_object(self):
|
||||
"""Test behavior when buffer fills up and needs cleanup."""
|
||||
# Fill up the buffer with many small objects
|
||||
stored_items = []
|
||||
|
||||
try:
|
||||
for i in range(1000): # Try to store many items
|
||||
key = f"item_{i}"
|
||||
value = f"data_{i}" * 100 # Make it reasonably sized
|
||||
address, monotonic_id = self.storage.put(key, value)
|
||||
stored_items.append((key, value, address, monotonic_id))
|
||||
except MemoryError:
|
||||
print(f"Buffer filled after {len(stored_items)} items")
|
||||
|
||||
# read all items except the first one
|
||||
# to simulate a blocking situation
|
||||
accessible_count = 0
|
||||
for key, original_value, address, monotonic_id in stored_items[1:]:
|
||||
for i in range(self.storage.n_readers):
|
||||
retrieved = self.storage.get(address, monotonic_id)
|
||||
if retrieved == original_value:
|
||||
accessible_count += 1
|
||||
|
||||
self.assertEqual(accessible_count, len(stored_items) - 1)
|
||||
|
||||
try:
|
||||
key = f"item_{len(stored_items)}"
|
||||
value = f"data_{len(stored_items)}" * 100
|
||||
address, monotonic_id = self.storage.put(key, value)
|
||||
except MemoryError:
|
||||
print(f"Buffer filled after {len(stored_items)} items")
|
||||
|
||||
# read the first item
|
||||
for i in range(self.storage.n_readers):
|
||||
key, original_value, address, monotonic_id = stored_items[0]
|
||||
retrieved = self.storage.get(address, monotonic_id)
|
||||
self.assertEqual(retrieved, original_value)
|
||||
|
||||
try:
|
||||
for i in range(len(stored_items), 1000): # Try to store many items
|
||||
key = f"item_{i}"
|
||||
value = f"data_{i}" * 100 # Make it reasonably sized
|
||||
address, monotonic_id = self.storage.put(key, value)
|
||||
stored_items.append((key, value, address, monotonic_id))
|
||||
except MemoryError:
|
||||
print(f"Buffer filled after {len(stored_items)} items")
|
||||
|
||||
# some items from the first batch may still be accessible
|
||||
self.assertGreaterEqual(len(stored_items), accessible_count + 10)
|
||||
|
||||
def test_invalid_get_operations(self):
|
||||
"""Test various invalid get operations."""
|
||||
# Test with non-existent address
|
||||
with self.assertRaises(ValueError): # Could be various exceptions
|
||||
self.storage.get(99999, 1)
|
||||
|
||||
# Store something first
|
||||
address, monotonic_id = self.storage.put("test", "value")
|
||||
|
||||
# Test with wrong monotonic_id
|
||||
with self.assertRaises(ValueError) as context:
|
||||
self.storage.get(address, monotonic_id + 100)
|
||||
|
||||
self.assertIn("has been modified or is invalid", str(context.exception))
|
||||
|
||||
def test_clear_storage(self):
|
||||
"""Test clearing the storage."""
|
||||
# Store some items
|
||||
for i in range(5):
|
||||
self.storage.put(f"item_{i}", f"value_{i}")
|
||||
|
||||
# Clear the storage
|
||||
self.storage.clear()
|
||||
|
||||
# Verify that all indices are empty
|
||||
self.assertEqual(len(self.storage.key_index), 0)
|
||||
self.assertEqual(len(self.storage.id_index), 0)
|
||||
self.assertEqual(len(self.storage.ring_buffer.metadata), 0)
|
||||
|
||||
# Verify that new items can be added after clearing
|
||||
address, monotonic_id = self.storage.put("new_item", "new_value")
|
||||
self.assertIn("new_item", self.storage.key_index)
|
||||
self.assertEqual((address, monotonic_id), (0, 0))
|
||||
|
||||
|
||||
# Reader process function
|
||||
def reader_process(process_id, storage_handle, items_to_read):
|
||||
"""Reader process that connects to existing shared memory and reads data."""
|
||||
reader_storage = SingleWriterShmObjectStorage.create_from_handle(storage_handle)
|
||||
|
||||
print(f"Reader {process_id} started")
|
||||
|
||||
errors = []
|
||||
|
||||
for key, original_value, address, monotonic_id in items_to_read:
|
||||
time.sleep(random.random() / 100)
|
||||
try:
|
||||
# Read data from shared memory
|
||||
retrieved_value = reader_storage.get(address, monotonic_id)
|
||||
|
||||
# Verify data integrity
|
||||
assert retrieved_value == original_value
|
||||
print(f"Reader {process_id} retrieved {key}: {retrieved_value}")
|
||||
except Exception as e:
|
||||
errors.append((key, str(e), type(e).__name__))
|
||||
|
||||
|
||||
def run_multiprocess_example():
|
||||
"""Run a minimal working example with real shared memory."""
|
||||
print("=== Minimal Object Storage Example ===")
|
||||
|
||||
try:
|
||||
# Create storage instance
|
||||
ring_buffer = SingleWriterShmRingBuffer(
|
||||
data_buffer_size=1024 * 100,
|
||||
create=True, # 10 MB buffer
|
||||
)
|
||||
storage = SingleWriterShmObjectStorage(
|
||||
max_object_size=1024,
|
||||
n_readers=3,
|
||||
ring_buffer=ring_buffer,
|
||||
serde_class=MsgpackSerde,
|
||||
reader_lock=Lock(),
|
||||
)
|
||||
|
||||
print(f"Created storage (writer: {storage.is_writer})")
|
||||
|
||||
# Test basic data types
|
||||
test_data = [
|
||||
("user_data", {"name": "Alice", "age": 30, "scores": [95, 87, 92]}),
|
||||
("simple_string", "Hello, World!"),
|
||||
("number", 42),
|
||||
("list_data", [1, 2, 3, "four", 5.0]),
|
||||
]
|
||||
|
||||
stored_items = []
|
||||
|
||||
# Store all data
|
||||
for key, value in test_data:
|
||||
print(f"Storing {key}: {value}")
|
||||
address, monotonic_id = storage.put(key, value)
|
||||
stored_items.append((key, value, address, monotonic_id))
|
||||
print(f" -> Stored at address {address}, ID {monotonic_id}")
|
||||
|
||||
print("\n--- Retrieving Data ---")
|
||||
processes = []
|
||||
handle = storage.handle()
|
||||
# initialize lock for reader processes
|
||||
handle.reader_lock = Lock()
|
||||
for i in range(storage.n_readers):
|
||||
p = multiprocessing.Process(
|
||||
target=reader_process, args=(i, handle, stored_items)
|
||||
)
|
||||
processes.append(p)
|
||||
p.start()
|
||||
|
||||
for p in processes:
|
||||
p.join(timeout=10)
|
||||
if p.is_alive():
|
||||
p.terminate()
|
||||
p.join()
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in minimal example: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the minimal example first
|
||||
run_multiprocess_example()
|
||||
print("\n" + "=" * 50 + "\n")
|
||||
|
||||
# Run the test suite
|
||||
print("Running comprehensive test suite...")
|
||||
unittest.main(verbosity=2, exit=False)
|
||||
@@ -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)
|
||||
@@ -0,0 +1,140 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import queue
|
||||
import random
|
||||
import typing
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
|
||||
from vllm.distributed.device_communicators.cuda_communicator import CudaCommunicator
|
||||
from vllm.distributed.parallel_state import (
|
||||
get_tp_group,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from vllm.engine.arg_utils import EngineArgs
|
||||
from vllm.engine.llm_engine import LLMEngine
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
torch.manual_seed(42)
|
||||
random.seed(44)
|
||||
|
||||
test_size_elements = 1024 * 1024
|
||||
|
||||
|
||||
def symm_mem_allreduce_worker(local_rank: int, world_size: int, q: mp.Queue):
|
||||
monkeypatch = pytest.MonkeyPatch()
|
||||
config = VllmConfig(parallel_config=ParallelConfig(tensor_parallel_size=world_size))
|
||||
|
||||
with monkeypatch.context() as m, set_current_vllm_config(config):
|
||||
m.delenv("CUDA_VISIBLE_DEVICES", raising=False)
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.accelerator.set_device_index(device)
|
||||
torch.set_default_device(device)
|
||||
torch.set_default_dtype(dtype)
|
||||
update_environment_variables(
|
||||
{
|
||||
"RANK": str(local_rank),
|
||||
"LOCAL_RANK": str(local_rank),
|
||||
"WORLD_SIZE": str(world_size),
|
||||
"MASTER_ADDR": "localhost",
|
||||
"MASTER_PORT": "12345",
|
||||
}
|
||||
)
|
||||
|
||||
init_distributed_environment()
|
||||
initialize_model_parallel(tensor_model_parallel_size=world_size)
|
||||
|
||||
cuda_communicator = typing.cast(
|
||||
CudaCommunicator, get_tp_group().device_communicator
|
||||
)
|
||||
symm_mem_comm = cuda_communicator.symm_mem_comm
|
||||
if symm_mem_comm is None or symm_mem_comm.disabled:
|
||||
# can't use skip under multiprocessing
|
||||
q.put("SymmMemCommunicator is not available or disabled.")
|
||||
return
|
||||
|
||||
inp_direct_symm_mem = torch.randint(
|
||||
1, 23, (test_size_elements,), dtype=dtype, device=device
|
||||
)
|
||||
if not symm_mem_comm.should_use_symm_mem(inp_direct_symm_mem):
|
||||
# can't use skip under multiprocessing
|
||||
q.put("SymmMemCommunicator isn't used for this world and input size.")
|
||||
return
|
||||
|
||||
original_inp_direct_symm_mem = inp_direct_symm_mem.clone()
|
||||
out_direct_symm_mem = symm_mem_comm.all_reduce(inp_direct_symm_mem)
|
||||
assert out_direct_symm_mem is not None
|
||||
|
||||
group = get_tp_group().device_group
|
||||
dist.all_reduce(original_inp_direct_symm_mem, group=group)
|
||||
torch.testing.assert_close(
|
||||
out_direct_symm_mem, original_inp_direct_symm_mem, atol=2.5, rtol=0.1
|
||||
)
|
||||
|
||||
# Test tensor_model_parallel_all_reduce which should use symm_mem
|
||||
inp_tensor_parallel = torch.randint(
|
||||
-23, 1, (test_size_elements,), dtype=dtype, device=device
|
||||
)
|
||||
original_inp_tensor_parallel = inp_tensor_parallel.clone()
|
||||
out_tensor_parallel = tensor_model_parallel_all_reduce(inp_tensor_parallel)
|
||||
dist.all_reduce(original_inp_tensor_parallel, group=group)
|
||||
torch.testing.assert_close(
|
||||
out_tensor_parallel, original_inp_tensor_parallel, atol=2.5, rtol=0.1
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="SymmMemAllreduce is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.parametrize("tp_size", [2])
|
||||
@pytest.mark.parametrize("pipeline_parallel_size", [1])
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_symm_mem_allreduce(
|
||||
monkeypatch: pytest.MonkeyPatch, tp_size, pipeline_parallel_size
|
||||
):
|
||||
world_size = tp_size * pipeline_parallel_size
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
q = mp.get_context("spawn").Queue()
|
||||
mp.spawn(symm_mem_allreduce_worker, args=(world_size, q), nprocs=world_size)
|
||||
try:
|
||||
val = q.get(timeout=1)
|
||||
except queue.Empty:
|
||||
val = None
|
||||
finally:
|
||||
cleanup_dist_env_and_memory()
|
||||
if val is not None:
|
||||
pytest.skip(val)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(),
|
||||
reason="SymmMemAllreduce is only available for CUDA platforms.",
|
||||
)
|
||||
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
|
||||
def test_dp_with_symm_mem_allreduce(monkeypatch: pytest.MonkeyPatch):
|
||||
world_size = 4
|
||||
if world_size > torch.accelerator.device_count():
|
||||
pytest.skip("Not enough GPUs to run the test.")
|
||||
# Verify that the DataParallel runs without error
|
||||
engine_args = EngineArgs(
|
||||
model="distilbert/distilgpt2",
|
||||
enforce_eager=True,
|
||||
enable_prefix_caching=True,
|
||||
data_parallel_size=2,
|
||||
tensor_parallel_size=2,
|
||||
data_parallel_backend="mp",
|
||||
)
|
||||
LLMEngine.from_engine_args(engine_args)
|
||||
@@ -0,0 +1,82 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# unit test for `examples/features/torchrun/torchrun_example_offline.py`
|
||||
import os
|
||||
import random
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import get_world_group
|
||||
|
||||
# By default, let PyTorch choose the WORLD backend for the current device
|
||||
# type (legacy lazy-init path). When VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1,
|
||||
# use the explicit eager-init pattern required by `split_group` (mixed
|
||||
# cpu:gloo,cuda:nccl backend + device_id binding).
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.accelerator.set_device_index(local_rank)
|
||||
dist.init_process_group(
|
||||
backend="cpu:gloo,cuda:nccl",
|
||||
device_id=torch.device(f"cuda:{local_rank}"),
|
||||
)
|
||||
else:
|
||||
dist.init_process_group()
|
||||
|
||||
# Create prompts
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
# set different `gpu_memory_utilization` for different ranks,
|
||||
# to test if all ranks agree on the same kv cache configuration.
|
||||
llm = LLM(
|
||||
model="facebook/opt-125m",
|
||||
tensor_parallel_size=2,
|
||||
pipeline_parallel_size=int(os.getenv("PP_SIZE", 1)),
|
||||
distributed_executor_backend="external_launcher",
|
||||
gpu_memory_utilization=random.uniform(0.8, 0.92),
|
||||
seed=0,
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
cpu_group = get_world_group().cpu_group
|
||||
|
||||
torch_rank = dist.get_rank(group=cpu_group)
|
||||
|
||||
|
||||
def test_consistent_across_ranks(obj):
|
||||
if torch_rank == 0:
|
||||
dist.broadcast_object_list([obj], src=0, group=cpu_group)
|
||||
else:
|
||||
container = [None]
|
||||
dist.broadcast_object_list(container, src=0, group=cpu_group)
|
||||
assert container[0] == obj
|
||||
|
||||
|
||||
test_consistent_across_ranks(llm.llm_engine.vllm_config.cache_config.num_cpu_blocks)
|
||||
test_consistent_across_ranks(llm.llm_engine.vllm_config.cache_config.num_gpu_blocks)
|
||||
|
||||
# make sure we can access the model parameters from the calling process
|
||||
# of the `LLM` instance.
|
||||
params = list(
|
||||
llm.llm_engine.model_executor.driver_worker.worker.model_runner.model.parameters()
|
||||
)
|
||||
test_consistent_across_ranks(len(params))
|
||||
|
||||
# all ranks should have the same outputs
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
test_consistent_across_ranks(prompt)
|
||||
test_consistent_across_ranks(generated_text)
|
||||
print(f"Rank {torch_rank}, Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
@@ -0,0 +1,91 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# unit test for `examples/features/torchrun/torchrun_example_offline.py`
|
||||
import os
|
||||
import random
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import LLM, SamplingParams
|
||||
from vllm.distributed.parallel_state import get_tp_group, get_world_group
|
||||
|
||||
# By default, let PyTorch choose the WORLD backend for the current device
|
||||
# type (legacy lazy-init path). When VLLM_DISTRIBUTED_USE_SPLIT_GROUP=1,
|
||||
# use the explicit eager-init pattern required by `split_group` (mixed
|
||||
# cpu:gloo,cuda:nccl backend + device_id binding).
|
||||
if envs.VLLM_DISTRIBUTED_USE_SPLIT_GROUP:
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
torch.accelerator.set_device_index(local_rank)
|
||||
dist.init_process_group(
|
||||
backend="cpu:gloo,cuda:nccl",
|
||||
device_id=torch.device(f"cuda:{local_rank}"),
|
||||
)
|
||||
else:
|
||||
dist.init_process_group()
|
||||
|
||||
# Create prompts
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
] * 10
|
||||
dp_size = int(os.getenv("DP_SIZE", "1"))
|
||||
dp_rank = int(os.getenv("DP_RANK", "0"))
|
||||
|
||||
if dp_size > 1:
|
||||
# distribute the prompts across the data parallel ranks
|
||||
prompts = [prompt for idx, prompt in enumerate(prompts) if idx % dp_size == dp_rank]
|
||||
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||||
|
||||
# set different `gpu_memory_utilization` for different ranks,
|
||||
# to test if all ranks agree on the same kv cache configuration.
|
||||
llm = LLM(
|
||||
model="microsoft/Phi-mini-MoE-instruct",
|
||||
tensor_parallel_size=int(os.getenv("TP_SIZE", "1")),
|
||||
pipeline_parallel_size=int(os.getenv("PP_SIZE", "1")),
|
||||
enable_expert_parallel=int(os.getenv("ENABLE_EP", "0")) == 1,
|
||||
distributed_executor_backend="external_launcher",
|
||||
gpu_memory_utilization=random.uniform(0.8, 0.92),
|
||||
seed=0,
|
||||
max_model_len=1024,
|
||||
max_num_seqs=16,
|
||||
)
|
||||
|
||||
outputs = llm.generate(prompts, sampling_params)
|
||||
|
||||
group = get_world_group() if dp_size == 1 else get_tp_group()
|
||||
cpu_group = group.cpu_group
|
||||
group_rank = dist.get_rank(group=cpu_group)
|
||||
|
||||
|
||||
def test_consistent_across_ranks(obj):
|
||||
if group_rank == 0:
|
||||
dist.broadcast_object_list([obj], src=group.ranks[0], group=cpu_group)
|
||||
else:
|
||||
container = [None]
|
||||
dist.broadcast_object_list(container, src=group.ranks[0], group=cpu_group)
|
||||
assert container[0] == obj
|
||||
|
||||
|
||||
test_consistent_across_ranks(llm.llm_engine.vllm_config.cache_config.num_cpu_blocks)
|
||||
test_consistent_across_ranks(llm.llm_engine.vllm_config.cache_config.num_gpu_blocks)
|
||||
|
||||
# make sure we can access the model parameters from the calling process
|
||||
# of the `LLM` instance.
|
||||
params = list(
|
||||
llm.llm_engine.model_executor.driver_worker.worker.model_runner.model.parameters()
|
||||
)
|
||||
test_consistent_across_ranks(len(params))
|
||||
|
||||
# all ranks should have the same outputs
|
||||
for output in outputs:
|
||||
prompt = output.prompt
|
||||
generated_text = output.outputs[0].text
|
||||
test_consistent_across_ranks(prompt)
|
||||
test_consistent_across_ranks(generated_text)
|
||||
print(f"Rank {group_rank}, Prompt: {prompt!r}, Generated text: {generated_text!r}")
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import socket
|
||||
|
||||
import pytest
|
||||
import ray
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
|
||||
from vllm.distributed.utils import StatelessProcessGroup
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
from vllm.utils.system_utils import update_environment_variables
|
||||
|
||||
from ..utils import multi_gpu_test
|
||||
|
||||
|
||||
@ray.remote
|
||||
class _CUDADeviceCountStatelessTestActor:
|
||||
def get_count(self):
|
||||
return current_platform.device_count()
|
||||
|
||||
def set_cuda_visible_devices(self, cuda_visible_devices: str):
|
||||
update_environment_variables({"CUDA_VISIBLE_DEVICES": cuda_visible_devices})
|
||||
|
||||
def get_cuda_visible_devices(self):
|
||||
return envs.CUDA_VISIBLE_DEVICES
|
||||
|
||||
|
||||
def test_cuda_device_count_stateless():
|
||||
"""Test that cuda_device_count_stateless changes return value if
|
||||
CUDA_VISIBLE_DEVICES is changed."""
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip("Skip for ROCm because Ray uses HIP_VISIBLE_DEVICES.")
|
||||
actor = _CUDADeviceCountStatelessTestActor.options( # type: ignore
|
||||
num_gpus=2
|
||||
).remote()
|
||||
assert len(sorted(ray.get(actor.get_cuda_visible_devices.remote()).split(","))) == 2
|
||||
assert ray.get(actor.get_count.remote()) == 2
|
||||
ray.get(actor.set_cuda_visible_devices.remote("0"))
|
||||
assert ray.get(actor.get_count.remote()) == 1
|
||||
ray.get(actor.set_cuda_visible_devices.remote(""))
|
||||
assert ray.get(actor.get_count.remote()) == 0
|
||||
|
||||
|
||||
def cpu_worker(rank, WORLD_SIZE, port1, port2):
|
||||
pg1 = StatelessProcessGroup.create(
|
||||
host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
|
||||
)
|
||||
if rank <= 2:
|
||||
pg2 = StatelessProcessGroup.create(
|
||||
host="127.0.0.1", port=port2, rank=rank, world_size=3
|
||||
)
|
||||
data = torch.tensor([rank])
|
||||
data = pg1.broadcast_obj(data, src=2)
|
||||
assert data.item() == 2
|
||||
if rank <= 2:
|
||||
data = torch.tensor([rank + 1])
|
||||
data = pg2.broadcast_obj(data, src=2)
|
||||
assert data.item() == 3
|
||||
pg2.barrier()
|
||||
pg1.barrier()
|
||||
|
||||
|
||||
def gpu_worker(rank, WORLD_SIZE, port1, port2):
|
||||
torch.accelerator.set_device_index(rank)
|
||||
pg1 = StatelessProcessGroup.create(
|
||||
host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
|
||||
)
|
||||
pynccl1 = PyNcclCommunicator(pg1, device=rank)
|
||||
if rank <= 2:
|
||||
pg2 = StatelessProcessGroup.create(
|
||||
host="127.0.0.1", port=port2, rank=rank, world_size=3
|
||||
)
|
||||
pynccl2 = PyNcclCommunicator(pg2, device=rank)
|
||||
data = torch.tensor([rank]).cuda()
|
||||
pynccl1.all_reduce(data)
|
||||
pg1.barrier()
|
||||
torch.accelerator.synchronize()
|
||||
if rank <= 2:
|
||||
pynccl2.all_reduce(data)
|
||||
pg2.barrier()
|
||||
torch.accelerator.synchronize()
|
||||
item = data[0].item()
|
||||
print(f"rank: {rank}, item: {item}")
|
||||
if rank == 3:
|
||||
assert item == 6
|
||||
else:
|
||||
assert item == 18
|
||||
|
||||
|
||||
def broadcast_worker(rank, WORLD_SIZE, port1, port2):
|
||||
pg1 = StatelessProcessGroup.create(
|
||||
host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
|
||||
)
|
||||
if rank == 2:
|
||||
pg1.broadcast_obj("secret", src=2)
|
||||
else:
|
||||
obj = pg1.broadcast_obj(None, src=2)
|
||||
assert obj == "secret"
|
||||
pg1.barrier()
|
||||
|
||||
|
||||
def allgather_worker(rank, WORLD_SIZE, port1, port2):
|
||||
pg1 = StatelessProcessGroup.create(
|
||||
host="127.0.0.1", port=port1, rank=rank, world_size=WORLD_SIZE
|
||||
)
|
||||
data = pg1.all_gather_obj(rank)
|
||||
assert data == list(range(WORLD_SIZE))
|
||||
pg1.barrier()
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="This test is flaky and prone to hang.")
|
||||
@multi_gpu_test(num_gpus=4)
|
||||
@pytest.mark.parametrize(
|
||||
"worker", [cpu_worker, gpu_worker, broadcast_worker, allgather_worker]
|
||||
)
|
||||
def test_stateless_process_group(worker):
|
||||
port1 = get_open_port()
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", port1))
|
||||
port2 = get_open_port()
|
||||
WORLD_SIZE = 4
|
||||
from multiprocessing import get_context
|
||||
|
||||
ctx = get_context("fork")
|
||||
processes = []
|
||||
for i in range(WORLD_SIZE):
|
||||
rank = i
|
||||
processes.append(
|
||||
ctx.Process(target=worker, args=(rank, WORLD_SIZE, port1, port2))
|
||||
)
|
||||
for p in processes:
|
||||
p.start()
|
||||
for p in processes:
|
||||
p.join()
|
||||
for p in processes:
|
||||
assert not p.exitcode
|
||||
print("All processes finished.")
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user