137 lines
6.4 KiB
Python
137 lines
6.4 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
|
|
import logging
|
|
UNSLOTH_ENABLE_LOGGING = os.environ.get("UNSLOTH_ENABLE_LOGGING", "0") == "1"
|
|
|
|
# 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())
|
|
pass
|
|
|
|
# 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
|
|
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)}")
|
|
pass
|
|
|
|
# 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)}")
|
|
pass
|
|
|
|
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
|
|
pass
|