|
| 1 | +""" |
| 2 | +extract factors the build is dependent on: |
| 3 | +[X] compute capability |
| 4 | + [ ] TODO: Q - What if we have multiple GPUs of different makes? |
| 5 | +- CUDA version |
| 6 | +- Software: |
| 7 | + - CPU-only: only CPU quantization functions (no optimizer, no matrix multiple) |
| 8 | + - CuBLAS-LT: full-build 8-bit optimizer |
| 9 | + - no CuBLAS-LT: no 8-bit matrix multiplication (`nomatmul`) |
| 10 | +
|
| 11 | +evaluation: |
| 12 | + - if paths faulty, return meaningful error |
| 13 | + - else: |
| 14 | + - determine CUDA version |
| 15 | + - determine capabilities |
| 16 | + - based on that set the default path |
| 17 | +""" |
| 18 | + |
1 | 19 | import ctypes as ct |
2 | | -from warnings import warn |
| 20 | +import logging |
| 21 | +import os |
| 22 | +from pathlib import Path |
3 | 23 |
|
4 | 24 | import torch |
5 | 25 |
|
6 | | -from bitsandbytes.cuda_setup.main import CUDASetup |
| 26 | +from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX, PACKAGE_DIR |
| 27 | +from bitsandbytes.cuda_specs import CUDASpecs, get_cuda_specs |
| 28 | + |
| 29 | +logger = logging.getLogger(__name__) |
| 30 | + |
| 31 | + |
| 32 | +def get_cuda_bnb_library_path(cuda_specs: CUDASpecs) -> Path: |
| 33 | + """ |
| 34 | + Get the disk path to the CUDA BNB native library specified by the |
| 35 | + given CUDA specs, taking into account the `BNB_CUDA_VERSION` override environment variable. |
| 36 | +
|
| 37 | + The library is not guaranteed to exist at the returned path. |
| 38 | + """ |
| 39 | + library_name = f"libbitsandbytes_cuda{cuda_specs.cuda_version_string}" |
| 40 | + if not cuda_specs.has_cublaslt: |
| 41 | + # if not has_cublaslt (CC < 7.5), then we have to choose _nocublaslt |
| 42 | + library_name += "_nocublaslt" |
| 43 | + library_name = f"{library_name}{DYNAMIC_LIBRARY_SUFFIX}" |
| 44 | + |
| 45 | + override_value = os.environ.get("BNB_CUDA_VERSION") |
| 46 | + if override_value: |
| 47 | + library_name_stem, _, library_name_ext = library_name.rpartition(".") |
| 48 | + # `library_name_stem` will now be e.g. `libbitsandbytes_cuda118`; |
| 49 | + # let's remove any trailing numbers: |
| 50 | + library_name_stem = library_name_stem.rstrip("0123456789") |
| 51 | + # `library_name_stem` will now be e.g. `libbitsandbytes_cuda`; |
| 52 | + # let's tack the new version number and the original extension back on. |
| 53 | + library_name = f"{library_name_stem}{override_value}.{library_name_ext}" |
| 54 | + logger.warning( |
| 55 | + f"WARNING: BNB_CUDA_VERSION={override_value} environment variable detected; loading {library_name}.\n" |
| 56 | + "This can be used to load a bitsandbytes version that is different from the PyTorch CUDA version.\n" |
| 57 | + "If this was unintended set the BNB_CUDA_VERSION variable to an empty string: export BNB_CUDA_VERSION=\n" |
| 58 | + "If you use the manual override make sure the right libcudart.so is in your LD_LIBRARY_PATH\n" |
| 59 | + "For example by adding the following to your .bashrc: export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:<path_to_cuda_dir/lib64\n" |
| 60 | + ) |
| 61 | + |
| 62 | + return PACKAGE_DIR / library_name |
| 63 | + |
| 64 | + |
| 65 | +class BNBNativeLibrary: |
| 66 | + _lib: ct.CDLL |
| 67 | + compiled_with_cuda = False |
| 68 | + |
| 69 | + def __init__(self, lib: ct.CDLL): |
| 70 | + self._lib = lib |
| 71 | + |
| 72 | + def __getattr__(self, item): |
| 73 | + return getattr(self._lib, item) |
| 74 | + |
| 75 | + |
| 76 | +class CudaBNBNativeLibrary(BNBNativeLibrary): |
| 77 | + compiled_with_cuda = True |
| 78 | + |
| 79 | + def __init__(self, lib: ct.CDLL): |
| 80 | + super().__init__(lib) |
| 81 | + lib.get_context.restype = ct.c_void_p |
| 82 | + lib.get_cusparse.restype = ct.c_void_p |
| 83 | + lib.cget_managed_ptr.restype = ct.c_void_p |
| 84 | + |
| 85 | + |
| 86 | +def get_native_library() -> BNBNativeLibrary: |
| 87 | + binary_path = PACKAGE_DIR / f"libbitsandbytes_cpu{DYNAMIC_LIBRARY_SUFFIX}" |
| 88 | + cuda_specs = get_cuda_specs() |
| 89 | + if cuda_specs: |
| 90 | + cuda_binary_path = get_cuda_bnb_library_path(cuda_specs) |
| 91 | + if cuda_binary_path.exists(): |
| 92 | + binary_path = cuda_binary_path |
| 93 | + else: |
| 94 | + logger.warning("Could not find the bitsandbytes CUDA binary at %r", cuda_binary_path) |
| 95 | + logger.debug(f"Loading bitsandbytes native library from: {binary_path}") |
| 96 | + dll = ct.cdll.LoadLibrary(str(binary_path)) |
| 97 | + |
| 98 | + if hasattr(dll, "get_context"): # only a CUDA-built library exposes this |
| 99 | + return CudaBNBNativeLibrary(dll) |
| 100 | + |
| 101 | + logger.warning( |
| 102 | + "The installed version of bitsandbytes was compiled without GPU support. " |
| 103 | + "8-bit optimizers, 8-bit multiplication, and GPU quantization are unavailable." |
| 104 | + ) |
| 105 | + return BNBNativeLibrary(dll) |
7 | 106 |
|
8 | | -setup = CUDASetup.get_instance() |
9 | | -if setup.initialized != True: |
10 | | - setup.run_cuda_setup() |
11 | 107 |
|
12 | | -lib = setup.lib |
13 | 108 | try: |
14 | | - if lib is None and torch.cuda.is_available(): |
15 | | - CUDASetup.get_instance().generate_instructions() |
16 | | - CUDASetup.get_instance().print_log_stack() |
17 | | - raise RuntimeError(''' |
18 | | - CUDA Setup failed despite GPU being available. Please run the following command to get more information: |
19 | | -
|
20 | | - python -m bitsandbytes |
21 | | -
|
22 | | - Inspect the output of the command and see if you can locate CUDA libraries. You might need to add them |
23 | | - to your LD_LIBRARY_PATH. If you suspect a bug, please take the information from python -m bitsandbytes |
24 | | - and open an issue at: https://github.com/TimDettmers/bitsandbytes/issues''') |
25 | | - _ = lib.cadam32bit_grad_fp32 # runs on an error if the library could not be found -> COMPILED_WITH_CUDA=False |
26 | | - lib.get_context.restype = ct.c_void_p |
27 | | - lib.get_cusparse.restype = ct.c_void_p |
28 | | - lib.cget_managed_ptr.restype = ct.c_void_p |
29 | | - COMPILED_WITH_CUDA = True |
30 | | -except AttributeError as ex: |
31 | | - warn("The installed version of bitsandbytes was compiled without GPU support. " |
32 | | - "8-bit optimizers, 8-bit multiplication, and GPU quantization are unavailable.") |
33 | | - COMPILED_WITH_CUDA = False |
34 | | - print(str(ex)) |
35 | | - |
36 | | - |
37 | | -# print the setup details after checking for errors so we do not print twice |
38 | | -#if 'BITSANDBYTES_NOWELCOME' not in os.environ or str(os.environ['BITSANDBYTES_NOWELCOME']) == '0': |
39 | | - #setup.print_log_stack() |
| 109 | + lib = get_native_library() |
| 110 | +except Exception as e: |
| 111 | + lib = None |
| 112 | + logger.error(f"Could not load bitsandbytes native library: {e}", exc_info=True) |
| 113 | + if torch.cuda.is_available(): |
| 114 | + logger.warning( |
| 115 | + """ |
| 116 | +CUDA Setup failed despite CUDA being available. Please run the following command to get more information: |
| 117 | +
|
| 118 | +python -m bitsandbytes |
| 119 | +
|
| 120 | +Inspect the output of the command and see if you can locate CUDA libraries. You might need to add them |
| 121 | +to your LD_LIBRARY_PATH. If you suspect a bug, please take the information from python -m bitsandbytes |
| 122 | +and open an issue at: https://github.com/TimDettmers/bitsandbytes/issues |
| 123 | +""" |
| 124 | + ) |
0 commit comments