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

Torch 2.1

Browse source
This commit is contained in:
Daniel Han-Chen 2023-11-30 17:21:29 +11:00
commit 7166dce0ae
2 changed files with 39 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

39
unsloth/__init__.py
View file

@ -12,5 +12,44 @@
# See the License for the specific language governing permissions and
# limitations under the License.
__version__ = "2023.11"
import os
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 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.")
# 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.")
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
pass
# Reload Pytorch with CUDA_VISIBLE_DEVICES
if reload_package:
import importlib
importlib.reload(torch)
pass
from .models import *

29
unsloth/models/__init__.py
View file

@ -12,33 +12,4 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
import os
# 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.")
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
pass
# Reload Pytorch with CUDA_VISIBLE_DEVICES
if reload_package:
import importlib
importlib.reload(torch)
pass
from .llama import FastLlamaModel