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 2025-10-18 17:32:01 -07:00
commit c394ff06de
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

@ -69,7 +69,28 @@ except Exception as exception:
raise exception
pass
from .device_type import (
# Check for unsloth_zoo
try:
unsloth_zoo_version = importlib_version("unsloth_zoo")
if Version(unsloth_zoo_version) < Version("2025.10.4"):
print(
"Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n"\
"Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`"
)
# if os.environ.get("UNSLOTH_DISABLE_AUTO_UPDATES", "0") == "0":
# try:
# os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo")
# except:
# try:
# os.system("pip install --upgrade --no-cache-dir --no-deps --user unsloth_zoo")
# except:
# raise ImportError("Unsloth: Please update unsloth_zoo via `pip install --upgrade --no-cache-dir --no-deps unsloth_zoo`")
import unsloth_zoo
except:
raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth_zoo`")
pass
from unsloth_zoo.device_type import (
is_hip,
get_device_type,
DEVICE_TYPE,
@ -84,6 +105,7 @@ if (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY",
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \
"expandable_segments:True,"\
"roundup_power2_divisions:[32:256,64:128,256:64,>:32]"
os.environ["PYTORCH_HIP_ALLOC_CONF"] = "expandable_segments:True"
elif (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY", "0")=="1") and \
("expandable_segments:True" in os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")):
warnings.warn(
@ -96,6 +118,11 @@ elif (DEVICE_TYPE in ("cuda", "hip")) and (os.environ.get("UNSLOTH_VLLM_STANDBY"
"",
os.environ["PYTORCH_CUDA_ALLOC_CONF"],
)
os.environ["PYTORCH_HIP_ALLOC_CONF"] = re.sub(
r"expandable\_segments\:True\,?",
"",
os.environ["PYTORCH_HIP_ALLOC_CONF"],
)
pass
# We support Pytorch 2
# Fixes https://github.com/unslothai/unsloth/issues/38
@ -108,6 +135,7 @@ if (major_torch < 2):
elif (major_torch == 2) and (minor_torch < 2):
# Disable expandable_segments
del os.environ["PYTORCH_CUDA_ALLOC_CONF"]
del os.environ["PYTORCH_HIP_ALLOC_CONF"]
pass
# CCE fails on Torch 2.8 and above
@ -225,27 +253,6 @@ elif DEVICE_TYPE == "xpu":
# TODO: check triton for intel installed properly.
pass
# Check for unsloth_zoo
try:
unsloth_zoo_version = importlib_version("unsloth_zoo")
if Version(unsloth_zoo_version) < Version("2025.10.2"):
print(
"Unsloth: Please update Unsloth and Unsloth-Zoo to the latest version!\n"\
"Do this via `pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo`"
)
# if os.environ.get("UNSLOTH_DISABLE_AUTO_UPDATES", "0") == "0":
# try:
# os.system("pip install --upgrade --no-cache-dir --no-deps unsloth_zoo")
# except:
# try:
# os.system("pip install --upgrade --no-cache-dir --no-deps --user unsloth_zoo")
# except:
# raise ImportError("Unsloth: Please update unsloth_zoo via `pip install --upgrade --no-cache-dir --no-deps unsloth_zoo`")
import unsloth_zoo
except:
raise ImportError("Unsloth: Please install unsloth_zoo via `pip install unsloth_zoo`")
pass
from .models import *
from .models import __version__
from .save import *