diff --git a/examples/offline_inference/rlhf_colocate.py b/examples/offline_inference/rlhf_colocate.py index 360fd79b5..241aa0ad8 100644 --- a/examples/offline_inference/rlhf_colocate.py +++ b/examples/offline_inference/rlhf_colocate.py @@ -30,6 +30,7 @@ https://docs.ray.io/en/latest/placement-groups.html import gc import os +import sys import ray import torch @@ -40,6 +41,10 @@ from torch.multiprocessing.reductions import reduce_tensor from vllm import LLM +if torch.version.hip is not None: + print("Skipping test for ROCm. Ray is unsupported on vLLM ROCm.") + sys.exit(0) + class MyLLM(LLM): """Configure the vLLM worker for Ray placement group execution. diff --git a/tests/distributed/test_utils.py b/tests/distributed/test_utils.py index 8289f697f..526b6749d 100644 --- a/tests/distributed/test_utils.py +++ b/tests/distributed/test_utils.py @@ -10,6 +10,7 @@ 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 vllm.utils.torch_utils import cuda_device_count_stateless @@ -32,6 +33,8 @@ class _CUDADeviceCountStatelessTestActor: 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()