|
|
|
|
|
|
|
import os |
|
import sys |
|
from warnings import warn |
|
|
|
import torch |
|
|
|
HEADER_WIDTH = 60 |
|
|
|
|
|
def print_header( |
|
txt: str, width: int = HEADER_WIDTH, filler: str = "+" |
|
) -> None: |
|
txt = f" {txt} " if txt else "" |
|
print(txt.center(width, filler)) |
|
|
|
|
|
def print_debug_info() -> None: |
|
print( |
|
"\nAbove we output some debug information. Please provide this info when " |
|
f"creating an issue via {PACKAGE_GITHUB_URL}/issues/new/choose ...\n" |
|
) |
|
|
|
|
|
print_header("") |
|
print_header("DEBUG INFORMATION") |
|
print_header("") |
|
print() |
|
|
|
|
|
from . import COMPILED_WITH_CUDA, PACKAGE_GITHUB_URL |
|
from .cuda_setup.main import get_compute_capabilities, get_cuda_lib_handle |
|
from .cuda_setup.env_vars import to_be_ignored |
|
|
|
print_header("POTENTIALLY LIBRARY-PATH-LIKE ENV VARS") |
|
for k, v in os.environ.items(): |
|
if "/" in v and not to_be_ignored(k, v): |
|
print(f"'{k}': '{v}'") |
|
print_header("") |
|
|
|
print( |
|
"\nWARNING: Please be sure to sanitize sensible info from any such env vars!\n" |
|
) |
|
|
|
print_header("OTHER") |
|
print(f"{COMPILED_WITH_CUDA = }") |
|
cuda = get_cuda_lib_handle() |
|
print(f"COMPUTE_CAPABILITIES_PER_GPU = {get_compute_capabilities(cuda)}") |
|
print_header("") |
|
print_header("DEBUG INFO END") |
|
print_header("") |
|
print( |
|
""" |
|
Running a quick check that: |
|
+ library is importable |
|
+ CUDA function is callable |
|
""" |
|
) |
|
|
|
try: |
|
from bitsandbytes.optim import Adam |
|
|
|
p = torch.nn.Parameter(torch.rand(10, 10).cuda()) |
|
a = torch.rand(10, 10).cuda() |
|
|
|
p1 = p.data.sum().item() |
|
|
|
adam = Adam([p]) |
|
|
|
out = a * p |
|
loss = out.sum() |
|
loss.backward() |
|
adam.step() |
|
|
|
p2 = p.data.sum().item() |
|
|
|
assert p1 != p2 |
|
print("SUCCESS!") |
|
print("Installation was successful!") |
|
sys.exit(0) |
|
|
|
except ImportError: |
|
print() |
|
warn( |
|
f"WARNING: {__package__} is currently running as CPU-only!\n" |
|
"Therefore, 8-bit optimizers and GPU quantization are unavailable.\n\n" |
|
f"If you think that this is so erroneously,\nplease report an issue!" |
|
) |
|
print_debug_info() |
|
sys.exit(0) |
|
except Exception as e: |
|
print(e) |
|
print_debug_info() |
|
sys.exit(1) |
|
|
