venv/lib/python3.11/site-packages/bitsandbytes/diagnostics/cuda.py

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}")