mrq/bitsandbytes-rocm
1
0
Fork 1
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

fix: Get device's compute capability

Browse Source
This commit is contained in:
rapsealk 2023-04-28 11:18:54 +09:00
parent f5110265ff
commit 2b4cc256f6
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
bitsandbytes/cuda_setup/main.py
View File

@ -279,9 +279,7 @@ def determine_cuda_runtime_lib_path() -> Union[Path, None]:
def get_cuda_version(cuda, cudart_path):
if cuda is None: return None
version = torch._C._cuda_getCompiledVersion()
major = version//1000
minor = (version-(major*1000))//10
major, minor = map(int, torch.version.cuda.split("."))
if major < 11:
CUDASetup.get_instance().add_log_entry('CUDA SETUP: CUDA version lower than 11 are currently not supported for LLM.int8(). You will be only to use 8-bit optimizers and quantization routines!!')
@ -302,7 +300,9 @@ def get_cuda_lib_handle():
def get_compute_capabilities(cuda):
ccs = []
ccs.append(torch.version.cuda)
for i in range(torch.cuda.device_count()):
cc_major, cc_minor = torch.cuda.get_device_capability(torch.cuda.device(i))
ccs.append(f"{cc_major}.{cc_minor}")
return ccs