From 57f044945f25d90d1b434014b2719ba6b06fdc44 Mon Sep 17 00:00:00 2001 From: zhaoyang-star Date: Fri, 23 Feb 2024 06:25:07 +0800 Subject: [PATCH] Fix nvcc not found in vlm-openai image (#2781) --- vllm/config.py | 2 +- vllm/utils.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/vllm/config.py b/vllm/config.py index 0b8a2a27f6d43..bd0dc89b585f7 100644 --- a/vllm/config.py +++ b/vllm/config.py @@ -319,7 +319,7 @@ def _verify_cache_dtype(self) -> None: pass elif self.cache_dtype == "fp8_e5m2": nvcc_cuda_version = get_nvcc_cuda_version() - if nvcc_cuda_version < Version("11.8"): + if nvcc_cuda_version and nvcc_cuda_version < Version("11.8"): raise ValueError( "FP8 is not supported when cuda version is lower than 11.8." ) diff --git a/vllm/utils.py b/vllm/utils.py index 6206879929061..8ca95e148eb39 100644 --- a/vllm/utils.py +++ b/vllm/utils.py @@ -181,13 +181,18 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None: os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, device_ids)) -def get_nvcc_cuda_version() -> Version: +def get_nvcc_cuda_version() -> Optional[Version]: cuda_home = os.environ.get('CUDA_HOME') if not cuda_home: cuda_home = '/usr/local/cuda' - logger.info( - f'CUDA_HOME is not found in the environment. Using {cuda_home} as CUDA_HOME.' - ) + if os.path.isfile(cuda_home + '/bin/nvcc'): + logger.info( + f'CUDA_HOME is not found in the environment. Using {cuda_home} as CUDA_HOME.' + ) + else: + logger.warning( + f'Not found nvcc in {cuda_home}. Skip cuda version check!') + return None nvcc_output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True) output = nvcc_output.split()