From 0ad16b7c1faa2f769d1775f35477d039cc3dd0c3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 8 Sep 2024 15:51:27 -0700
Subject: [PATCH] Update __init__.py

---
 unsloth/__init__.py | 93 +++++++++++++++++++++++----------------------
 1 file changed, 48 insertions(+), 45 deletions(-)

diff --git a/unsloth/__init__.py b/unsloth/__init__.py
index f6ed999530..e7db41ce2c 100644
--- a/unsloth/__init__.py
+++ b/unsloth/__init__.py
@@ -96,59 +96,62 @@ 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
+if "SPACE_AUTHOR_NAME" not in os.environ and "SPACE_REPO_NAME" not in os.environ:
 
-try:
-    cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
-    libcuda_dirs()
-except:
-    warnings.warn(
-        "Unsloth: Running `ldconfig /usr/lib64-nvidia` to link CUDA."\
-    )
+    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
 
-    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]
-
-        # Try linking cuda folder, or everything in local
-        if len(possible_cudas) == 0:
-            os.system(f"ldconfig /usr/local/")
-        else:
-            find_number = re.compile(r"([\d\.]{2,})")
-            latest_cuda = np.argsort([float(find_number.search(x).group(1)) for x in possible_cudas])[::-1][0]
-            latest_cuda = possible_cudas[latest_cuda]
-            os.system(f"ldconfig /usr/local/{latest_cuda}")
-    pass
-
-    importlib.reload(bnb)
-    importlib.reload(triton)
     try:
-        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
         cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
         libcuda_dirs()
     except:
         warnings.warn(
-            "Unsloth: CUDA is not linked properly.\n"\
-            "Try running `python -m bitsandbytes` then `python -m xformers.info`\n"\
-            "We tried running `ldconfig /usr/lib64-nvidia` ourselves, but it didn't work.\n"\
-            "You need to run in your terminal `sudo ldconfig /usr/lib64-nvidia` yourself, then import Unsloth.\n"\
-            "Also try `sudo ldconfig /usr/local/cuda-xx.x` - find the latest cuda version.\n"\
-            "Unsloth will still run for now, but maybe it might crash - let's hope it works!"
+            "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]
+
+            # Try linking cuda folder, or everything in local
+            if len(possible_cudas) == 0:
+                os.system(f"ldconfig /usr/local/")
+            else:
+                find_number = re.compile(r"([\d\.]{2,})")
+                latest_cuda = np.argsort([float(find_number.search(x).group(1)) for x in possible_cudas])[::-1][0]
+                latest_cuda = possible_cudas[latest_cuda]
+                os.system(f"ldconfig /usr/local/{latest_cuda}")
+        pass
+
+        importlib.reload(bnb)
+        importlib.reload(triton)
+        try:
+            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
+            cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
+            libcuda_dirs()
+        except:
+            warnings.warn(
+                "Unsloth: CUDA is not linked properly.\n"\
+                "Try running `python -m bitsandbytes` then `python -m xformers.info`\n"\
+                "We tried running `ldconfig /usr/lib64-nvidia` ourselves, but it didn't work.\n"\
+                "You need to run in your terminal `sudo ldconfig /usr/lib64-nvidia` yourself, then import Unsloth.\n"\
+                "Also try `sudo ldconfig /usr/local/cuda-xx.x` - find the latest cuda version.\n"\
+                "Unsloth will still run for now, but maybe it might crash - let's hope it works!"
+            )
+    pass
 pass
 
 from .models import *