Repository: unslothai/unsloth
File: unsloth/__init__.py (Python)

Code before:
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
pass

try:
    import torch
except:
    raise ImportError("Pytorch is not installed. Go to https://pytorch.org/.\n"\
                      "We have some installation instructions on our Github page.")

# We support torch 2.1 and 2.1.1
# Fixes https://github.com/unslothai/unsloth/issues/38
torch_version = torch.__version__.split(".")
major_torch, minor_torch = torch_version[0], torch_version[1]
major_torch, minor_torch = int(major_torch), int(minor_torch)
if (major_torch != 2):# or (major_torch == 2 and minor_torch < 1):
    raise ImportError("Unsloth only supports Pytorch 2.1 for now. Please update your Pytorch to 2.1.\n"\
                      "We have some installation instructions on our Github page.")


# Try loading bitsandbytes and triton
import bitsandbytes as bnb
import triton
triton_version = triton.__version__.split(".")
triton_major = int(triton_version[0])
if triton_major >= 3:
    from triton.backends.nvidia.driver import libcuda_dirs
else:
    from triton.common.build import libcuda_dirs

import os
import re
import numpy as np
import subprocess

try:
    cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
    libcuda_dirs()
except:
    warnings.warn(
        "Unsloth: Running `ldconfig /usr/lib64-nvidia` to link CUDA."\
    )

    if os.path.exists("/usr/lib64-nvidia"):
        os.system("ldconfig /usr/lib64-nvidia")
    elif os.path.exists("/usr/local"):
        # Sometimes bitsandbytes cannot be linked properly in Runpod for example
        possible_cudas = subprocess.check_output(["ls", "-al", "/usr/local"]).decode("utf-8").split("\n")
        find_cuda = re.compile(r"[\s](cuda\-[\d\.]{2,})$")
        possible_cudas = [find_cuda.search(x) for x in possible_cudas]
        possible_cudas = [x.group(1) for x in possible_cudas if x is not None]


Review comment (suggestion):
This didn't work for me with `pytorch-triton==3.0.0+901819d2b6`:

```python
    import unsloth
  File "/home/carmocca/git/unsloth/unsloth/__init__.py", line 64, in <module>
    from triton.backends.nvidia.driver import libcuda_dirs
ImportError: cannot import name 'libcuda_dirs' from 'triton.backends.nvidia.driver' (/home/carmocca/git/nightly-venv/lib/python3.10/site-packages/triton/backends/nvidia/driver.py)
```

Also, you would want to update L98 which also imports this once there' a working solution

Code after:
        )
        os.environ["CUDA_VISIBLE_DEVICES"] = str(first_id)
else:
    # warnings.warn("Unsloth: 'CUDA_VISIBLE_DEVICES' is not set. We shall set it ourselves.")
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
pass

# Reduce VRAM usage by reducing fragmentation
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

try:
    import torch
except:
    raise ImportError("Pytorch is not installed. Go to https://pytorch.org/.\n"\
                      "We have some installation instructions on our Github page.")
pass

# Fix up is_bf16_supported https://github.com/unslothai/unsloth/issues/504
major_version, minor_version = torch.cuda.get_device_capability()
SUPPORTS_BFLOAT16 = (major_version >= 8)
def is_bf16_supported(): return SUPPORTS_BFLOAT16
torch.cuda.is_bf16_supported = is_bf16_supported

# We support Pytorch 2
# Fixes https://github.com/unslothai/unsloth/issues/38
torch_version = torch.__version__.split(".")
major_torch, minor_torch = torch_version[0], torch_version[1]
major_torch, minor_torch = int(major_torch), int(minor_torch)
if (major_torch < 2):
    raise ImportError("Unsloth only supports Pytorch 2 for now. Please update your Pytorch to 2.1.\n"\
                      "We have some installation instructions on our Github page.")
elif (major_torch == 2) and (minor_torch < 2):
    # Disable expandable_segments
    del os.environ["PYTORCH_CUDA_ALLOC_CONF"]
pass


# Try loading bitsandbytes and triton
import bitsandbytes as bnb

import triton
libcuda_dirs = lambda: None
if Version(triton.__version__) >= Version("3.0.0"):
    try: from triton.backends.nvidia.driver import libcuda_dirs
    except: pass
else: from triton.common.build import libcuda_dirs

import os
import re
import numpy as np