wehub-resource-sync
135 lines
4.6 KiB
Python
135 lines
4.6 KiB
Python
# 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)
|