xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-19 05:27:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add TPU to DeviceConfig

Browse Source
This commit is contained in:
Woosuk Kwon 2024-04-01 03:19:17 +00:00
parent 3b8f43024f
commit 824521c987
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
vllm/config.py
View File

@ -10,7 +10,8 @@ from transformers import PretrainedConfig
from vllm.logger import init_logger from vllm.logger import init_logger
from vllm.transformers_utils.config import get_config, get_hf_text_config from vllm.transformers_utils.config import get_config, get_hf_text_config
from vllm.utils import get_cpu_memory, get_nvcc_cuda_version, is_hip, is_neuron from vllm.utils import (get_cpu_memory, get_nvcc_cuda_version, is_hip,
is_neuron, is_tpu)
if TYPE_CHECKING: if TYPE_CHECKING:
from ray.util.placement_group import PlacementGroup from ray.util.placement_group import PlacementGroup
@ -582,6 +583,8 @@ class DeviceConfig:
# Automated device type detection # Automated device type detection
if is_neuron(): if is_neuron():
self.device_type = "neuron" self.device_type = "neuron"
elif is_tpu():
self.device_type = "tpu"
else: else:
# We don't call torch.cuda.is_available() here to # We don't call torch.cuda.is_available() here to
# avoid initializing CUDA before workers are forked # avoid initializing CUDA before workers are forked