A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 22:45 UTC
[rocm][ray] Fix: Unify Ray device visibility handling across CUDA and ROCm (#33308)
Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>
kourosh hakhamaneshi committed
Feb 4, 2026 at 10:09 UTC
2f6d17cb2f4a49e29aae5c3c1ff64623d66d0257
8 files changed
+43
-8
docker/Dockerfile.rocm
-2
@@ -15,8 +15,6 @@ FROM ${BASE_IMAGE} AS base
15
16
ARG ARG_PYTORCH_ROCM_ARCH
17
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
18
-ENV RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1
19
-ENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1
18
19
# Install some basic utilities
20
RUN apt-get update -q -y && apt-get install -q -y \
tests/config/test_config_generation.py
+5
@@ -60,6 +60,11 @@ def test_ray_runtime_env(monkeypatch: pytest.MonkeyPatch):
60
runtime_env = {
61
"env_vars": {
62
"TEST_ENV_VAR": "test_value",
63
+ # In future ray versions, this will be default, so when setting a
64
+ # task or actor with num_gpus=None/0, the visible devices env var
65
+ # won't be overridden resulting in no GPUs being visible on a gpu
66
+ # machine.
67
+ "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO": "0",
68
},
69
}
70
vllm/platforms/cuda.py
+3
@@ -102,6 +102,9 @@ class CudaPlatformBase(Platform):
102
ray_device_key: str = "GPU"
103
dist_backend: str = "nccl"
104
device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
105
+ ray_noset_device_env_vars: list[str] = [
106
+ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
107
+ ]
108
109
@property
110
def supported_dtypes(self) -> list[torch.dtype]:
vllm/platforms/interface.py
+5
@@ -116,6 +116,11 @@ class Platform:
116
# https://github.com/ray-project/ray/tree/master/python/ray/_private/accelerators # noqa
117
device_control_env_var: str = "VLLM_DEVICE_CONTROL_ENV_VAR_PLACEHOLDER"
118
119
+ # environment variables that need to be set to 1 to prevent ray from
120
+ # setting the visible devices e.g.
121
+ # RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES
122
+ ray_noset_device_env_vars: list[str] = []
123
+
124
# The torch.compile backend for compiling simple and
125
# standalone functions. The default value is "inductor" to keep
126
# the same behavior as PyTorch.
vllm/platforms/rocm.py
+5
@@ -194,6 +194,11 @@ class RocmPlatform(Platform):
194
dist_backend: str = "nccl"
195
# rocm shares the same device control env var as CUDA
196
device_control_env_var: str = "CUDA_VISIBLE_DEVICES"
197
+ ray_noset_device_env_vars: list[str] = [
198
+ "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES",
199
+ "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES",
200
+ "RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES",
201
+ ]
202
203
supported_quantization: list[str] = [
204
"awq",
vllm/v1/executor/ray_executor.py
+24
@@ -69,6 +69,8 @@ class RayDistributedExecutor(Executor):
69
"VLLM_HOST_PORT",
70
"LOCAL_RANK",
71
"CUDA_VISIBLE_DEVICES",
72
+ "HIP_VISIBLE_DEVICES",
73
+ "ROCR_VISIBLE_DEVICES",
74
}
75
76
# These non-vLLM env vars are copied from the driver to workers
@@ -146,6 +148,14 @@ class RayDistributedExecutor(Executor):
148
149
return ray_remote_kwargs
150
151
+ def _update_noset_device_env_vars(self, ray_remote_kwargs):
152
+ runtime_env = ray_remote_kwargs.setdefault("runtime_env", {})
153
+ env_vars = runtime_env.setdefault("env_vars", {})
154
+ env_vars.update(
155
+ {env_var: "1" for env_var in current_platform.ray_noset_device_env_vars}
156
+ )
157
+ return ray_remote_kwargs
158
+
159
# child class could overwrite this to return actual env vars.
160
def _get_env_vars_to_be_updated(self):
161
return self._env_vars_for_all_workers
@@ -169,6 +179,11 @@ class RayDistributedExecutor(Executor):
179
ray_remote_kwargs
180
)
181
182
+ # The way ray actors are setup in vllm is that the visible devices are
183
+ # not set by actors, they are left unset by ray. Internally we index
184
+ # the right gpu with local_rank. This is similar to how mp mode works.
185
+ self._update_noset_device_env_vars(ray_remote_kwargs)
186
+
187
# Create the workers.
188
bundle_indices: list[int]
189
if envs.VLLM_RAY_BUNDLE_INDICES:
@@ -303,6 +318,15 @@ class RayDistributedExecutor(Executor):
318
)
319
320
# Set environment variables for the driver and workers.
321
+ # We set CUDA_VISIBLE_DEVICES to ALL GPUs on the node for each worker.
322
+ # This is needed because:
323
+ # 1. Ray's compiled DAG needs to find the allocated GPU in
324
+ # CUDA_VISIBLE_DEVICES.
325
+ # 2. vLLM's communication layer (NCCL, CustomAllreduce) needs to see
326
+ # all GPUs for P2P checks and communication setup. Though if it was
327
+ # just this reason, we could have also just kept the visible devices
328
+ # unset.
329
+ # Each worker will use local_rank to index into the visible devices.
330
all_args_to_update_environment_variables = [
331
{
332
current_platform.device_control_env_var: ",".join(
vllm/v1/worker/gpu_worker.py
+1
@@ -209,6 +209,7 @@ class Worker(WorkerBase):
209
f"be less than or equal to the number of visible devices "
210
f"({visible_device_count})."
211
)
212
+
213
self.device = torch.device(f"cuda:{self.local_rank}")
214
current_platform.set_device(self.device)
215
vllm/v1/worker/worker_base.py
-6
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import os
4
from collections.abc import Callable
5
from typing import TYPE_CHECKING, Any, TypeVar
6
@@ -221,11 +220,6 @@ class WorkerWrapperBase:
220
envs_list: list[dict[str, str]],
221
) -> None:
222
envs = envs_list[self.rpc_rank]
224
- key = "CUDA_VISIBLE_DEVICES"
225
- if key in envs and key in os.environ:
226
- # overwriting CUDA_VISIBLE_DEVICES is desired behavior
227
- # suppress the warning in `update_environment_variables`
228
- del os.environ[key]
223
update_environment_variables(envs)
224
225
def init_worker(self, all_kwargs: list[dict[str, Any]]) -> None: