File size: 7,256 Bytes
ea35292 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | from collections.abc import Iterable, Iterator
import logging
import os
from pathlib import Path
import torch
from bitsandbytes.cextension import HIP_ENVIRONMENT, get_cuda_bnb_library_path
from bitsandbytes.cuda_specs import CUDASpecs
from bitsandbytes.diagnostics.utils import print_dedented
CUDART_PATH_PREFERRED_ENVVARS = ("CONDA_PREFIX", "LD_LIBRARY_PATH")
CUDART_PATH_IGNORED_ENVVARS = {
"DBUS_SESSION_BUS_ADDRESS", # hardware related
"GOOGLE_VM_CONFIG_LOCK_FILE", # GCP: requires elevated permissions, causing problems in VMs and Jupyter notebooks
"HOME", # Linux shell default
"LESSCLOSE",
"LESSOPEN", # related to the `less` command
"MAIL", # something related to emails
"OLDPWD",
"PATH", # this is for finding binaries, not libraries
"PWD", # PWD: this is how the shell keeps track of the current working dir
"SHELL", # binary for currently invoked shell
"SSH_AUTH_SOCK", # SSH stuff, therefore unrelated
"SSH_TTY",
"TMUX", # Terminal Multiplexer
"XDG_DATA_DIRS", # XDG: Desktop environment stuff
"XDG_GREETER_DATA_DIR", # XDG: Desktop environment stuff
"XDG_RUNTIME_DIR",
"_", # current Python interpreter
}
CUDA_RUNTIME_LIB_PATTERNS = (
(
"libamdhip64.so*", # Linux
"amdhip64*.dll", # Windows
)
if HIP_ENVIRONMENT
else (
"cudart64*.dll", # Windows
"libcudart*.so*", # libcudart.so, libcudart.so.11.0, libcudart.so.12.0, libcudart.so.12.1, libcudart.so.12.2 etc.
"nvcuda*.dll", # Windows
)
)
logger = logging.getLogger(__name__)
def find_cuda_libraries_in_path_list(paths_list_candidate: str) -> Iterable[Path]:
for dir_string in paths_list_candidate.split(os.pathsep):
if not dir_string:
continue
if os.sep not in dir_string:
continue
try:
dir = Path(dir_string)
try:
if not dir.exists():
logger.warning(f"The directory listed in your path is found to be non-existent: {dir}")
continue
except OSError: # Assume an esoteric error trying to poke at the directory
pass
for lib_pattern in CUDA_RUNTIME_LIB_PATTERNS:
for pth in dir.glob(lib_pattern):
if pth.is_file() and not pth.is_symlink():
yield pth
except (OSError, PermissionError):
pass
def is_relevant_candidate_env_var(env_var: str, value: str) -> bool:
return (
env_var in CUDART_PATH_PREFERRED_ENVVARS # is a preferred location
or (
os.sep in value # might contain a path
and env_var not in CUDART_PATH_IGNORED_ENVVARS # not ignored
and "CONDA" not in env_var # not another conda envvar
and "BASH_FUNC" not in env_var # not a bash function defined via envvar
and "\n" not in value # likely e.g. a script or something?
)
)
def get_potentially_lib_path_containing_env_vars() -> dict[str, str]:
return {env_var: value for env_var, value in os.environ.items() if is_relevant_candidate_env_var(env_var, value)}
def find_cudart_libraries() -> Iterator[Path]:
"""
Searches for a cuda installations, in the following order of priority:
1. active conda env
2. LD_LIBRARY_PATH
3. any other env vars, while ignoring those that
- are known to be unrelated
- don't contain the path separator `/`
If multiple libraries are found in part 3, we optimistically try one,
while giving a warning message.
"""
candidate_env_vars = get_potentially_lib_path_containing_env_vars()
for envvar in CUDART_PATH_PREFERRED_ENVVARS:
if envvar in candidate_env_vars:
directory = candidate_env_vars[envvar]
yield from find_cuda_libraries_in_path_list(directory)
candidate_env_vars.pop(envvar)
for env_var, value in candidate_env_vars.items():
yield from find_cuda_libraries_in_path_list(value)
def _print_cuda_diagnostics(cuda_specs: CUDASpecs) -> None:
print(
f"PyTorch settings found: CUDA_VERSION={cuda_specs.cuda_version_string}, "
f"Highest Compute Capability: {cuda_specs.highest_compute_capability}.",
)
binary_path = get_cuda_bnb_library_path(cuda_specs)
if not binary_path.exists():
print_dedented(
f"""
No compatible CUDA library found (tried: {binary_path.name}). You may need to compile from source:
https://huggingface.co/docs/bitsandbytes/main/en/installation#cuda-compile
""",
)
# 7.5 is the minimum CC for int8 tensor cores
if not cuda_specs.has_imma:
print_dedented(
"""
WARNING: Compute capability < 7.5 detected! Only slow 8-bit matmul is supported for your GPU!
If you run into issues with 8-bit matmul, you can try 4-bit quantization:
https://huggingface.co/blog/4bit-transformers-bitsandbytes
""",
)
def _print_hip_diagnostics(cuda_specs: CUDASpecs) -> None:
print(f"PyTorch settings found: ROCM_VERSION={cuda_specs.cuda_version_string}")
rocm_override = os.environ.get("BNB_ROCM_VERSION")
if rocm_override:
print(f"BNB_ROCM_VERSION override: {rocm_override}")
binary_path = get_cuda_bnb_library_path(cuda_specs)
if not binary_path.exists():
print_dedented(
f"""
No compatible ROCm library found (tried: {binary_path.name}). You may need to compile from source:
https://huggingface.co/docs/bitsandbytes/main/en/installation#rocm-compile
Use BNB_ROCM_VERSION to force a specific version if needed.
""",
)
hip_major, hip_minor = cuda_specs.cuda_version_tuple
if (hip_major, hip_minor) < (6, 1):
print_dedented(
"""
WARNING: bitsandbytes is fully supported only from ROCm 6.1.
""",
)
def print_diagnostics(cuda_specs: CUDASpecs) -> None:
if HIP_ENVIRONMENT:
_print_hip_diagnostics(cuda_specs)
else:
_print_cuda_diagnostics(cuda_specs)
def print_runtime_diagnostics() -> None:
backend = "ROCm" if HIP_ENVIRONMENT else "CUDA"
runtime_version = torch.version.hip if HIP_ENVIRONMENT else torch.version.cuda
override_var = "BNB_ROCM_VERSION" if HIP_ENVIRONMENT else "BNB_CUDA_VERSION"
override_example = "72" if HIP_ENVIRONMENT else "122"
cudart_paths = list(find_cudart_libraries())
if not cudart_paths:
print(f"{backend} SETUP: WARNING! {backend} runtime files not found in any environmental path.")
elif len(cudart_paths) > 1:
print_dedented(
f"""
Found duplicate {backend} runtime files (see below).
bitsandbytes will use PyTorch's {backend} runtime ({runtime_version}) and auto-select
the closest available library version. If you need to force a specific version,
set {override_var}, e.g.:
export {override_var}={override_example}
""",
)
for pth in cudart_paths:
print(f"* Found {backend} runtime at: {pth}")
|