Unverified Commit 57f04494 authored by zhaoyang-star's avatar zhaoyang-star Committed by GitHub
Browse files

Fix nvcc not found in vlm-openai image (#2781)

parent 4caf7044
...@@ -319,7 +319,7 @@ class CacheConfig: ...@@ -319,7 +319,7 @@ class CacheConfig:
pass pass
elif self.cache_dtype == "fp8_e5m2": elif self.cache_dtype == "fp8_e5m2":
nvcc_cuda_version = get_nvcc_cuda_version() 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( raise ValueError(
"FP8 is not supported when cuda version is lower than 11.8." "FP8 is not supported when cuda version is lower than 11.8."
) )
......
...@@ -181,13 +181,18 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None: ...@@ -181,13 +181,18 @@ def set_cuda_visible_devices(device_ids: List[int]) -> None:
os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, device_ids)) 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') cuda_home = os.environ.get('CUDA_HOME')
if not cuda_home: if not cuda_home:
cuda_home = '/usr/local/cuda' cuda_home = '/usr/local/cuda'
logger.info( if os.path.isfile(cuda_home + '/bin/nvcc'):
f'CUDA_HOME is not found in the environment. Using {cuda_home} as CUDA_HOME.' 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"], nvcc_output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"],
universal_newlines=True) universal_newlines=True)
output = nvcc_output.split() output = nvcc_output.split()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment