From d3d08f20d690c576ce378bebad4c2d80fe2132a9 Mon Sep 17 00:00:00 2001 From: bojiang Date: Tue, 2 Jul 2024 14:49:55 +0800 Subject: [PATCH] add warning for nvidia compute_capability --- openllm_next/accelerator_spec.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/openllm_next/accelerator_spec.py b/openllm_next/accelerator_spec.py index fe1d6632..1168127c 100644 --- a/openllm_next/accelerator_spec.py +++ b/openllm_next/accelerator_spec.py @@ -78,6 +78,7 @@ def get_local_machine_spec(): from pynvml import ( nvmlDeviceGetCount, + nvmlDeviceGetCudaComputeCapability, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo, nvmlDeviceGetName, @@ -98,6 +99,14 @@ def get_local_machine_spec(): model=name, memory_size=math.ceil(int(memory_info.total) / 1024**3) ) ) + compute_capability = nvmlDeviceGetCudaComputeCapability(handle) + if compute_capability < (7, 5): + output( + f"GPU {name.decode()} with compute capability {compute_capability} " + "may not be supported, 7.5 or higher is recommended. check " + "https://developer.nvidia.com/cuda-gpus for more information", + style="yellow", + ) nvmlShutdown() return DeploymentTarget( accelerators=accelerators,