xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-06-10 09:42:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BugFix] Avoid initializing CUDA too early (#3487)

Browse Source
This commit is contained in:
Nick Hill 2024-03-18 23:05:20 -07:00 committed by GitHub
parent ef65dcfa6f
commit 7341c77d69
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/config.py
View File

@ -577,12 +577,12 @@ class DeviceConfig:
def __init__(self, device: str = "auto") -> None: def __init__(self, device: str = "auto") -> None:
if device == "auto": if device == "auto":
# Automated device type detection # Automated device type detection
if torch.cuda.is_available(): if is_neuron():
self.device_type = "cuda"
elif is_neuron():
self.device_type = "neuron" self.device_type = "neuron"
else: else:
raise RuntimeError("No supported device detected.") # We don't call torch.cuda.is_available() here to
# avoid initializing CUDA before workers are forked
self.device_type = "cuda"
else: else:
# Device type is assigned explicitly # Device type is assigned explicitly
self.device_type = device self.device_type = device