forked from Karylab-cklius/vllm
Stop setting CUDA_VISIBLE_DEVICES internally in vLLM, add device_ids arg (#45026)
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: kourosh hakhamaneshi <kouroshHakha@users.noreply.github.com>
This commit is contained in:
co-authored by
Claude
Codex
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
kourosh hakhamaneshi
parent
e9de72fe6c
commit
ebfbcfe46a
@@ -268,9 +268,14 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
|
||||
using TileShapeD = typename MlaSm100Type::TileShapeD;
|
||||
arguments.problem_shape =
|
||||
cute::make_tuple(TileShapeH{}, static_cast<int>(max_seq_len), TileShapeD{}, static_cast<int>(num_batches));
|
||||
// Assumes device 0 when getting sm_count.
|
||||
arguments.hw_info.sm_count =
|
||||
sm_count <= 0 ? cutlass::KernelHardwareInfo::query_device_multiprocessor_count(/*device_id=*/0) : sm_count;
|
||||
if (sm_count <= 0) {
|
||||
int current_device = 0;
|
||||
cudaGetDevice(¤t_device);
|
||||
arguments.hw_info.sm_count =
|
||||
cutlass::KernelHardwareInfo::query_device_multiprocessor_count(current_device);
|
||||
} else {
|
||||
arguments.hw_info.sm_count = sm_count;
|
||||
}
|
||||
arguments.split_kv = static_cast<int>(num_kv_splits);
|
||||
MlaSm100Type::Fmha::set_split_kv(arguments);
|
||||
|
||||
|
||||
@@ -649,3 +649,196 @@ def test_cloud_storage_tokenizer_skips_get_model_path(monkeypatch):
|
||||
args = EngineArgs(model="s3://bucket/model", tokenizer="s3://bucket/tokenizer")
|
||||
assert args.model == "s3://bucket/model"
|
||||
assert args.tokenizer == "s3://bucket/tokenizer"
|
||||
|
||||
|
||||
class TestDeviceIds:
|
||||
def test_device_ids_with_cvd_out_of_range(self, monkeypatch):
|
||||
"""--device-ids index beyond the CVD set raises ValueError."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
key = current_platform.device_control_env_var
|
||||
monkeypatch.setenv(key, "4,5")
|
||||
args = EngineArgs(model="m", device_ids=[0, 2])
|
||||
with pytest.raises(ValueError, match="out of range"):
|
||||
args._resolve_device_ids()
|
||||
|
||||
def test_device_ids_with_cvd_resolve_to_physical_ids(self, monkeypatch):
|
||||
"""--device-ids are CVD-local indices resolved to physical ids."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
key = current_platform.device_control_env_var
|
||||
monkeypatch.setenv(key, "4,5")
|
||||
args = EngineArgs(model="m", device_ids=[0, 1])
|
||||
assert args._resolve_device_ids() == [4, 5]
|
||||
|
||||
def test_device_ids_with_uuid_cvd_resolve_to_physical_ids(self, monkeypatch):
|
||||
"""--device-ids support UUID CVD values resolved by the platform."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
key = current_platform.device_control_env_var
|
||||
monkeypatch.setenv(key, "GPU-abcd1234,GPU-ef567890")
|
||||
monkeypatch.setattr(
|
||||
type(current_platform),
|
||||
"device_control_id_to_physical_device_id",
|
||||
classmethod(
|
||||
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
|
||||
),
|
||||
)
|
||||
|
||||
args = EngineArgs(model="m", device_ids=[0, 1])
|
||||
assert args._resolve_device_ids() == [4, 5]
|
||||
|
||||
def test_device_ids_with_uuid_args_resolve_to_physical_ids(self, monkeypatch):
|
||||
"""UUID --device-ids are resolved to physical IDs immediately."""
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(
|
||||
type(current_platform),
|
||||
"device_control_id_to_physical_device_id",
|
||||
classmethod(lambda cls, device_id: {"GPU-abcd1234": 4}[device_id]),
|
||||
)
|
||||
|
||||
args = EngineArgs(model="m", device_ids=["GPU-abcd1234"])
|
||||
assert args._resolve_device_ids() == [4]
|
||||
|
||||
def test_device_ids_reject_mixed_integer_and_uuid_args(self):
|
||||
"""--device-ids must not mix CVD indices and UUIDs."""
|
||||
args = EngineArgs(model="m", device_ids=[0, "GPU-abcd1234"])
|
||||
with pytest.raises(ValueError, match="must not mix"):
|
||||
args._resolve_device_ids()
|
||||
|
||||
def test_no_device_ids(self):
|
||||
"""No --device-ids returns None."""
|
||||
args = EngineArgs(model="m")
|
||||
assert args._resolve_device_ids() is None
|
||||
|
||||
def test_cli_parsing(self):
|
||||
"""--device-ids parses comma-separated string from CLI."""
|
||||
parser = FlexibleArgumentParser()
|
||||
EngineArgs.add_cli_args(parser)
|
||||
parsed = parser.parse_args(["--model", "m", "--device-ids", "0,2,4"])
|
||||
assert parsed.device_ids == [0, 2, 4]
|
||||
|
||||
def test_cli_parsing_uuid(self):
|
||||
"""--device-ids parses comma-separated UUID strings from CLI."""
|
||||
parser = FlexibleArgumentParser()
|
||||
EngineArgs.add_cli_args(parser)
|
||||
parsed = parser.parse_args(
|
||||
["--model", "m", "--device-ids", "GPU-abcd1234,GPU-ef567890"]
|
||||
)
|
||||
assert parsed.device_ids == ["GPU-abcd1234", "GPU-ef567890"]
|
||||
|
||||
def test_assigned_physical_gpu_ids_are_physical_with_cvd(self, monkeypatch):
|
||||
"""assigned_physical_gpu_ids are already physical and not composed with CVD."""
|
||||
import vllm.platforms.interface as platform_interface
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [4, 5])
|
||||
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
|
||||
|
||||
assert current_platform.device_id_to_physical_device_id(0) == 4
|
||||
assert current_platform.device_id_to_physical_device_id(1) == 5
|
||||
assert current_platform.logical_device_id_to_visible_device_id(0) == 0
|
||||
assert current_platform.logical_device_id_to_visible_device_id(1) == 1
|
||||
|
||||
def test_assigned_physical_gpu_ids_map_to_visible_uuid_cvd(self, monkeypatch):
|
||||
"""Physical IDs map back to visible ordinals when CVD uses UUIDs."""
|
||||
import vllm.platforms.interface as platform_interface
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [5])
|
||||
monkeypatch.setenv(
|
||||
current_platform.device_control_env_var,
|
||||
"GPU-abcd1234,GPU-ef567890",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
type(current_platform),
|
||||
"device_control_id_to_physical_device_id",
|
||||
classmethod(
|
||||
lambda cls, device_id: {"GPU-abcd1234": 4, "GPU-ef567890": 5}[device_id]
|
||||
),
|
||||
)
|
||||
|
||||
assert current_platform.logical_device_id_to_visible_device_id(0) == 1
|
||||
|
||||
def test_device_ids_reject_duplicates(self):
|
||||
"""--device-ids must not contain duplicate entries."""
|
||||
args = EngineArgs(model="m", device_ids=[2, 2])
|
||||
with pytest.raises(ValueError, match="duplicates"):
|
||||
args._resolve_device_ids()
|
||||
|
||||
def test_cli_parsing_strips_whitespace(self):
|
||||
"""--device-ids tolerates whitespace around commas."""
|
||||
parser = FlexibleArgumentParser()
|
||||
EngineArgs.add_cli_args(parser)
|
||||
parsed = parser.parse_args(["--model", "m", "--device-ids", "0, 2, 4"])
|
||||
assert parsed.device_ids == [0, 2, 4]
|
||||
|
||||
def test_visible_ordinal_to_physical_ignores_assigned_ids(self, monkeypatch):
|
||||
"""visible_device_id_to_physical_device_id maps torch device ordinals,
|
||||
independent of the logical-to-physical mapping.
|
||||
|
||||
Regression test: CustomAllreduce passes device.index (a visible
|
||||
ordinal) and must not index into assigned_physical_gpu_ids, which
|
||||
raised IndexError for non-identity --device-ids like [2, 3].
|
||||
"""
|
||||
import vllm.platforms.interface as platform_interface
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
monkeypatch.setattr(platform_interface, "_assigned_physical_gpu_ids", [2, 3])
|
||||
monkeypatch.delenv(current_platform.device_control_env_var, raising=False)
|
||||
|
||||
# CVD unset: visible ordinal == physical ID, even beyond the
|
||||
# assigned list's length.
|
||||
assert current_platform.visible_device_id_to_physical_device_id(2) == 2
|
||||
assert current_platform.visible_device_id_to_physical_device_id(3) == 3
|
||||
|
||||
monkeypatch.setenv(current_platform.device_control_env_var, "4,5")
|
||||
assert current_platform.visible_device_id_to_physical_device_id(1) == 5
|
||||
with pytest.raises(IndexError, match="out of range"):
|
||||
current_platform.visible_device_id_to_physical_device_id(2)
|
||||
|
||||
|
||||
class TestDpDeviceIdSharding:
|
||||
def test_dp_supervisor_device_ids_stay_env_relative(self):
|
||||
"""Regression test: the DP supervisor must pass env-relative indices,
|
||||
not physical IDs, because each child re-resolves --device-ids
|
||||
against its inherited device-control env var."""
|
||||
import argparse
|
||||
|
||||
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
|
||||
|
||||
args = argparse.Namespace(
|
||||
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=None
|
||||
)
|
||||
assert _build_device_ids(args, local_rank=0) == [0, 1]
|
||||
assert _build_device_ids(args, local_rank=1) == [2, 3]
|
||||
|
||||
def test_dp_supervisor_shards_user_device_ids(self):
|
||||
"""User-provided --device-ids are sharded across DP children."""
|
||||
import argparse
|
||||
|
||||
from vllm.entrypoints.openai.dp_supervisor import _build_device_ids
|
||||
|
||||
args = argparse.Namespace(
|
||||
tensor_parallel_size=2, pipeline_parallel_size=1, device_ids=[4, 5, 6, 7]
|
||||
)
|
||||
assert _build_device_ids(args, local_rank=0) == [4, 5]
|
||||
assert _build_device_ids(args, local_rank=1) == [6, 7]
|
||||
with pytest.raises(ValueError, match="needs devices"):
|
||||
_build_device_ids(args, local_rank=2)
|
||||
|
||||
def test_dp_rank_shards_user_assigned_gpu_ids(self):
|
||||
"""get_physical_gpu_ids_for_local_dp_rank slices the user-provided
|
||||
--device-ids list instead of recomputing from the env var."""
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.v1.engine.utils import get_physical_gpu_ids_for_local_dp_rank
|
||||
|
||||
evar = current_platform.device_control_env_var
|
||||
assert get_physical_gpu_ids_for_local_dp_rank(
|
||||
evar, local_dp_rank=1, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
|
||||
) == [6, 7]
|
||||
with pytest.raises(ValueError, match="needs devices"):
|
||||
get_physical_gpu_ids_for_local_dp_rank(
|
||||
evar, local_dp_rank=2, world_size=2, user_assigned_gpu_ids=[4, 5, 6, 7]
|
||||
)
|
||||
|
||||
@@ -364,7 +364,7 @@ class MockVLLMServer:
|
||||
await self._serve_task
|
||||
|
||||
|
||||
def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]):
|
||||
def launch_mock_vllm(child_args: argparse.Namespace):
|
||||
logger.info("Launching mock vLLM on port %s", child_args.port)
|
||||
mock_vllm = MockVLLMServer(
|
||||
port=child_args.port,
|
||||
@@ -375,7 +375,7 @@ def launch_mock_vllm(child_args: argparse.Namespace, env_updates: dict[str, str]
|
||||
|
||||
|
||||
def launch_mock_vllm_with_drain(
|
||||
child_args: argparse.Namespace, env_updates: dict[str, str]
|
||||
child_args: argparse.Namespace,
|
||||
):
|
||||
logger.info("Launching mock vLLM with 15s drain on port %s", child_args.port)
|
||||
mock_vllm = MockVLLMServer(
|
||||
|
||||
@@ -302,6 +302,14 @@ class ParallelConfig:
|
||||
Each entry must use `numactl --physcpubind` CPU-list syntax, for example
|
||||
`"0-3"` or `"0,2,4-7"`.
|
||||
"""
|
||||
assigned_physical_gpu_ids: list[int] | None = None
|
||||
"""Mapping from vLLM-local logical GPU IDs to physical GPU IDs.
|
||||
|
||||
For example, ``[2, 3]`` means logical GPU 0 maps to physical GPU 2,
|
||||
and logical GPU 1 maps to physical GPU 3. Physical IDs are used only
|
||||
at platform/topology boundaries such as NVML, NIC affinity, P2P
|
||||
checks, and final CUDA device selection when needed. When None,
|
||||
logical IDs map to visible device IDs in order."""
|
||||
|
||||
distributed_timeout_seconds: int | None = None
|
||||
"""Timeout in seconds for distributed operations (e.g., init_process_group).
|
||||
@@ -772,6 +780,7 @@ class ParallelConfig:
|
||||
"numa_bind",
|
||||
"numa_bind_nodes",
|
||||
"numa_bind_cpus",
|
||||
"assigned_physical_gpu_ids",
|
||||
}
|
||||
|
||||
from vllm.config.utils import get_hash_factors, hash_factors
|
||||
|
||||
@@ -704,7 +704,14 @@ class FlashInferNVLinkOneSidedManager(All2AllManagerBase):
|
||||
self.num_experts = num_experts
|
||||
|
||||
self.cleanup()
|
||||
gpus_per_node = torch.accelerator.device_count()
|
||||
from vllm.platforms.interface import get_assigned_physical_gpu_ids
|
||||
|
||||
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
|
||||
gpus_per_node = (
|
||||
len(assigned_physical_gpu_ids)
|
||||
if assigned_physical_gpu_ids is not None
|
||||
else torch.accelerator.device_count()
|
||||
)
|
||||
logger.debug(
|
||||
"Making One-sided NVLink mapping: rank=%d, world size=%d",
|
||||
self.rank,
|
||||
|
||||
@@ -320,13 +320,21 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
|
||||
|
||||
is_distributed = dist.is_initialized()
|
||||
|
||||
num_dev = current_platform.device_count()
|
||||
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
|
||||
if cuda_visible_devices is None:
|
||||
cuda_visible_devices = ",".join(str(i) for i in range(num_dev))
|
||||
from vllm.platforms.interface import get_assigned_physical_gpu_ids
|
||||
|
||||
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
# Key by the ordered list: the cache stores directed local-index
|
||||
# pairs, so permutations of the same set are distinct mappings.
|
||||
cache_key = ",".join(str(i) for i in assigned_physical_gpu_ids)
|
||||
num_dev = len(assigned_physical_gpu_ids)
|
||||
else:
|
||||
num_dev = current_platform.device_count()
|
||||
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
|
||||
cache_key = cuda_visible_devices or ",".join(str(i) for i in range(num_dev))
|
||||
|
||||
path = os.path.join(
|
||||
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cuda_visible_devices}.json"
|
||||
envs.VLLM_CACHE_ROOT, f"gpu_p2p_access_cache_for_{cache_key}.json"
|
||||
)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
from vllm.distributed.parallel_state import get_world_group
|
||||
@@ -338,7 +346,15 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
|
||||
# enter this block to calculate the cache
|
||||
logger.info("generating GPU P2P access cache in %s", path)
|
||||
cache: dict[str, bool] = {}
|
||||
ids = list(range(num_dev))
|
||||
# The probe subprocesses inherit this process's device-control env
|
||||
# var, so they must be given visible ordinals, not physical IDs.
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
ids = [
|
||||
current_platform.logical_device_id_to_visible_device_id(local)
|
||||
for local in range(num_dev)
|
||||
]
|
||||
else:
|
||||
ids = list(range(num_dev))
|
||||
# batch of all pairs of GPUs
|
||||
batch_src, batch_tgt = zip(*list(product(ids, ids)))
|
||||
# NOTE: we use `subprocess` rather than `multiprocessing` here
|
||||
@@ -368,8 +384,11 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:
|
||||
) from e
|
||||
with open(output_file.name, "rb") as f:
|
||||
result = pickle.load(f)
|
||||
# Cache entries must be keyed by local indices (0..N-1) because
|
||||
# gpu_p2p_access_check() is called with local ranks.
|
||||
id_to_local = {device_id: local for local, device_id in enumerate(ids)}
|
||||
for _i, _j, r in zip(batch_src, batch_tgt, result):
|
||||
cache[f"{_i}->{_j}"] = r
|
||||
cache[f"{id_to_local[_i]}->{id_to_local[_j]}"] = r
|
||||
with open(path, "w") as f:
|
||||
json.dump(cache, f, indent=4)
|
||||
if is_distributed:
|
||||
|
||||
@@ -34,7 +34,12 @@ def _can_p2p(rank: int, world_size: int) -> bool:
|
||||
continue
|
||||
if envs.VLLM_SKIP_P2P_CHECK:
|
||||
logger.debug("Skipping P2P check and trusting the driver's P2P report.")
|
||||
return torch.cuda.can_device_access_peer(rank, i)
|
||||
# can_device_access_peer takes visible device ordinals, while
|
||||
# rank and i are logical local IDs.
|
||||
return torch.cuda.can_device_access_peer(
|
||||
current_platform.logical_device_id_to_visible_device_id(rank),
|
||||
current_platform.logical_device_id_to_visible_device_id(i),
|
||||
)
|
||||
if not gpu_p2p_access_check(rank, i):
|
||||
return False
|
||||
return True
|
||||
@@ -126,13 +131,10 @@ class CustomAllreduce:
|
||||
CUSTOM_ALL_REDUCE_MAX_SIZES[device_capability_str][world_size],
|
||||
max_size,
|
||||
)
|
||||
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
|
||||
if cuda_visible_devices:
|
||||
device_ids = list(map(int, cuda_visible_devices.split(",")))
|
||||
else:
|
||||
device_ids = list(range(current_platform.device_count()))
|
||||
|
||||
physical_device_id = device_ids[device.index]
|
||||
# device.index is a visible ordinal, not a logical local ID.
|
||||
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
|
||||
device.index
|
||||
)
|
||||
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
|
||||
gather_list = [
|
||||
torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
|
||||
|
||||
@@ -129,12 +129,10 @@ class QuickAllReduce:
|
||||
assert isinstance(device, torch.device)
|
||||
self.device = device
|
||||
|
||||
cuda_visible_devices = envs.CUDA_VISIBLE_DEVICES
|
||||
if cuda_visible_devices:
|
||||
device_ids = list(map(int, cuda_visible_devices.split(",")))
|
||||
else:
|
||||
device_ids = list(range(current_platform.device_count()))
|
||||
physical_device_id = device_ids[device.index]
|
||||
# device.index is a visible ordinal, not a logical local ID.
|
||||
physical_device_id = current_platform.visible_device_id_to_physical_device_id(
|
||||
device.index
|
||||
)
|
||||
tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
|
||||
gather_list = [
|
||||
torch.tensor([0], dtype=torch.int, device="cpu")
|
||||
|
||||
@@ -840,7 +840,13 @@ class MessageQueue:
|
||||
The MessageQueue instance for the calling process,
|
||||
and a list of handles (only non-empty for the reader process).
|
||||
"""
|
||||
local_size = current_platform.device_count()
|
||||
from vllm.platforms.interface import get_assigned_physical_gpu_ids
|
||||
|
||||
assigned_physical_gpu_ids = get_assigned_physical_gpu_ids()
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
local_size = len(assigned_physical_gpu_ids)
|
||||
else:
|
||||
local_size = current_platform.device_count()
|
||||
rank = dist.get_rank()
|
||||
same_node = rank // local_size == reader_rank // local_size
|
||||
buffer_io = MessageQueue(
|
||||
|
||||
@@ -482,10 +482,11 @@ def _init_lmcache_engine(
|
||||
)
|
||||
|
||||
# Change current device.
|
||||
num_gpus = torch.accelerator.device_count()
|
||||
local_rank = parallel_config.rank % num_gpus
|
||||
torch.accelerator.set_device_index(local_rank)
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
from vllm.distributed.parallel_state import get_world_group
|
||||
|
||||
device_index = get_world_group().device_index
|
||||
torch.accelerator.set_device_index(device_index)
|
||||
device = torch.device(f"cuda:{device_index}")
|
||||
metadata = LMCacheEngineMetadata(
|
||||
model_config.model,
|
||||
parallel_config.world_size,
|
||||
|
||||
@@ -392,6 +392,14 @@ class GroupCoordinator:
|
||||
|
||||
self.rank = torch.distributed.get_rank()
|
||||
self.local_rank = local_rank
|
||||
self.device_index: int
|
||||
if _WORLD is not None:
|
||||
self.device_index = _WORLD.device_index
|
||||
else:
|
||||
assert local_rank >= 0, (
|
||||
"local_rank must be provided when creating the world group"
|
||||
)
|
||||
self.device_index = local_rank
|
||||
|
||||
self_device_group = None
|
||||
self_cpu_group = None
|
||||
@@ -442,11 +450,18 @@ class GroupCoordinator:
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_cuda_alike():
|
||||
self.device = torch.device(f"cuda:{local_rank}")
|
||||
visible_device_index = (
|
||||
current_platform.logical_device_id_to_visible_device_id(
|
||||
self.device_index
|
||||
)
|
||||
)
|
||||
self.device = torch.device(f"cuda:{visible_device_index}")
|
||||
elif current_platform.is_xpu():
|
||||
self.device = torch.device(f"xpu:{local_rank}")
|
||||
self.device = torch.device(f"xpu:{self.device_index}")
|
||||
elif current_platform.is_out_of_tree():
|
||||
self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
|
||||
self.device = torch.device(
|
||||
f"{current_platform.device_name}:{self.device_index}"
|
||||
)
|
||||
else:
|
||||
self.device = torch.device("cpu")
|
||||
|
||||
@@ -1438,7 +1453,12 @@ def _init_process_group_for_split_group(
|
||||
"""
|
||||
if torch.accelerator.is_available() and backend != "gloo":
|
||||
init_backend = "cpu:gloo,cuda:nccl"
|
||||
device_id: torch.device | None = torch.device(f"cuda:{local_rank}")
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
visible_device_index = current_platform.logical_device_id_to_visible_device_id(
|
||||
local_rank
|
||||
)
|
||||
device_id: torch.device | None = torch.device(f"cuda:{visible_device_index}")
|
||||
else:
|
||||
init_backend = "gloo"
|
||||
device_id = None
|
||||
|
||||
@@ -86,6 +86,15 @@ class StatelessGroupCoordinator(GroupCoordinator):
|
||||
|
||||
self.rank = global_rank
|
||||
self.local_rank = local_rank
|
||||
from vllm.distributed.parallel_state import _WORLD
|
||||
|
||||
if _WORLD is not None:
|
||||
self.device_index = _WORLD.device_index
|
||||
else:
|
||||
assert local_rank >= 0, (
|
||||
"local_rank must be provided when creating the world group"
|
||||
)
|
||||
self.device_index = local_rank
|
||||
|
||||
self_device_group = None
|
||||
self_cpu_group = None
|
||||
@@ -152,11 +161,18 @@ class StatelessGroupCoordinator(GroupCoordinator):
|
||||
self.tcp_store_group = self_tcp_store_group
|
||||
|
||||
if current_platform.is_cuda_alike():
|
||||
self.device = torch.device(f"cuda:{local_rank}")
|
||||
visible_device_index = (
|
||||
current_platform.logical_device_id_to_visible_device_id(
|
||||
self.device_index
|
||||
)
|
||||
)
|
||||
self.device = torch.device(f"cuda:{visible_device_index}")
|
||||
elif current_platform.is_xpu():
|
||||
self.device = torch.device(f"xpu:{local_rank}")
|
||||
self.device = torch.device(f"xpu:{self.device_index}")
|
||||
elif current_platform.is_out_of_tree():
|
||||
self.device = torch.device(f"{current_platform.device_name}:{local_rank}")
|
||||
self.device = torch.device(
|
||||
f"{current_platform.device_name}:{self.device_index}"
|
||||
)
|
||||
else:
|
||||
self.device = torch.device("cpu")
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ import copy
|
||||
import dataclasses
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass
|
||||
@@ -465,6 +466,7 @@ class EngineArgs:
|
||||
numa_bind: bool = ParallelConfig.numa_bind
|
||||
numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes
|
||||
numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus
|
||||
device_ids: list[int | str] | None = None
|
||||
tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
|
||||
prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
|
||||
decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
|
||||
@@ -979,6 +981,20 @@ class EngineArgs:
|
||||
parallel_group.add_argument(
|
||||
"--numa-bind-cpus", **parallel_kwargs["numa_bind_cpus"]
|
||||
)
|
||||
parallel_group.add_argument(
|
||||
"--device-ids",
|
||||
type=lambda s: [
|
||||
int(device_id) if device_id.isdigit() else device_id
|
||||
for device_id in (part.strip() for part in s.split(","))
|
||||
],
|
||||
default=None,
|
||||
help="Comma-separated physical GPU device IDs or UUIDs to use "
|
||||
'(e.g. --device-ids "2,3,5,7"). Avoids setting '
|
||||
"CUDA_VISIBLE_DEVICES, preserving full GPU topology "
|
||||
"visibility for GPU-NIC affinity and DeepGEMM. "
|
||||
"Note: has no effect with Ray executors; use Ray "
|
||||
"placement groups for GPU selection instead.",
|
||||
)
|
||||
parallel_group.add_argument(
|
||||
"--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"]
|
||||
)
|
||||
@@ -1716,6 +1732,47 @@ class EngineArgs:
|
||||
)
|
||||
return SpeculativeConfig(**self.speculative_config)
|
||||
|
||||
def _resolve_device_ids(self) -> list[int] | None:
|
||||
if not self.device_ids:
|
||||
return None
|
||||
if self.distributed_executor_backend == "ray":
|
||||
logger.warning(
|
||||
"--device-ids has no effect when using the Ray executor. "
|
||||
"Use Ray placement groups for GPU selection instead."
|
||||
)
|
||||
ids = self.device_ids
|
||||
if len(set(ids)) != len(ids):
|
||||
raise ValueError(f"--device-ids must not contain duplicates: {ids}")
|
||||
if all(isinstance(i, str) for i in ids):
|
||||
return [
|
||||
current_platform.device_control_id_to_physical_device_id(i)
|
||||
for i in cast(list[str], ids)
|
||||
]
|
||||
if any(isinstance(i, str) for i in ids):
|
||||
raise ValueError("--device-ids must not mix integer IDs and UUIDs")
|
||||
int_ids = cast(list[int], ids)
|
||||
# Compose with CUDA_VISIBLE_DEVICES: if CVD is set, treat
|
||||
# --device-ids values as indices into the CVD-visible set.
|
||||
cvd = getattr(
|
||||
envs,
|
||||
current_platform.device_control_env_var,
|
||||
os.environ.get(current_platform.device_control_env_var),
|
||||
)
|
||||
if cvd:
|
||||
cvd_ids = [
|
||||
current_platform.device_control_id_to_physical_device_id(x)
|
||||
for x in cvd.split(",")
|
||||
]
|
||||
for i in int_ids:
|
||||
if i >= len(cvd_ids):
|
||||
raise ValueError(
|
||||
f"--device-ids index {i} is out of range for "
|
||||
f"{current_platform.device_control_env_var}"
|
||||
f"={cvd} ({len(cvd_ids)} devices visible)"
|
||||
)
|
||||
return [cvd_ids[i] for i in int_ids]
|
||||
return int_ids
|
||||
|
||||
def create_diffusion_config(self) -> DiffusionConfig | None:
|
||||
if self.diffusion_config is None:
|
||||
return None
|
||||
@@ -2029,6 +2086,7 @@ class EngineArgs:
|
||||
cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
|
||||
_api_process_count=self._api_process_count,
|
||||
_api_process_rank=self._api_process_rank,
|
||||
assigned_physical_gpu_ids=self._resolve_device_ids(),
|
||||
numa_bind=self.numa_bind,
|
||||
numa_bind_nodes=self.numa_bind_nodes,
|
||||
numa_bind_cpus=self.numa_bind_cpus,
|
||||
|
||||
@@ -23,12 +23,10 @@ import uvloop
|
||||
from fastapi import FastAPI, Response
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.system_utils import (
|
||||
decorate_logs,
|
||||
kill_process_tree,
|
||||
set_process_title,
|
||||
update_environment_variables,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -127,22 +125,29 @@ def _build_vllm_dp_server_args(
|
||||
child_args.data_parallel_multi_port_external_lb = False
|
||||
child_args.data_parallel_supervisor_port = None
|
||||
child_args.api_server_count = 1
|
||||
child_args.device_ids = _build_device_ids(args, local_rank)
|
||||
return child_args
|
||||
|
||||
|
||||
def _build_vllm_dp_server_env(
|
||||
args: argparse.Namespace, local_rank: int
|
||||
) -> dict[str, str]:
|
||||
# set visible devices for the child process
|
||||
def _build_device_ids(args: argparse.Namespace, local_rank: int) -> list[int | str]:
|
||||
"""Build the --device-ids value for a DP child process.
|
||||
|
||||
The child resolves these against its own inherited device-control env
|
||||
var (e.g. CUDA_VISIBLE_DEVICES), so integer IDs must stay env-relative
|
||||
here rather than being translated to physical IDs.
|
||||
"""
|
||||
devices_per_rank = args.tensor_parallel_size * args.pipeline_parallel_size
|
||||
start = local_rank * devices_per_rank
|
||||
stop = start + devices_per_rank
|
||||
device_env = current_platform.device_control_env_var
|
||||
visible_devices = ",".join(
|
||||
str(current_platform.device_id_to_physical_device_id(idx))
|
||||
for idx in range(start, stop)
|
||||
)
|
||||
return {device_env: visible_devices}
|
||||
device_ids = getattr(args, "device_ids", None)
|
||||
if device_ids is not None:
|
||||
if stop > len(device_ids):
|
||||
raise ValueError(
|
||||
f"--device-ids has {len(device_ids)} entries, but DP rank "
|
||||
f"{local_rank} needs devices [{start}, {stop})"
|
||||
)
|
||||
return device_ids[start:stop]
|
||||
return list(range(start, stop))
|
||||
|
||||
|
||||
def _child_base_url(args: argparse.Namespace, port: int) -> str:
|
||||
@@ -228,9 +233,7 @@ def _build_dp_supervisor_app(supervisor: DPSupervisor) -> FastAPI:
|
||||
return app
|
||||
|
||||
|
||||
def _run_vllm_dp_server(
|
||||
child_args: argparse.Namespace, env_updates: dict[str, str]
|
||||
) -> None:
|
||||
def _run_vllm_dp_server(child_args: argparse.Namespace) -> None:
|
||||
"""
|
||||
Entrypoint function for the vLLM DP Server.
|
||||
"""
|
||||
@@ -241,7 +244,6 @@ def _run_vllm_dp_server(
|
||||
os.setpgrp()
|
||||
|
||||
name = f"APIServer_DP{child_args.data_parallel_rank}"
|
||||
update_environment_variables(env_updates)
|
||||
set_process_title(name)
|
||||
decorate_logs(name)
|
||||
uvloop.run(run_server(child_args))
|
||||
@@ -345,11 +347,10 @@ class DPSupervisor:
|
||||
context = multiprocessing.get_context("spawn")
|
||||
for local_rank in range(self.args.data_parallel_size_local):
|
||||
child_args = _build_vllm_dp_server_args(self.args, local_rank)
|
||||
child_env = _build_vllm_dp_server_env(self.args, local_rank)
|
||||
process = context.Process(
|
||||
target=_run_vllm_dp_server,
|
||||
name=f"APIServer_DPRank_{child_args.data_parallel_rank}",
|
||||
args=(child_args, child_env),
|
||||
args=(child_args,),
|
||||
)
|
||||
process.start()
|
||||
self._processes.append(process)
|
||||
|
||||
@@ -685,6 +685,15 @@ class CudaPlatformBase(Platform):
|
||||
# all the related functions work on real physical device ids.
|
||||
# the major benefit of using NVML is that it will not initialize CUDA
|
||||
class NvmlCudaPlatform(CudaPlatformBase):
|
||||
@classmethod
|
||||
@with_nvml_context
|
||||
def device_control_id_to_physical_device_id(cls, device_id: str) -> int:
|
||||
try:
|
||||
return int(device_id)
|
||||
except ValueError:
|
||||
handle = pynvml.nvmlDeviceGetHandleByUUID(device_id)
|
||||
return pynvml.nvmlDeviceGetIndex(handle)
|
||||
|
||||
@classmethod
|
||||
@cache
|
||||
@with_nvml_context
|
||||
|
||||
+102
-1
@@ -30,6 +30,33 @@ else:
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
_assigned_physical_gpu_ids: list[int] | None = None
|
||||
|
||||
|
||||
def set_assigned_physical_gpu_ids(ids: list[int]) -> None:
|
||||
"""Set the physical GPU IDs assigned to this worker process.
|
||||
Called during worker init so that device_id_to_physical_device_id()
|
||||
can map local_rank to the correct physical device without relying
|
||||
on CUDA_VISIBLE_DEVICES.
|
||||
|
||||
Idempotent: a second call with the same value is a no-op.
|
||||
Raises RuntimeError if called again with a different value.
|
||||
|
||||
This is expected to run during single-threaded worker initialization."""
|
||||
global _assigned_physical_gpu_ids
|
||||
if _assigned_physical_gpu_ids is not None:
|
||||
if _assigned_physical_gpu_ids != ids:
|
||||
raise RuntimeError(
|
||||
f"set_assigned_physical_gpu_ids called with conflicting values: "
|
||||
f"existing={_assigned_physical_gpu_ids}, new={ids}"
|
||||
)
|
||||
return
|
||||
_assigned_physical_gpu_ids = ids
|
||||
|
||||
|
||||
def get_assigned_physical_gpu_ids() -> list[int] | None:
|
||||
return _assigned_physical_gpu_ids
|
||||
|
||||
|
||||
@functools.cache
|
||||
def in_wsl() -> bool:
|
||||
@@ -233,8 +260,34 @@ class Platform:
|
||||
"""
|
||||
import vllm.kernels # noqa: F401
|
||||
|
||||
@classmethod
|
||||
def device_control_id_to_physical_device_id(cls, device_id: str) -> int:
|
||||
"""Map one device-control env entry to an integer physical device ID."""
|
||||
try:
|
||||
return int(device_id)
|
||||
except ValueError as e:
|
||||
raise ValueError(
|
||||
f"Non-integer device ID {device_id!r} is not supported by "
|
||||
f"{cls.device_name}."
|
||||
) from e
|
||||
|
||||
@classmethod
|
||||
def device_id_to_physical_device_id(cls, device_id: int):
|
||||
"""Map a vLLM-local logical device ID to a physical device ID.
|
||||
|
||||
The input is a logical local ID (e.g. a local rank), NOT a visible
|
||||
device ordinal; for the latter use
|
||||
visible_device_id_to_physical_device_id(). The two coincide only
|
||||
when no logical-to-physical mapping is in effect.
|
||||
"""
|
||||
if _assigned_physical_gpu_ids is not None:
|
||||
if device_id >= len(_assigned_physical_gpu_ids):
|
||||
raise IndexError(
|
||||
f"device_id {device_id} is out of range for "
|
||||
f"assigned_physical_gpu_ids {_assigned_physical_gpu_ids} "
|
||||
f"({len(_assigned_physical_gpu_ids)} devices assigned)"
|
||||
)
|
||||
return _assigned_physical_gpu_ids[device_id]
|
||||
# Treat empty device control env var as unset. This is a valid
|
||||
# configuration in Ray setups where the engine is launched in
|
||||
# a CPU-only placement group located on a GPU node.
|
||||
@@ -244,10 +297,58 @@ class Platform:
|
||||
):
|
||||
device_ids = os.environ[cls.device_control_env_var].split(",")
|
||||
physical_device_id = device_ids[device_id]
|
||||
return int(physical_device_id)
|
||||
return cls.device_control_id_to_physical_device_id(physical_device_id)
|
||||
else:
|
||||
return device_id
|
||||
|
||||
@classmethod
|
||||
def logical_device_id_to_visible_device_id(cls, device_id: int) -> int:
|
||||
"""Map a vLLM-local logical device ID to the current process's
|
||||
visible accelerator ordinal.
|
||||
|
||||
vLLM internals use logical local IDs. Physical IDs are used only
|
||||
at platform/topology boundaries. This helper performs the final
|
||||
translation needed by APIs such as ``torch.device("cuda:N")``.
|
||||
"""
|
||||
physical_device_id = cls.device_id_to_physical_device_id(device_id)
|
||||
device_control_env = os.environ.get(cls.device_control_env_var, "")
|
||||
if not device_control_env:
|
||||
return physical_device_id
|
||||
|
||||
visible_physical_device_ids = [
|
||||
cls.device_control_id_to_physical_device_id(physical_id)
|
||||
for physical_id in device_control_env.split(",")
|
||||
]
|
||||
if physical_device_id not in visible_physical_device_ids:
|
||||
raise RuntimeError(
|
||||
f"Physical device {physical_device_id} for logical device "
|
||||
f"{device_id} is not visible in {cls.device_control_env_var}="
|
||||
f"{device_control_env}"
|
||||
)
|
||||
return visible_physical_device_ids.index(physical_device_id)
|
||||
|
||||
@classmethod
|
||||
def visible_device_id_to_physical_device_id(cls, device_id: int) -> int:
|
||||
"""Map a visible accelerator ordinal (e.g. ``torch.device.index``)
|
||||
to a physical device ID.
|
||||
|
||||
This is the inverse of the env-var translation performed by
|
||||
logical_device_id_to_visible_device_id() and is independent of any
|
||||
logical-to-physical mapping set via set_assigned_physical_gpu_ids().
|
||||
"""
|
||||
device_control_env = os.environ.get(cls.device_control_env_var, "")
|
||||
if not device_control_env:
|
||||
return device_id
|
||||
visible_device_ids = device_control_env.split(",")
|
||||
if device_id >= len(visible_device_ids):
|
||||
raise IndexError(
|
||||
f"visible device ordinal {device_id} is out of range for "
|
||||
f"{cls.device_control_env_var}={device_control_env}"
|
||||
)
|
||||
return cls.device_control_id_to_physical_device_id(
|
||||
visible_device_ids[device_id]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def import_kernels(cls) -> None:
|
||||
"""Import any platform-specific C kernels."""
|
||||
|
||||
+16
-9
@@ -74,7 +74,7 @@ from vllm.v1.engine.utils import (
|
||||
EngineHandshakeMetadata,
|
||||
EngineZmqAddresses,
|
||||
SignalCallback,
|
||||
get_device_indices,
|
||||
get_physical_gpu_ids_for_local_dp_rank,
|
||||
)
|
||||
from vllm.v1.executor import Executor
|
||||
from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind
|
||||
@@ -2175,23 +2175,30 @@ class EngineCoreActorMixin:
|
||||
pass
|
||||
else:
|
||||
device_control_env_var = current_platform.device_control_env_var
|
||||
self._set_cuda_visible_devices(
|
||||
self._set_assigned_physical_gpu_ids(
|
||||
vllm_config, local_dp_rank, device_control_env_var
|
||||
)
|
||||
|
||||
def _set_cuda_visible_devices(
|
||||
self, vllm_config: VllmConfig, local_dp_rank: int, device_control_env_var: str
|
||||
def _set_assigned_physical_gpu_ids(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
local_dp_rank: int,
|
||||
device_control_env_var: str,
|
||||
):
|
||||
world_size = vllm_config.parallel_config.world_size
|
||||
# Set CUDA_VISIBLE_DEVICES or equivalent.
|
||||
try:
|
||||
value = get_device_indices(
|
||||
device_control_env_var, local_dp_rank, world_size
|
||||
physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
|
||||
device_control_env_var,
|
||||
local_dp_rank,
|
||||
world_size,
|
||||
user_assigned_gpu_ids=(
|
||||
vllm_config.parallel_config.assigned_physical_gpu_ids
|
||||
),
|
||||
)
|
||||
os.environ[device_control_env_var] = value
|
||||
vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids
|
||||
except IndexError as e:
|
||||
raise Exception(
|
||||
f"Error setting {device_control_env_var}: "
|
||||
f"Error computing assigned_physical_gpu_ids: "
|
||||
f"local range: [{local_dp_rank * world_size}, "
|
||||
f"{(local_dp_rank + 1) * world_size}) "
|
||||
f'base value: "{os.getenv(device_control_env_var)}"'
|
||||
|
||||
+66
-43
@@ -12,7 +12,6 @@ from multiprocessing import Process, connection
|
||||
from multiprocessing.process import BaseProcess
|
||||
from multiprocessing.queues import Queue
|
||||
from typing import TYPE_CHECKING, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import msgspec
|
||||
import zmq
|
||||
@@ -175,38 +174,38 @@ class CoreEngineProcManager:
|
||||
self.manager_stopped = threading.Event()
|
||||
self.failed_proc_name: str | None = None
|
||||
|
||||
# All ranks share this config object: capture the user-provided
|
||||
# --device-ids list before the per-rank shard overwrites it. Mutating
|
||||
# the config before each proc.start() works because the spawn method
|
||||
# pickles process args at start() time, sequentially per rank.
|
||||
user_assigned_gpu_ids = vllm_config.parallel_config.assigned_physical_gpu_ids
|
||||
try:
|
||||
for proc, local_dp_rank in zip(self.processes, local_dp_ranks):
|
||||
# Adjust device control in DP for platforms that cannot rely
|
||||
# on torch.accelerator.set_device_index(), and for Ray launchers.
|
||||
device_control_context: contextlib.AbstractContextManager[None] = (
|
||||
contextlib.nullcontext()
|
||||
)
|
||||
# Populate the logical-to-physical GPU mapping in DP for
|
||||
# platforms that cannot rely on
|
||||
# torch.accelerator.set_device_index(), and for Ray.
|
||||
needs_device_env_isolation = not (
|
||||
current_platform.is_cuda_alike() or current_platform.is_xpu()
|
||||
)
|
||||
if is_dp and (
|
||||
needs_device_env_isolation or vllm_config.parallel_config.use_ray
|
||||
):
|
||||
device_control_context = set_device_control_env_var(
|
||||
vllm_config, local_dp_rank
|
||||
set_assigned_physical_gpu_ids_for_dp_rank(
|
||||
vllm_config, local_dp_rank, user_assigned_gpu_ids
|
||||
)
|
||||
|
||||
with (
|
||||
device_control_context,
|
||||
numa_utils.configure_subprocess(
|
||||
# EngineCore itself does not have a TP/PP-local rank.
|
||||
# When DP is enabled, set_device_control_env_var()
|
||||
# narrows visible devices to this DP shard first, so
|
||||
# local_rank=0 means "the first local GPU in this
|
||||
# shard". The actual TP/PP worker processes spawned by
|
||||
# the executor are bound separately with their own
|
||||
# local_rank values.
|
||||
vllm_config,
|
||||
local_rank=0,
|
||||
dp_local_rank=local_dp_rank,
|
||||
process_kind="EngineCore",
|
||||
),
|
||||
with numa_utils.configure_subprocess(
|
||||
# EngineCore itself does not have a TP/PP-local rank.
|
||||
# When DP is enabled, set_assigned_physical_gpu_ids_for_dp_rank()
|
||||
# populates the logical-to-physical mapping for this DP
|
||||
# shard, so local_rank=0 means "the first local GPU in
|
||||
# this shard". The actual TP/PP worker processes spawned
|
||||
# by the executor are bound separately with their own
|
||||
# local_rank values.
|
||||
vllm_config,
|
||||
local_rank=0,
|
||||
dp_local_rank=local_dp_rank,
|
||||
process_kind="EngineCore",
|
||||
):
|
||||
proc.start()
|
||||
finally:
|
||||
@@ -281,55 +280,79 @@ class SignalCallback:
|
||||
self._event.set()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def set_device_control_env_var(
|
||||
vllm_config: VllmConfig, local_dp_rank: int
|
||||
) -> Iterator[None]:
|
||||
def set_assigned_physical_gpu_ids_for_dp_rank(
|
||||
vllm_config: VllmConfig,
|
||||
local_dp_rank: int,
|
||||
user_assigned_gpu_ids: list[int] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Temporarily set CUDA_VISIBLE_DEVICES or equivalent
|
||||
for engine subprocess.
|
||||
Populate assigned_physical_gpu_ids on the config for the given DP rank.
|
||||
|
||||
user_assigned_gpu_ids is the full (un-sharded) --device-ids list, if the
|
||||
user provided one; this DP rank's shard is sliced from it. It is passed
|
||||
explicitly rather than read from the config because callers may reuse
|
||||
one config object across DP ranks, overwriting the field each time.
|
||||
"""
|
||||
world_size = vllm_config.parallel_config.world_size
|
||||
local_world_size = vllm_config.parallel_config.local_world_size
|
||||
evar = current_platform.device_control_env_var
|
||||
|
||||
value = get_device_indices(evar, local_dp_rank, world_size, local_world_size)
|
||||
with patch.dict(os.environ, values=((evar, value),)):
|
||||
yield
|
||||
physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
|
||||
evar,
|
||||
local_dp_rank,
|
||||
world_size,
|
||||
local_world_size,
|
||||
user_assigned_gpu_ids=user_assigned_gpu_ids,
|
||||
)
|
||||
vllm_config.parallel_config.assigned_physical_gpu_ids = physical_gpu_ids
|
||||
|
||||
|
||||
def get_device_indices(
|
||||
def get_physical_gpu_ids_for_local_dp_rank(
|
||||
device_control_env_var: str,
|
||||
local_dp_rank: int,
|
||||
world_size: int,
|
||||
local_world_size: int | None = None,
|
||||
):
|
||||
user_assigned_gpu_ids: list[int] | None = None,
|
||||
) -> list[int]:
|
||||
"""
|
||||
Returns a comma-separated string of device indices for the specified
|
||||
Returns list of physical GPU IDs for the specified
|
||||
data parallel rank.
|
||||
|
||||
For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
|
||||
this will select devices 2 and 3 for local_dp_rank=1.
|
||||
this will return [2, 3] for local_dp_rank=1.
|
||||
|
||||
If user_assigned_gpu_ids is provided (e.g. from --device-ids), this DP
|
||||
rank's shard is sliced from it instead of being derived from the
|
||||
device-control env var.
|
||||
"""
|
||||
if local_world_size is None:
|
||||
local_world_size = world_size
|
||||
if user_assigned_gpu_ids is not None:
|
||||
start = local_dp_rank * world_size
|
||||
stop = start + local_world_size
|
||||
if stop > len(user_assigned_gpu_ids):
|
||||
raise ValueError(
|
||||
f"--device-ids provides {len(user_assigned_gpu_ids)} devices, "
|
||||
f"but DP rank {local_dp_rank} needs devices [{start}, {stop})"
|
||||
)
|
||||
return user_assigned_gpu_ids[start:stop]
|
||||
try:
|
||||
value = ",".join(
|
||||
str(current_platform.device_id_to_physical_device_id(i))
|
||||
return [
|
||||
current_platform.device_id_to_physical_device_id(i)
|
||||
for i in range(
|
||||
local_dp_rank * world_size,
|
||||
local_dp_rank * world_size + local_world_size,
|
||||
)
|
||||
)
|
||||
]
|
||||
except IndexError as e:
|
||||
raise Exception(
|
||||
f"Error setting {device_control_env_var}: "
|
||||
f"Error computing device indices for "
|
||||
f"{device_control_env_var}: "
|
||||
f"local range: [{local_dp_rank * world_size}, "
|
||||
f"{(local_dp_rank + 1) * world_size}) "
|
||||
"base value: "
|
||||
f'"{os.getenv(device_control_env_var)}"'
|
||||
) from e
|
||||
return value
|
||||
|
||||
|
||||
def _apply_dp_identity_suffix(dp_vllm_config, dp_rank: int) -> None:
|
||||
@@ -453,11 +476,11 @@ class CoreEngineActorManager:
|
||||
# https://github.com/ray-project/ray/blob/master/python/ray/_private/accelerators/intel_gpu.py#L56 # noqa: E501
|
||||
if current_platform.is_xpu():
|
||||
device_evar = current_platform.device_control_env_var
|
||||
device_indices = get_device_indices(
|
||||
physical_gpu_ids = get_physical_gpu_ids_for_local_dp_rank(
|
||||
device_evar, local_index, world_size
|
||||
)
|
||||
actor_env_vars = self.env_vars_dict.copy()
|
||||
actor_env_vars[device_evar] = device_indices
|
||||
actor_env_vars[device_evar] = ",".join(str(d) for d in physical_gpu_ids)
|
||||
runtime_env = RuntimeEnv(env_vars=actor_env_vars)
|
||||
|
||||
actor = (
|
||||
|
||||
@@ -826,6 +826,16 @@ class WorkerProc:
|
||||
signal.signal(signal.SIGTERM, signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
|
||||
# Publish the logical-to-physical mapping early so topology helpers
|
||||
# work before init_device (needed by set_worker_net_device below).
|
||||
assigned_physical_gpu_ids = kwargs[
|
||||
"vllm_config"
|
||||
].parallel_config.assigned_physical_gpu_ids
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
from vllm.platforms.interface import set_assigned_physical_gpu_ids
|
||||
|
||||
set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
|
||||
|
||||
# Set net device env vars for the worker if VLLM_GPU_NIC_PCIE_MAPPING is set
|
||||
set_worker_net_device(kwargs.get("local_rank", 0), kwargs["vllm_config"])
|
||||
|
||||
|
||||
@@ -258,30 +258,35 @@ class RayDistributedExecutor(Executor):
|
||||
}
|
||||
self.collective_rpc("adjust_rank", args=(rerank_mapping,))
|
||||
|
||||
# Get the set of GPU IDs used on each node.
|
||||
worker_node_and_gpu_ids = []
|
||||
# Get the set of physical GPU IDs used on each node.
|
||||
worker_node_and_physical_gpu_ids = []
|
||||
for worker in [self.driver_dummy_worker] + self.workers:
|
||||
if worker is None:
|
||||
# driver_dummy_worker can be None when using ray spmd worker.
|
||||
continue
|
||||
worker_node_and_gpu_ids.append(
|
||||
ray.get(worker.get_node_and_gpu_ids.remote()) # type: ignore[attr-defined]
|
||||
worker_node_and_physical_gpu_ids.append(
|
||||
ray.get(worker.get_node_and_physical_gpu_ids.remote()) # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
node_workers = defaultdict(list) # node id -> list of worker ranks
|
||||
node_gpus = defaultdict(list) # node id -> list of gpu ids
|
||||
node_physical_gpu_ids = defaultdict(list) # node id -> physical GPU IDs
|
||||
|
||||
for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
|
||||
for i, (node_id, physical_gpu_ids) in enumerate(
|
||||
worker_node_and_physical_gpu_ids
|
||||
):
|
||||
node_workers[node_id].append(i)
|
||||
# `gpu_ids` can be a list of strings or integers.
|
||||
# `physical_gpu_ids` can be a list of strings or integers.
|
||||
# convert them to integers for consistency.
|
||||
# NOTE: gpu_ids can be larger than 9 (e.g. 16 GPUs),
|
||||
# NOTE: physical GPU IDs can be larger than 9 (e.g. 16 GPUs),
|
||||
# string sorting is not sufficient.
|
||||
# see https://github.com/vllm-project/vllm/issues/5590
|
||||
gpu_ids = [int(x) for x in gpu_ids]
|
||||
node_gpus[node_id].extend(gpu_ids)
|
||||
for node_id, gpu_ids in node_gpus.items():
|
||||
node_gpus[node_id] = sorted(gpu_ids)
|
||||
physical_gpu_ids = [
|
||||
current_platform.device_control_id_to_physical_device_id(str(x))
|
||||
for x in physical_gpu_ids
|
||||
]
|
||||
node_physical_gpu_ids[node_id].extend(physical_gpu_ids)
|
||||
for node_id, physical_gpu_ids in node_physical_gpu_ids.items():
|
||||
node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)
|
||||
|
||||
all_ips = set(worker_ips + [driver_ip])
|
||||
n_ips = len(all_ips)
|
||||
@@ -297,23 +302,8 @@ class RayDistributedExecutor(Executor):
|
||||
" each node."
|
||||
)
|
||||
|
||||
# Set environment variables for the driver and workers.
|
||||
# We set CUDA_VISIBLE_DEVICES to ALL GPUs on the node for each worker.
|
||||
# This is needed because:
|
||||
# 1. Ray's compiled DAG needs to find the allocated GPU in
|
||||
# CUDA_VISIBLE_DEVICES.
|
||||
# 2. vLLM's communication layer (NCCL, CustomAllreduce) needs to see
|
||||
# all GPUs for P2P checks and communication setup. Though if it was
|
||||
# just this reason, we could have also just kept the visible devices
|
||||
# unset.
|
||||
# Each worker will use local_rank to index into the visible devices.
|
||||
all_args_to_update_environment_variables = [
|
||||
{
|
||||
current_platform.device_control_env_var: ",".join(
|
||||
map(str, node_gpus[node_id])
|
||||
),
|
||||
}
|
||||
for (node_id, _) in worker_node_and_gpu_ids
|
||||
all_args_to_update_environment_variables: list[dict[str, str]] = [
|
||||
{} for _ in worker_node_and_physical_gpu_ids
|
||||
]
|
||||
|
||||
# Environment variables to copy from driver to workers
|
||||
@@ -336,7 +326,7 @@ class RayDistributedExecutor(Executor):
|
||||
"update_environment_variables", args=(self._get_env_vars_to_be_updated(),)
|
||||
)
|
||||
|
||||
if len(node_gpus) == 1:
|
||||
if len(node_physical_gpu_ids) == 1:
|
||||
# in single node case, we don't need to get the IP address.
|
||||
# the loopback address is sufficient
|
||||
# NOTE: a node may have several IP addresses, one for each
|
||||
@@ -352,10 +342,11 @@ class RayDistributedExecutor(Executor):
|
||||
|
||||
# Initialize the actual workers inside worker wrapper.
|
||||
all_kwargs = []
|
||||
for rank, (node_id, _) in enumerate(worker_node_and_gpu_ids):
|
||||
for rank, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):
|
||||
local_rank = node_workers[node_id].index(rank)
|
||||
kwargs = dict(
|
||||
vllm_config=self.vllm_config,
|
||||
assigned_physical_gpu_ids=sorted(node_physical_gpu_ids[node_id]),
|
||||
local_rank=local_rank,
|
||||
rank=rank,
|
||||
distributed_init_method=distributed_init_method,
|
||||
|
||||
@@ -79,24 +79,25 @@ class RayWorkerProc(WorkerProc):
|
||||
1. __init__: lightweight setup, stores init args (no device/model init)
|
||||
2. initialize_worker: called after GPU IDs are discovered, completes
|
||||
the full WorkerProc initialization with the correct local_rank and
|
||||
CUDA_VISIBLE_DEVICES.
|
||||
logical-to-physical GPU mapping.
|
||||
|
||||
CUDA_VISIBLE_DEVICES setup flow:
|
||||
GPU assignment flow:
|
||||
|
||||
1. RayExecutorV2 enables RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES so Ray does
|
||||
not set CUDA_VISIBLE_DEVICES on RayWorkerProc actors at creation time.
|
||||
2. Each actor is scheduled with a placement group and bundle index; Ray resolves
|
||||
the physical GPU ID for that bundle at placement time.
|
||||
3. After placement, the worker discovers that GPU ID and sets
|
||||
CUDA_VISIBLE_DEVICES before finishing WorkerProc initialization.
|
||||
3. After placement, the executor discovers each worker's GPU ID and passes the
|
||||
node's logical-to-physical mapping (assigned_physical_gpu_ids) to
|
||||
initialize_worker(); CUDA_VISIBLE_DEVICES is never modified.
|
||||
|
||||
There is no workaround for this unset-and-reset sequence when the placement group
|
||||
is externally managed: scheduling must complete before CUDA_VISIBLE_DEVICES can
|
||||
match the GPU tied to the worker's bundle.
|
||||
Scheduling must complete before the mapping is known when the placement
|
||||
group is externally managed: only then is the GPU tied to the worker's
|
||||
bundle resolved.
|
||||
|
||||
This sequence allows multiple vLLM instances to coexist on the same node:
|
||||
each instance is unaware which physical devices others hold, and the
|
||||
externally managed placement group avoids CUDA_VISIBLE_DEVICES conflicts
|
||||
externally managed placement group avoids device assignment conflicts
|
||||
by binding workers to specific placement group bundles.
|
||||
"""
|
||||
|
||||
@@ -120,28 +121,33 @@ class RayWorkerProc(WorkerProc):
|
||||
is_driver_worker=is_driver_worker,
|
||||
)
|
||||
|
||||
def get_node_and_gpu_ids(self) -> tuple[str, list[int]]:
|
||||
"""Return (node_id, gpu_ids) assigned to this actor by Ray."""
|
||||
def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:
|
||||
"""Return (node_id, physical_gpu_ids) assigned to this actor by Ray."""
|
||||
node_id = ray.get_runtime_context().get_node_id()
|
||||
device_key = current_platform.ray_device_key
|
||||
if not device_key:
|
||||
raise RuntimeError(
|
||||
f"current platform {current_platform.device_name} does not support ray."
|
||||
)
|
||||
gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
|
||||
return node_id, [int(x) for x in gpu_ids]
|
||||
physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
|
||||
return node_id, [
|
||||
current_platform.device_control_id_to_physical_device_id(str(x))
|
||||
for x in physical_gpu_ids
|
||||
]
|
||||
|
||||
def initialize_worker(
|
||||
self,
|
||||
local_rank: int,
|
||||
env_vars: dict[str, str],
|
||||
driver_env_vars: dict[str, str] | None = None,
|
||||
assigned_physical_gpu_ids: list[int] | None = None,
|
||||
) -> None:
|
||||
"""Complete initialization after GPU assignment is known.
|
||||
|
||||
*driver_env_vars* are applied with ``setdefault`` — they fill
|
||||
in missing vars but never overwrite node-local values.
|
||||
*env_vars* (e.g. CUDA_VISIBLE_DEVICES) always overwrite.
|
||||
*env_vars* always overwrite.
|
||||
*assigned_physical_gpu_ids* maps local_rank to physical CUDA device ID.
|
||||
"""
|
||||
if driver_env_vars:
|
||||
for key, value in driver_env_vars.items():
|
||||
@@ -149,6 +155,13 @@ class RayWorkerProc(WorkerProc):
|
||||
for key, value in env_vars.items():
|
||||
os.environ[key] = value
|
||||
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
vllm_config = self._init_kwargs["vllm_config"]
|
||||
assert isinstance(vllm_config, VllmConfig)
|
||||
vllm_config.parallel_config.assigned_physical_gpu_ids = (
|
||||
assigned_physical_gpu_ids
|
||||
)
|
||||
|
||||
self.local_rank = local_rank
|
||||
super().__init__(
|
||||
local_rank=local_rank,
|
||||
@@ -365,36 +378,48 @@ class RayExecutorV2(MultiprocExecutor):
|
||||
)
|
||||
self.ray_worker_handles.append(handle)
|
||||
|
||||
# Step 6: Discover GPU IDs assigned to each worker via Ray runtime context.
|
||||
worker_node_and_gpu_ids = ray.get(
|
||||
[h.actor.get_node_and_gpu_ids.remote() for h in self.ray_worker_handles]
|
||||
# Step 6: Discover physical GPU IDs assigned to each worker via Ray
|
||||
# runtime context.
|
||||
worker_node_and_physical_gpu_ids = ray.get(
|
||||
[
|
||||
h.actor.get_node_and_physical_gpu_ids.remote()
|
||||
for h in self.ray_worker_handles
|
||||
]
|
||||
)
|
||||
|
||||
node_workers: dict[str, list[int]] = defaultdict(list)
|
||||
node_gpus: dict[str, list[int]] = defaultdict(list)
|
||||
for i, (node_id, gpu_ids) in enumerate(worker_node_and_gpu_ids):
|
||||
node_physical_gpu_ids: dict[str, list[int]] = defaultdict(list)
|
||||
for i, (node_id, physical_gpu_ids) in enumerate(
|
||||
worker_node_and_physical_gpu_ids
|
||||
):
|
||||
node_workers[node_id].append(i)
|
||||
node_gpus[node_id].extend(gpu_ids)
|
||||
for node_id, gpu_ids in node_gpus.items():
|
||||
node_gpus[node_id] = sorted(gpu_ids)
|
||||
node_physical_gpu_ids[node_id].extend(physical_gpu_ids)
|
||||
for node_id, physical_gpu_ids in node_physical_gpu_ids.items():
|
||||
node_physical_gpu_ids[node_id] = sorted(physical_gpu_ids)
|
||||
|
||||
# Step 7: Initialize workers with correct local_rank and
|
||||
# CUDA_VISIBLE_DEVICES. Each worker sees all GPUs assigned to
|
||||
# this executor on its node; local_rank indexes into that set.
|
||||
# Step 7: Initialize workers with local logical ranks and the
|
||||
# logical-to-physical GPU mapping discovered from Ray placement.
|
||||
init_worker_refs = []
|
||||
for i, (node_id, _) in enumerate(worker_node_and_gpu_ids):
|
||||
for i, (node_id, _) in enumerate(worker_node_and_physical_gpu_ids):
|
||||
local_rank = node_workers[node_id].index(i)
|
||||
worker_env_vars = {
|
||||
current_platform.device_control_env_var: ",".join(
|
||||
map(str, node_gpus[node_id])
|
||||
),
|
||||
}
|
||||
assigned_physical_gpu_ids = sorted(node_physical_gpu_ids[node_id])
|
||||
worker_env_vars: dict[str, str] = {}
|
||||
self.ray_worker_handles[i].local_rank = local_rank
|
||||
init_worker_refs.append(
|
||||
self.ray_worker_handles[i].actor.initialize_worker.remote(
|
||||
local_rank, worker_env_vars, self.driver_env_vars
|
||||
local_rank,
|
||||
worker_env_vars,
|
||||
self.driver_env_vars,
|
||||
assigned_physical_gpu_ids=assigned_physical_gpu_ids,
|
||||
)
|
||||
)
|
||||
# Also set on the executor-side config for consistency. The mapping
|
||||
# is per-node, so only do this when all workers share one node.
|
||||
if len(node_physical_gpu_ids) == 1:
|
||||
node_id_0 = worker_node_and_physical_gpu_ids[0][0]
|
||||
self.vllm_config.parallel_config.assigned_physical_gpu_ids = sorted(
|
||||
node_physical_gpu_ids[node_id_0]
|
||||
)
|
||||
ray.get(init_worker_refs)
|
||||
|
||||
# Step 8: Collect response MQ handles
|
||||
|
||||
@@ -93,7 +93,7 @@ try:
|
||||
def get_node_ip(self) -> str:
|
||||
return get_ip()
|
||||
|
||||
def get_node_and_gpu_ids(self) -> tuple[str, list[int]]:
|
||||
def get_node_and_physical_gpu_ids(self) -> tuple[str, list[int]]:
|
||||
node_id = ray.get_runtime_context().get_node_id()
|
||||
device_key = vllm.platforms.current_platform.ray_device_key
|
||||
if not device_key:
|
||||
@@ -101,8 +101,10 @@ try:
|
||||
"current platform %s does not support ray.",
|
||||
vllm.platforms.current_platform.device_name,
|
||||
)
|
||||
gpu_ids = ray.get_runtime_context().get_accelerator_ids()[device_key]
|
||||
return node_id, gpu_ids
|
||||
physical_gpu_ids = ray.get_runtime_context().get_accelerator_ids()[
|
||||
device_key
|
||||
]
|
||||
return node_id, physical_gpu_ids
|
||||
|
||||
def setup_device_if_necessary(self):
|
||||
# TODO(swang): This is needed right now because Ray CG executes
|
||||
|
||||
@@ -270,19 +270,47 @@ class Worker(WorkerBase):
|
||||
|
||||
# DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK
|
||||
self.local_rank += dp_local_rank * tp_pp_world_size
|
||||
|
||||
# Publish the logical-to-physical mapping for topology queries
|
||||
# such as NIC affinity and P2P checks.
|
||||
assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
from vllm.platforms.interface import set_assigned_physical_gpu_ids
|
||||
|
||||
set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
|
||||
assert self.local_rank < len(assigned_physical_gpu_ids), (
|
||||
f"local_rank {self.local_rank} is out of bounds for "
|
||||
f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
|
||||
)
|
||||
# NOTE(patch pr45026): local_world_size is derived from
|
||||
# parallel_config.nnodes, which is only set for the "mp"
|
||||
# multi-node backend. With the "ray"/"external_launcher"
|
||||
# backends nnodes stays 1, so local_world_size collapses to
|
||||
# the full world_size and this check wrongly fires on
|
||||
# cross-node deployments. assigned_physical_gpu_ids is already
|
||||
# per-node and the local_rank bound above fully validates the
|
||||
# mapping for these backends, so skip the check for them.
|
||||
if parallel_config.distributed_executor_backend not in (
|
||||
"ray",
|
||||
"external_launcher",
|
||||
):
|
||||
assert self.parallel_config.local_world_size <= len(
|
||||
assigned_physical_gpu_ids
|
||||
), (
|
||||
f"local_world_size ({self.parallel_config.local_world_size})"
|
||||
" exceeds assigned_physical_gpu_ids count "
|
||||
f"({len(assigned_physical_gpu_ids)})"
|
||||
)
|
||||
else:
|
||||
assert self.local_rank < torch.accelerator.device_count(), (
|
||||
f"DP adjusted local rank {self.local_rank} is out of bounds. "
|
||||
)
|
||||
visible_device_count = (
|
||||
torch.accelerator.device_count() if torch.cuda.is_available() else 0
|
||||
)
|
||||
assert self.parallel_config.local_world_size <= visible_device_count, (
|
||||
f"local_world_size ({self.parallel_config.local_world_size}) must "
|
||||
f"be less than or equal to the number of visible devices "
|
||||
f"({visible_device_count})."
|
||||
f"DP adjusted local rank {self.local_rank} is out of "
|
||||
f"bounds for {torch.accelerator.device_count()} devices."
|
||||
)
|
||||
|
||||
self.device = torch.device(f"cuda:{self.local_rank}")
|
||||
visible_device_index = (
|
||||
current_platform.logical_device_id_to_visible_device_id(self.local_rank)
|
||||
)
|
||||
self.device = torch.device(f"cuda:{visible_device_index}")
|
||||
torch.accelerator.set_device_index(self.device)
|
||||
|
||||
current_platform.check_if_supports_dtype(self.model_config.dtype)
|
||||
|
||||
@@ -286,6 +286,12 @@ class WorkerWrapperBase:
|
||||
extended_calls,
|
||||
)
|
||||
|
||||
assigned_physical_gpu_ids = kwargs.pop("assigned_physical_gpu_ids", None)
|
||||
if assigned_physical_gpu_ids is not None:
|
||||
vllm_config.parallel_config.assigned_physical_gpu_ids = (
|
||||
assigned_physical_gpu_ids
|
||||
)
|
||||
|
||||
shared_worker_lock = kwargs.pop("shared_worker_lock", None)
|
||||
if shared_worker_lock is None:
|
||||
msg = (
|
||||
|
||||
Reference in New Issue
Block a user