[ROCm][AMD] unify CUDA_VISIBLE_DEVICES usage in cuda/rocm (#6352)

This commit is contained in:
Hongxia Yang
2024-07-12 00:30:46 -04:00
committed by GitHub
parent d26a8b3f1f
commit b6c16cf8ff
5 changed files with 10 additions and 34 deletions

View File

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, ClassVar, List, Optional, Tuple, Union
import torch
from transformers import PretrainedConfig
import vllm.envs as envs
from vllm.logger import init_logger
from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS
from vllm.model_executor.models import ModelRegistry
@@ -14,7 +13,7 @@ from vllm.tracing import is_otel_installed
from vllm.transformers_utils.config import get_config, get_hf_text_config
from vllm.utils import (cuda_device_count_stateless, get_cpu_memory, is_cpu,
is_hip, is_neuron, is_openvino, is_tpu, is_xpu,
print_warning_once, update_environment_variables)
print_warning_once)
if TYPE_CHECKING:
from ray.util.placement_group import PlacementGroup
@@ -695,12 +694,6 @@ class ParallelConfig:
self.distributed_executor_backend = backend
logger.info("Defaulting to use %s for distributed inference",
backend)
# If CUDA_VISIBLE_DEVICES is set on ROCm prior to vLLM init,
# propagate changes to HIP_VISIBLE_DEVICES (conversion handled by
# the update_environment_variables function)
if is_hip() and envs.CUDA_VISIBLE_DEVICES:
update_environment_variables(
{"CUDA_VISIBLE_DEVICES": envs.CUDA_VISIBLE_DEVICES})
self._verify_args()
self.rank = 0