317 lines
13 KiB
Python
317 lines
13 KiB
Python
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import importlib.util
|
|
from pathlib import Path
|
|
from importlib.metadata import version as importlib_version
|
|
from packaging.version import Version as TrueVersion
|
|
import re
|
|
import logging
|
|
|
|
UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
|
|
|
|
def Version(version):
|
|
try:
|
|
new_version = str(version)
|
|
new_version = re.match(r"[0-9\.]{1,}", new_version)
|
|
if new_version is None:
|
|
raise Exception(str(e))
|
|
new_version = new_version.group(0).rstrip(".")
|
|
if new_version != version:
|
|
new_version += ".1" # Add .1 for dev / alpha / beta / rc
|
|
return TrueVersion(new_version)
|
|
except:
|
|
from inspect import getframeinfo, stack
|
|
caller = getframeinfo(stack()[1][0])
|
|
raise RuntimeError(
|
|
f"Unsloth: Could not get version for `{version}`\n"\
|
|
f"File name = [{caller.filename}] Line number = [{caller.lineno}]"
|
|
)
|
|
|
|
|
|
# Ignore logging messages
|
|
class HideLoggingMessage(logging.Filter):
|
|
__slots__ = ("text",)
|
|
|
|
def __init__(self, text):
|
|
self.text = text
|
|
|
|
def filter(self, x):
|
|
return not (self.text in x.getMessage())
|
|
|
|
|
|
# Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'
|
|
# MUST do this at the start primarily due to tensorflow causing issues
|
|
def fix_message_factory_issue():
|
|
try:
|
|
import google.protobuf.message_factory
|
|
|
|
class MessageFactory:
|
|
def CreatePrototype(self, *args, **kwargs):
|
|
return
|
|
|
|
def GetMessages(self, *args, **kwargs):
|
|
return
|
|
|
|
def GetPrototype(self, *args, **kwargs):
|
|
return
|
|
|
|
if not hasattr(google.protobuf.message_factory, "MessageFactory"):
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
|
google.protobuf.message_factory.MessageFactory = MessageFactory
|
|
elif (
|
|
hasattr(google.protobuf.message_factory, "MessageFactory")
|
|
and not hasattr(
|
|
google.protobuf.message_factory.MessageFactory, "GetPrototype"
|
|
)
|
|
and not hasattr(google.protobuf.message_factory, "GetMessageClass")
|
|
):
|
|
google.protobuf.message_factory.MessageFactory = MessageFactory
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print("Unsloth: Patching protobuf.MessageFactory as it doesn't exist")
|
|
elif (
|
|
hasattr(google.protobuf.message_factory, "MessageFactory")
|
|
and not hasattr(
|
|
google.protobuf.message_factory.MessageFactory, "GetPrototype"
|
|
)
|
|
and hasattr(google.protobuf.message_factory, "GetMessageClass")
|
|
):
|
|
GetMessageClass = google.protobuf.message_factory.GetMessageClass
|
|
|
|
def GetPrototype(self, descriptor):
|
|
return GetMessageClass(descriptor)
|
|
|
|
google.protobuf.message_factory.MessageFactory.GetPrototype = GetPrototype
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print("Unsloth: Patching protobuf.MessageFactory.GetPrototype")
|
|
pass
|
|
except:
|
|
pass
|
|
|
|
|
|
# Fix Xformers performance issues since 0.0.25
|
|
def fix_xformers_performance_issue():
|
|
if importlib.util.find_spec("xformers") is None:
|
|
return
|
|
xformers_version = importlib_version("xformers")
|
|
if Version(xformers_version) < Version("0.0.29"):
|
|
xformers_location = importlib.util.find_spec("xformers").origin
|
|
xformers_location = os.path.split(xformers_location)[0]
|
|
cutlass = Path(xformers_location) / "ops" / "fmha" / "cutlass.py"
|
|
try:
|
|
if cutlass.exists():
|
|
with open(cutlass, "r+", encoding = "utf-8") as f:
|
|
text = f.read()
|
|
# See https://github.com/facebookresearch/xformers/issues/1176#issuecomment-2545829591
|
|
if "num_splits_key=-1," in text:
|
|
text = text.replace(
|
|
"num_splits_key=-1,",
|
|
"num_splits_key=None,",
|
|
)
|
|
f.seek(0)
|
|
f.write(text)
|
|
f.truncate()
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print(
|
|
"Unsloth: Patching Xformers to fix some performance issues."
|
|
)
|
|
except Exception as e:
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print(f"Unsloth: Failed patching Xformers with error = {str(e)}")
|
|
|
|
|
|
# ValueError: 'aimv2' is already used by a Transformers config, pick another name.
|
|
def fix_vllm_aimv2_issue():
|
|
if importlib.util.find_spec("vllm") is None:
|
|
return
|
|
vllm_version = importlib_version("vllm")
|
|
if Version(vllm_version) < Version("0.10.1"):
|
|
vllm_version = importlib.util.find_spec("vllm").origin
|
|
vllm_version = os.path.split(vllm_version)[0]
|
|
ovis_config = Path(vllm_version) / "transformers_utils" / "configs" / "ovis.py"
|
|
try:
|
|
if ovis_config.exists():
|
|
with open(ovis_config, "r+", encoding = "utf-8") as f:
|
|
text = f.read()
|
|
# See https://github.com/vllm-project/vllm-ascend/issues/2046
|
|
if 'AutoConfig.register("aimv2", AIMv2Config)' in text:
|
|
text = text.replace(
|
|
'AutoConfig.register("aimv2", AIMv2Config)',
|
|
"",
|
|
)
|
|
text = text.replace(
|
|
"""backbone_config.pop('model_type')
|
|
backbone_config = AutoConfig.for_model(model_type,
|
|
**backbone_config)""",
|
|
"""if model_type != "aimv2":
|
|
backbone_config.pop('model_type')
|
|
backbone_config = AutoConfig.for_model(model_type, **backbone_config)
|
|
else:
|
|
backbone_config = AIMv2Config(**backbone_config)""",
|
|
)
|
|
f.seek(0)
|
|
f.write(text)
|
|
f.truncate()
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print(
|
|
"Unsloth: Patching vLLM to fix `'aimv2' is already used by a Transformers config, pick another name.`"
|
|
)
|
|
except Exception as e:
|
|
if UNSLOTH_ENABLE_LOGGING:
|
|
print(f"Unsloth: Failed patching vLLM with error = {str(e)}")
|
|
|
|
|
|
def fix_vllm_guided_decoding_params():
|
|
if importlib.util.find_spec("vllm") is None:
|
|
return
|
|
# GuidedDecodingParmas is renamed to StructuredOutputsParams in vLLM
|
|
# https://github.com/vllm-project/vllm/pull/22772/files
|
|
# trl still wants to use GuidedDecodingParams. This is a temporary patch till trl updates
|
|
import vllm
|
|
|
|
try:
|
|
from vllm.sampling_params import GuidedDecodingParams
|
|
except ImportError:
|
|
vllm.sampling_params.GuidedDecodingParams = (
|
|
vllm.sampling_params.StructuredOutputsParams
|
|
)
|
|
|
|
|
|
def ignore_logger_messages():
|
|
# Ignore Environment variable `HF_TOKEN` is set
|
|
try:
|
|
from huggingface_hub._login import logger as huggingface_hub_logger
|
|
|
|
huggingface_hub_logger.addFilter(HideLoggingMessage("`HF_TOKEN`"))
|
|
del huggingface_hub_logger
|
|
except:
|
|
pass
|
|
|
|
|
|
def patch_ipykernel_hf_xet():
|
|
# HF-XET == 1.1.10 and ipykernel == 7.0.0 / 7.0.1 causes issues
|
|
# See https://github.com/huggingface/xet-core/issues/526
|
|
# 2025-10-13T20:37:33.028737Z ERROR Python exception updating progress:, error: PyErr { type: <class 'LookupError'>, value: LookupError(<ContextVar name='shell_parent' at 0x7535b4cebd80>), traceback: Some(<traceback object at 0x753408489f40>) }, caller: "src/progress_update.rs:313"
|
|
# at /home/runner/work/xet-core/xet-core/error_printer/src/lib.rs:28
|
|
if importlib.util.find_spec("hf_xet") is None:
|
|
return
|
|
if importlib.util.find_spec("ipykernel") is None:
|
|
return
|
|
if importlib.util.find_spec("huggingface_hub") is None:
|
|
return
|
|
|
|
ipykernel_version = Version(importlib_version("ipykernel"))
|
|
if (
|
|
(Version(importlib_version("hf_xet")) == Version("1.1.10"))
|
|
and (
|
|
(ipykernel_version == Version("7.0.0"))
|
|
or (
|
|
ipykernel_version == Version("7.0.1")
|
|
) # 7.0.1 seems to also break with LookupError: <ContextVar name='shell_parent' at 0x7a9775143ec0>
|
|
)
|
|
):
|
|
print(
|
|
"#### Unsloth: `hf_xet==1.1.10` and `ipykernel==7.0.0` or `ipykernel==7.0.1` breaks progress bars. Using ASCII progress bars.\n"
|
|
"#### Unsloth: To re-enable progress bars, please upgrade to `ipykernel>=7.1.0` or wait for a fix to\n"
|
|
"https://github.com/huggingface/xet-core/issues/526"
|
|
)
|
|
from huggingface_hub.utils import disable_progress_bars
|
|
|
|
disable_progress_bars()
|
|
|
|
|
|
def patch_trackio():
|
|
# Set some environment variables to customize the Trackio dashboard for experiment tracking
|
|
# See https://github.com/unslothai/notebooks/pull/110
|
|
os.environ["TRACKIO_LOGO_LIGHT_URL"] = (
|
|
"https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20black%20text.png"
|
|
)
|
|
os.environ["TRACKIO_LOGO_DARK_URL"] = (
|
|
"https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20logo%20white%20text.png"
|
|
)
|
|
os.environ["TRACKIO_PLOT_ORDER"] = "train/reward"
|
|
|
|
|
|
def patch_datasets():
|
|
# Datasets 4.4.0 and 4.4.1 weirdly have some weird `_thread.RLock_recursion_count` issues
|
|
if importlib.util.find_spec("datasets") is None:
|
|
return
|
|
|
|
datasets_version = Version(importlib_version("datasets"))
|
|
if (datasets_version <= Version("4.5.0")) and (
|
|
datasets_version >= Version("4.4.0")
|
|
):
|
|
raise NotImplementedError(
|
|
f"#### Unsloth: Using `datasets = {str(datasets_version)}` will cause recursion errors.\n"
|
|
"Please downgrade datasets to `datasets==4.3.0"
|
|
)
|
|
|
|
|
|
def check_fbgemm_gpu_version():
|
|
if importlib.util.find_spec("fbgemm_gpu") is None:
|
|
return
|
|
try:
|
|
fbgemm_gpu_version = importlib_version("fbgemm_gpu_genai")
|
|
except:
|
|
return
|
|
# We noticed some SegFault or bad alloc errors on lower versions of fbgemm_gpu.
|
|
if Version(fbgemm_gpu_version) < Version("1.4.0"):
|
|
raise ImportError(
|
|
f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected. It might cause unexpected issues like segmentation faults. Please uninstall the current one by doing `pip uninstall fbgemm-gpu` && `pip install fbgemm-gpu` to install fbgemm-gpu 1.4.0 or newer!"
|
|
)
|
|
elif UNSLOTH_ENABLE_LOGGING:
|
|
print(f"Unsloth: fbgemm_gpu_genai=={fbgemm_gpu_version} detected.")
|
|
|
|
|
|
def torchvision_compatibility_check():
|
|
if importlib.util.find_spec("torch") is None:
|
|
raise ImportError("Unsloth: torch not found. Please install torch first.")
|
|
if importlib.util.find_spec("torchvision") is None:
|
|
return
|
|
torch_version = importlib_version("torch")
|
|
torchvision_version = importlib_version("torchvision")
|
|
|
|
# Torch version -> minimum required torchvision version
|
|
# See https://pytorch.org/get-started/previous-versions/
|
|
TORCH_TORCHVISION_COMPAT = [
|
|
("2.9.0", "0.24.0"),
|
|
("2.8.0", "0.23.0"),
|
|
("2.7.0", "0.22.0"),
|
|
("2.6.0", "0.21.0"),
|
|
("2.5.0", "0.20.0"),
|
|
("2.4.0", "0.19.0"),
|
|
]
|
|
|
|
required_torchvision = None
|
|
for min_torch, min_torchvision in TORCH_TORCHVISION_COMPAT:
|
|
if Version(torch_version) >= Version(min_torch):
|
|
required_torchvision = min_torchvision
|
|
break
|
|
|
|
if required_torchvision is None:
|
|
# Torch version not in compatibility table, skip check
|
|
return
|
|
|
|
if Version(torchvision_version) < Version(required_torchvision):
|
|
raise ImportError(
|
|
f"Unsloth: torch=={torch_version} requires torchvision>={required_torchvision}, "
|
|
f"but found torchvision=={torchvision_version}. "
|
|
f"Please refer to https://pytorch.org/get-started/previous-versions/ for more information."
|
|
)
|
|
elif UNSLOTH_ENABLE_LOGGING:
|
|
print(
|
|
f"Unsloth: torch=={torch_version} and torchvision=={torchvision_version} are compatible."
|
|
)
|