ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update __init__.py

Browse source
This commit is contained in:
Daniel Han-Chen 2023-11-30 17:30:10 +11:00
commit b9b9fca05c
1 changed files with 29 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

51
unsloth/__init__.py
View file

@ -13,50 +13,57 @@
# limitations under the License.
__version__ = "2023.11"
import os
import warnings
import importlib
try:
import torch
import torch
except:
raise ImportError("Pytorch is not installed. Go to https://pytorch.org/.\n"\
"We have some installation instructions on our Github page.")
raise ImportError("Pytorch is not installed. Go to https://pytorch.org/.\n"\
"We have some installation instructions on our Github page.")
# We only support torch 2.1
major_torch, minor_torch, _ = torch.__version__.split(".")
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.")
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.")
# Currently only supports 1 GPU, or else seg faults will occur.
reload_package = False
n_gpus = torch.cuda.device_count()
if n_gpus == 0:
raise RuntimeError("Unsloth: Requires at least 1 GPU. Found 0.")
raise RuntimeError("Unsloth: Requires at least 1 GPU. Found 0.")
elif n_gpus > 1:
if "CUDA_VISIBLE_DEVICES" in os.environ:
device = os.environ["CUDA_VISIBLE_DEVICES"]
if not device.isdigit():
print(f"Unsloth: 'CUDA_VISIBLE_DEVICES' is currently {device} "\
"but we require 'CUDA_VISIBLE_DEVICES=0'\n"\
"We shall set it ourselves.")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
reload_package = True
else:
print("Unsloth: 'CUDA_VISIBLE_DEVICES' is not set. We shall set it ourselves.")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
reload_package = True
if "CUDA_VISIBLE_DEVICES" in os.environ:
device = os.environ["CUDA_VISIBLE_DEVICES"]
if not device.isdigit():
warnings.warn(
f"Unsloth: 'CUDA_VISIBLE_DEVICES' is currently {device} "\
"but we require 'CUDA_VISIBLE_DEVICES=0'\n"\
"We shall set it ourselves."
)
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
reload_package = True
else:
warnings.warn("Unsloth: 'CUDA_VISIBLE_DEVICES' is not set. We shall set it ourselves.")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
reload_package = True
pass
# Reload Pytorch with CUDA_VISIBLE_DEVICES
if reload_package:
import importlib
importlib.reload(torch)
importlib.reload(torch)
pass
# Try loading bitsandbytes
import bitsandbytes as bnb
try:
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
cdequantize_blockwise_fp32 = bnb.functional.lib.cdequantize_blockwise_fp32
except:
raise ImportError("CUDA is not linked properly. Try running `ldconfig /usr/lib64-nvidia` first.")
warnings.warn("CUDA is not linked properly. We shall run `ldconfig /usr/lib64-nvidia` to try to fix it.")
os.system("ldconfig /usr/lib64-nvidia")
importlib.reload(bnb)
pass
from .models import *