Skip to content

Commit

Permalink
[Device] Defer the CPU detection (#3126)
Browse files Browse the repository at this point in the history
This PR moves the CPU detection to the last device after
checking all GPUs.
  • Loading branch information
MasterJH5574 authored Feb 10, 2025
1 parent ca927c7 commit 16fe3f5
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion python/mlc_llm/support/auto_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@

FOUND = green("Found")
NOT_FOUND = red("Not found")
AUTO_DETECT_DEVICES = ["cpu", "cuda", "rocm", "metal", "vulkan", "opencl"]
AUTO_DETECT_DEVICES = ["cuda", "rocm", "metal", "vulkan", "opencl", "cpu"]
_RESULT_CACHE: Dict[str, bool] = {}


Expand Down Expand Up @@ -78,6 +78,8 @@ def _device_exists(device: Device) -> bool:
for i in subproc_outputs[0].split(","):
logger.info("%s device: %s:%s", FOUND, device_type, i)
_RESULT_CACHE[f"{device_type}:{i}"] = True
if device.device_type == Device.kDLCPU:
break
else:
logger.error(
"GPU device detection failed. Please report this issue with the output of command: %s",
Expand Down

0 comments on commit 16fe3f5

Please sign in to comment.