54d8983401**Summary:** The existing QAT + LoRA path only applied fake quantization to the original slow path, but the default is the fast path that calls unsloth's fast LoRA primitives. This commit integrates fake quantization into these fast primitives as well, and add unit tests to assert that fake quantization is actually taking place. **Test Plan:** Unit tests: ``` pytest tests/utils/test_qat.py ``` End-to-end test: https://gist.github.com/andrewor14/6360dd69b5784c71c46e80c14f53e6b6 Full fine-tuning Llama3.1-8B with and without QAT + LoRA on yahma/alpaca-cleaned for 1 epoch: - Batch size = 8 (no grad accum) - Learning rate = 2e-4 - Quantization scheme = int4 weight only (with bf16 activations) Wikitext perplexity: - Baseline = int4 quantized model finetuned without QAT - QAT int4 quantized model (with this PR) achieved 33% lower perplexity than the int4 baseline - QAT int4 quantized model without this PR was worse than the int4 baseline ``` ==> unsloth_model_lora_baseline_output/lm_eval_float.log <== | | |none | 0|word_perplexity|↓ |7.5551|± | N/A| ==> unsloth_model_lora_baseline_output/lm_eval_quantized.log <== | | |none | 0|word_perplexity|↓ |8.7655|± | N/A| ==> unsloth_model_lora_qat_int4_output/lm_eval_quantized.log <== | | |none | 0|word_perplexity|↓ |8.3548|± | N/A| ```
1680 lines
67 KiB
Python
1680 lines
67 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.
|
|
|
|
__version__ = "2025.9.6"
|
|
|
|
__all__ = [
|
|
"SUPPORTS_BFLOAT16",
|
|
"is_bfloat16_supported",
|
|
"is_vLLM_available",
|
|
|
|
"prepare_model_for_kbit_training",
|
|
"xformers",
|
|
"xformers_attention",
|
|
"xformers_version",
|
|
"__version__",
|
|
"importlib_version",
|
|
"HAS_FLASH_ATTENTION",
|
|
"HAS_FLASH_ATTENTION_SOFTCAPPING",
|
|
"USE_MODELSCOPE",
|
|
"platform_system",
|
|
"patch_tokenizer",
|
|
"get_statistics",
|
|
"Unsloth_Offloaded_Gradient_Checkpointer",
|
|
"offload_to_disk",
|
|
"offload_input_embeddings",
|
|
"offload_output_embeddings",
|
|
"unsloth_offloaded_gradient_checkpoint",
|
|
"torch_compile_options",
|
|
"patch_linear_scaling",
|
|
"patch_llama_rope_scaling",
|
|
"create_boolean_mask",
|
|
"torch_amp_custom_fwd",
|
|
"torch_amp_custom_bwd",
|
|
# "accelerate_old_send_to_device",
|
|
# "accelerate_new_send_to_device",
|
|
"patch_gradient_accumulation_fix",
|
|
"patch_compiling_bitsandbytes",
|
|
"patch_regional_compilation",
|
|
"patch_layernorm",
|
|
"patch_torch_compile",
|
|
"patch_model_and_tokenizer",
|
|
|
|
"patch_unsloth_gradient_checkpointing",
|
|
"unpatch_unsloth_gradient_checkpointing",
|
|
"patch_gradient_checkpointing",
|
|
"unpatch_gradient_checkpointing",
|
|
|
|
"HAS_CUT_CROSS_ENTROPY",
|
|
"EMPTY_LOGITS",
|
|
"fused_linear_cross_entropy",
|
|
"unsloth_fused_ce_loss",
|
|
"patch_unsloth_smart_gradient_checkpointing",
|
|
"unpatch_unsloth_smart_gradient_checkpointing",
|
|
|
|
"patch_compiled_autograd",
|
|
"process_vision_info",
|
|
"unsloth_compile_transformers",
|
|
"patch_fast_lora",
|
|
"validate_loftq_config",
|
|
"RaiseUninitialized",
|
|
"fast_inference_setup",
|
|
"patch_peft_fast_inference",
|
|
"error_out_no_vllm",
|
|
"dequantize_module_weight",
|
|
]
|
|
|
|
import torch
|
|
from typing import Union, Optional, List, Any, Callable, Tuple
|
|
from platform import system as platform_system
|
|
platform_system = platform_system()
|
|
import numpy as np
|
|
import contextlib
|
|
import re
|
|
import functools
|
|
import warnings, subprocess, re, inspect, psutil, os, math
|
|
from unsloth_zoo.utils import Version
|
|
from importlib.metadata import version as importlib_version
|
|
from unsloth import DEVICE_TYPE, DEVICE_COUNT
|
|
from unsloth_zoo.log import logger
|
|
from unsloth_zoo.tokenizer_utils import (
|
|
patch_tokenizer as _patch_tokenizer,
|
|
)
|
|
from unsloth_zoo.patching_utils import (
|
|
patch_compiling_bitsandbytes,
|
|
patch_layernorm,
|
|
patch_torch_compile,
|
|
patch_model_and_tokenizer,
|
|
patch_compiled_autograd,
|
|
)
|
|
from unsloth_zoo.gradient_checkpointing import (
|
|
Unsloth_Offloaded_Gradient_Checkpointer,
|
|
unsloth_offloaded_gradient_checkpoint,
|
|
patch_unsloth_gradient_checkpointing,
|
|
unpatch_unsloth_gradient_checkpointing,
|
|
|
|
Unsloth_Gradient_Checkpointer,
|
|
unsloth_gradient_checkpoint,
|
|
patch_gradient_checkpointing,
|
|
unpatch_gradient_checkpointing,
|
|
|
|
patch_unsloth_smart_gradient_checkpointing,
|
|
unpatch_unsloth_smart_gradient_checkpointing,
|
|
)
|
|
from unsloth_zoo.loss_utils import (
|
|
HAS_CUT_CROSS_ENTROPY,
|
|
fused_linear_cross_entropy,
|
|
_unsloth_get_batch_samples,
|
|
unsloth_fused_ce_loss,
|
|
)
|
|
from unsloth_zoo.vision_utils import (
|
|
process_vision_info,
|
|
)
|
|
from unsloth_zoo.compiler import (
|
|
get_transformers_model_type,
|
|
unsloth_compile_transformers as _unsloth_compile_transformers,
|
|
)
|
|
from unsloth_zoo.training_utils import (
|
|
prepare_model_for_training,
|
|
)
|
|
from unsloth_zoo.temporary_patches import (
|
|
TEMPORARY_PATCHES,
|
|
)
|
|
for temporary_patch in TEMPORARY_PATCHES:
|
|
temporary_patch()
|
|
|
|
# =============================================
|
|
# Disable some warnings which can get annoying
|
|
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "torch")
|
|
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "huggingface_hub")
|
|
warnings.filterwarnings(action = "ignore", category = FutureWarning, module = "huggingface_hub")
|
|
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "trl")
|
|
warnings.filterwarnings(action = "ignore", category = FutureWarning, module = "trl")
|
|
warnings.filterwarnings(action = "ignore", category = FutureWarning, module = "xformers")
|
|
warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = "subprocess")
|
|
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "transformers")
|
|
warnings.filterwarnings(action = "ignore", category = FutureWarning, module = "accelerate")
|
|
warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = "multiprocessing")
|
|
warnings.filterwarnings(action = "ignore", category = RuntimeWarning, module = "multiprocess")
|
|
warnings.filterwarnings(action = "ignore", category = UserWarning, module = "triton")
|
|
|
|
# Stop "Special tokens have been added in the vocabulary, ..."
|
|
import logging
|
|
logging.getLogger("transformers.tokenization_utils_base").setLevel(logging.CRITICAL+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
|
|
|
|
# Stop vLLM messages
|
|
if os.environ.get('UNSLOTH_ENABLE_LOGGING', '0') != '1':
|
|
try:
|
|
from vllm.worker.worker import logger as vllm_worker_logger
|
|
vllm_worker_logger.addFilter(HideLoggingMessage("Sleep mode freed"))
|
|
del vllm_worker_logger
|
|
except:
|
|
pass
|
|
try:
|
|
from vllm.v1.worker.gpu_worker import logger as vllm_gpu_worker_logger
|
|
vllm_gpu_worker_logger.addFilter(HideLoggingMessage("Sleep mode freed"))
|
|
del vllm_gpu_worker_logger
|
|
except:
|
|
pass
|
|
try:
|
|
from vllm.executor.executor_base import logger as vllm_executor_logger
|
|
vllm_executor_logger.addFilter(HideLoggingMessage("to fall asleep"))
|
|
vllm_executor_logger.addFilter(HideLoggingMessage("to wake up"))
|
|
vllm_executor_logger.addFilter(HideLoggingMessage("Executor is not sleeping"))
|
|
del vllm_executor_logger
|
|
except:
|
|
pass
|
|
try:
|
|
from vllm.core.block.prefix_caching_block import logger as vllm_prefix_caching_logger
|
|
vllm_prefix_caching_logger.addFilter(HideLoggingMessage("reset prefix cache"))
|
|
del vllm_prefix_caching_logger
|
|
except:
|
|
pass
|
|
try:
|
|
from vllm.v1.core.block_pool import logger as vllm_block_pool_logger
|
|
vllm_block_pool_logger.addFilter(HideLoggingMessage("reset prefix cache"))
|
|
del vllm_block_pool_logger
|
|
except:
|
|
pass
|
|
try:
|
|
from vllm.lora.models import logger as vllm_lora_model_logger
|
|
vllm_lora_model_logger.addFilter(HideLoggingMessage("Regarding multimodal models, vLLM currently only supports adding"))
|
|
del vllm_lora_model_logger
|
|
except:
|
|
pass
|
|
pass
|
|
|
|
# The speedups for torchdynamo mostly come with GPU Ampere or higher and which is not detected here.
|
|
from transformers.training_args import logger as transformers_training_args_logger
|
|
transformers_training_args_logger.addFilter(HideLoggingMessage("The speedups"))
|
|
# torch.distributed process group is initialized, but parallel_mode != ParallelMode.DISTRIBUTED.
|
|
transformers_training_args_logger.addFilter(HideLoggingMessage("torch.distributed"))
|
|
# average_tokens_across_devices is set to True but it is invalid when world size is1
|
|
transformers_training_args_logger.addFilter(HideLoggingMessage("average_tokens_across_devices"))
|
|
del transformers_training_args_logger
|
|
|
|
# No label_names provided for model class
|
|
from transformers.trainer import logger as transformers_trainer_logger
|
|
transformers_trainer_logger.addFilter(HideLoggingMessage("No label_names"))
|
|
del transformers_trainer_logger
|
|
|
|
# Using the default loss: `ForCausalLMLoss`.
|
|
try:
|
|
from transformers.modeling_utils import logger as transformers_modeling_utils_logger
|
|
transformers_modeling_utils_logger.addFilter(HideLoggingMessage("ForCausalLMLoss"))
|
|
del transformers_modeling_utils_logger
|
|
except:
|
|
pass
|
|
|
|
# The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function.
|
|
try:
|
|
from accelerate.utils.modeling import logger as accelerate_utils_modeling_logger
|
|
accelerate_utils_modeling_logger.addFilter(HideLoggingMessage("The model weights are not tied"))
|
|
del accelerate_utils_modeling_logger
|
|
except:
|
|
pass
|
|
|
|
# Setting `pad_token_id` to `eos_token_id`
|
|
try:
|
|
from transformers.generation.utils import logger as transformers_generation_utils_logger
|
|
transformers_generation_utils_logger.addFilter(HideLoggingMessage("Setting `pad_token_id` to `eos_token_id`"))
|
|
# "You have set `compile_config`
|
|
transformers_generation_utils_logger.addFilter(HideLoggingMessage("compile_config"))
|
|
del transformers_generation_utils_logger
|
|
except:
|
|
pass
|
|
|
|
# The following generation flags are not valid and may be ignored:
|
|
try:
|
|
from transformers.generation.configuration_utils import logger as configuration_logger
|
|
configuration_logger.addFilter(HideLoggingMessage("following generation flags"))
|
|
del configuration_logger
|
|
except:
|
|
pass
|
|
|
|
# Gemma3 It is strongly recommended to train Gemma3 models with the `eager`
|
|
try:
|
|
from transformers.models.gemma3.modeling_gemma3 import logger as gemma3_logger
|
|
gemma3_logger.addFilter(HideLoggingMessage("strongly recommended"))
|
|
del gemma3_logger
|
|
except:
|
|
pass
|
|
|
|
# Xet Storage is enabled for this repo, but the 'hf_xet' package is not installed.
|
|
try:
|
|
from huggingface_hub.file_download import logger as hub_logger
|
|
hub_logger.addFilter(HideLoggingMessage("hf_xet"))
|
|
del hub_logger
|
|
except:
|
|
pass
|
|
|
|
# MXFP4 quantization requires triton >= 3.4.0
|
|
try:
|
|
from transformers.quantizers.quantizer_mxfp4 import logger as mxfp4_logger
|
|
mxfp4_logger.addFilter(HideLoggingMessage("requires triton"))
|
|
del mxfp4_logger
|
|
except:
|
|
pass
|
|
|
|
# You passed `quantization_config` or equivalent parameters
|
|
try:
|
|
warnings.filterwarnings(
|
|
action = "ignore",
|
|
message = r".*quantization_config.*",
|
|
category = UserWarning,
|
|
append = True,
|
|
)
|
|
except:
|
|
pass
|
|
|
|
# UserWarning: Logical operators 'and' and 'or' are deprecated for non-scalar tensors; please use '&' or '|' instead
|
|
# Will be fixed in torch 2.8.1 https://github.com/pytorch/pytorch/issues/158463
|
|
try:
|
|
warnings.filterwarnings(
|
|
action = "ignore",
|
|
message = r".*Logical operators 'and' and 'or'.*",
|
|
category = UserWarning,
|
|
append = True,
|
|
)
|
|
except:
|
|
pass
|
|
|
|
# Using a slow image processor as `use_fast`
|
|
try:
|
|
from transformers.processing_utils import logger as processing_utils_logger
|
|
processing_utils_logger.addFilter(HideLoggingMessage("`use_fast`"))
|
|
del processing_utils_logger
|
|
except:
|
|
pass
|
|
|
|
# Using a slow image processor as `use_fast`
|
|
try:
|
|
from transformers.models.auto.image_processing_auto import logger as processing_utils_logger
|
|
processing_utils_logger.addFilter(HideLoggingMessage("`use_fast`"))
|
|
del processing_utils_logger
|
|
except:
|
|
pass
|
|
|
|
# `use_cache=True` is incompatible with gradient checkpointing
|
|
try:
|
|
from transformers.trainer import logger as trainer_logger
|
|
trainer_logger.addFilter(HideLoggingMessage("`use_cache=True`"))
|
|
del trainer_logger
|
|
except:
|
|
pass
|
|
|
|
# `use_cache=True` is incompatible with gradient checkpointing
|
|
try:
|
|
from transformers.utils.generic import logger as trainer_logger
|
|
trainer_logger.addFilter(HideLoggingMessage("`use_cache=True`"))
|
|
del trainer_logger
|
|
except:
|
|
pass
|
|
|
|
# Errors out on
|
|
# Some weights of Gemma3nForConditionalGeneration were not initialized from the model checkpoint
|
|
from transformers.modeling_utils import logger as transformers_logger
|
|
class _RaiseUninitialized(logging.Handler):
|
|
def __init__(self):
|
|
super().__init__()
|
|
def emit(self, record):
|
|
record_lower = str(record).lower()
|
|
if ("some weights of" in record_lower) and \
|
|
("score.weight" not in record_lower) and \
|
|
("classifier.weight" not in record_lower):
|
|
raise Exception(
|
|
f"Unsloth: Critical error since some weights are not initialized.\n"\
|
|
f"Please try updating Unsloth, transformers and timm via:\n"\
|
|
f"`pip install --upgrade --force-reinstall --no-cache-dir --no-deps unsloth unsloth_zoo transformers timm`\n"\
|
|
f"{str(record)}"
|
|
)
|
|
pass
|
|
class RaiseUninitialized:
|
|
def __init__(self):
|
|
self.error_handler = _RaiseUninitialized()
|
|
transformers_logger.addHandler(self.error_handler)
|
|
def remove(self):
|
|
transformers_logger.removeHandler(self.error_handler)
|
|
pass
|
|
|
|
# Patch get_model_param_count to record correct 4bit / 8bit
|
|
from transformers.trainer_pt_utils import is_deepspeed_zero3_enabled
|
|
|
|
def extract_quant_model_param_count(model):
|
|
"""
|
|
Calculate quant model param count based on difference in param class. Returns int for param count.
|
|
"""
|
|
count: int = 0
|
|
for name, p in model.named_parameters():
|
|
if p.__class__.__name__ == "Params4bit":
|
|
count += 2 * p.numel()
|
|
else:
|
|
count += p.numel()
|
|
return count
|
|
pass
|
|
|
|
def get_model_param_count(model, trainable_only = False):
|
|
"""
|
|
Calculate model's total param count. If trainable_only is True then count only those requiring grads
|
|
"""
|
|
if is_deepspeed_zero3_enabled():
|
|
def numel(p):
|
|
return p.ds_numel if hasattr(p, "ds_numel") else p.numel()
|
|
else:
|
|
def numel(p):
|
|
return p.numel()
|
|
s = sum(numel(p) for p in model.parameters() if not trainable_only or p.requires_grad)
|
|
if (not trainable_only) and \
|
|
hasattr(model, "config") and \
|
|
hasattr(model.config, "quantization_config"):
|
|
approx = extract_quant_model_param_count(model)
|
|
if approx is not None:
|
|
s = approx
|
|
return s
|
|
pass
|
|
import transformers.trainer_pt_utils
|
|
transformers.trainer_pt_utils.get_model_param_count = get_model_param_count
|
|
import transformers.trainer
|
|
transformers.trainer.get_model_param_count = get_model_param_count
|
|
# =============================================
|
|
|
|
# =============================================
|
|
# Edits all Config files to enable RoPE Scaling for all models
|
|
|
|
# Transformers had to update for Mistral Nemo 12b since Attention is (5120, 4096) now.
|
|
def patch_mistral_nemo_config(config):
|
|
if "head_dim (" not in config:
|
|
add_head_dim = "If it is not specified, will default to `8`.\n"\
|
|
" head_dim (`int`, *optional*, defaults to `hidden_size // num_attention_heads`):\n"\
|
|
" The attention head dimension."
|
|
config = config.replace("If it is not specified, will default to `8`.", add_head_dim)
|
|
|
|
add_head_dim = "num_key_value_heads=8,\n head_dim=None,"
|
|
config = config.replace("num_key_value_heads=8,", add_head_dim)
|
|
|
|
add_head_dim = "self.sliding_window = sliding_window\n self.head_dim = head_dim or hidden_size // num_attention_heads\n"
|
|
config = config.replace("self.sliding_window = sliding_window", add_head_dim)
|
|
pass
|
|
return config
|
|
pass
|
|
|
|
try:
|
|
# Some Config files use layer_type_validation
|
|
# for eg Gemma-2, so we must import it to stop errors.
|
|
from transformers.configuration_utils import layer_type_validation
|
|
except:
|
|
pass
|
|
from transformers import __version__ as transformers_version
|
|
from transformers import PretrainedConfig
|
|
model_architectures = ["llama", "mistral", "gemma", "gemma2", "qwen2", "granite", "qwen3", "qwen3_moe", "falcon_h1"]
|
|
|
|
for model_name in model_architectures:
|
|
config_filepath = f"transformers.models.{model_name}.configuration_{model_name}"
|
|
model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
|
|
config_filename = f"{model_name.title().replace('_','')}Config" # qwen3 arch folder is qwen3_moe but config is Qwen3Config. Need to remove underscore(_) for now
|
|
try:
|
|
exec(f"from {config_filepath} import {config_filename}", globals())
|
|
except:
|
|
continue
|
|
|
|
try:
|
|
config = inspect.getsource(eval(config_filename))
|
|
except:
|
|
continue
|
|
if "rope_scaling" in config: continue
|
|
config = re.sub(
|
|
r"(\*\*kwargs)[\s]{0,}\,[\s]{0,}\)[\s]{0,}\:",
|
|
r"rope_scaling=None,"\
|
|
r"\n **kwargs):\n"\
|
|
r"\n self.rope_scaling = rope_scaling\n",
|
|
config,
|
|
)
|
|
|
|
# Just for Mistral Nemo
|
|
if model_name == "mistral":
|
|
if Version(transformers_version) <= Version("4.42.4"):
|
|
config = patch_mistral_nemo_config(config)
|
|
pass
|
|
|
|
exec(config, globals())
|
|
exec(f"import {config_filepath}", globals())
|
|
exec(f"{config_filepath}.{config_filename} = {config_filename}", globals())
|
|
pass
|
|
# =============================================
|
|
|
|
# =============================================
|
|
# torch.cuda.amp.custom_fwd is deprecated >= 2.4
|
|
torch_version = torch.__version__
|
|
if DEVICE_TYPE in ("cuda", "hip"):
|
|
if Version(torch_version) < Version("2.4.0"):
|
|
torch_amp_custom_fwd = torch.cuda.amp.custom_fwd
|
|
torch_amp_custom_bwd = torch.cuda.amp.custom_bwd
|
|
else:
|
|
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "cuda")
|
|
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "cuda")
|
|
pass
|
|
elif DEVICE_TYPE == "xpu":
|
|
if Version(torch_version) < Version("2.6.0"):
|
|
raise RuntimeError("torch.xpu currently only supports torch.version >= 2.6.0")
|
|
else:
|
|
torch_amp_custom_fwd = torch.amp.custom_fwd(device_type = "xpu")
|
|
torch_amp_custom_bwd = torch.amp.custom_bwd(device_type = "xpu")
|
|
# =============================================
|
|
|
|
# =============================================
|
|
# Fix KeyError: 'Cache only has 0 layers, attempted to access layer with index 0'
|
|
# import transformers.cache_utils
|
|
# if hasattr(transformers.cache_utils, "DynamicCache") and \
|
|
# transformers.cache_utils.DynamicCache.__getitem__.__name__ != "__cache_utils_getitem__":
|
|
|
|
# source = inspect.getsource(transformers.cache_utils.DynamicCache.__getitem__)
|
|
# start = source.find("def")
|
|
# spaces = start*" "
|
|
# source = source.split("\n")
|
|
# source = "\n".join(x[start:] for x in source)
|
|
# where = source.find("raise KeyError")
|
|
# source = source[:where] + \
|
|
# f"if len(self) == 0:\n{spaces}{spaces}"\
|
|
# " raise RuntimeError('Unsloth: You must call `FastLanguageModel.for_inference(model)` before doing inference for Unsloth models.')\n" + \
|
|
# f"{spaces}{spaces}else:\n{spaces}{spaces}{spaces}" + source[where:]
|
|
# source = source.replace("__getitem__", "__cache_utils_getitem__", 1)
|
|
# exec(source)
|
|
# transformers.cache_utils.DynamicCache.__getitem__ = __cache_utils_getitem__
|
|
# pass
|
|
# =============================================
|
|
|
|
# =============================================
|
|
# Weird Databricks errors
|
|
from transformers.utils import is_openai_available
|
|
if is_openai_available():
|
|
try:
|
|
from openai import OpenAI
|
|
except:
|
|
print("Unsloth: OpenAI failed to import - ignoring for now.")
|
|
import transformers.utils
|
|
def _is_openai_available(): return False
|
|
transformers.utils.is_openai_available = _is_openai_available
|
|
pass
|
|
pass
|
|
|
|
# =============================================
|
|
# Get Flash Attention v2 if Ampere (RTX 30xx, A100)
|
|
if DEVICE_TYPE in ("cuda", "hip"):
|
|
import bitsandbytes as bnb
|
|
|
|
from transformers import AutoTokenizer
|
|
from transformers.utils.import_utils import _is_package_available
|
|
|
|
SUPPORTS_BFLOAT16 = False
|
|
HAS_FLASH_ATTENTION = False
|
|
HAS_FLASH_ATTENTION_SOFTCAPPING = False
|
|
|
|
if DEVICE_TYPE == "cuda":
|
|
major_version, minor_version = torch.cuda.get_device_capability()
|
|
torch.cuda.get_device_capability = functools.cache(torch.cuda.get_device_capability)
|
|
|
|
if major_version >= 8:
|
|
SUPPORTS_BFLOAT16 = True
|
|
if _is_package_available("flash_attn"):
|
|
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
|
|
try:
|
|
try:
|
|
# See https://github.com/unslothai/unsloth/issues/1437
|
|
from flash_attn.flash_attn_interface import flash_attn_gpu
|
|
except:
|
|
from flash_attn.flash_attn_interface import flash_attn_cuda
|
|
HAS_FLASH_ATTENTION = True
|
|
|
|
# Also check for softcapping
|
|
from flash_attn import __version__ as flash_attn_version
|
|
HAS_FLASH_ATTENTION_SOFTCAPPING = Version(flash_attn_version) >= Version("2.6.3")
|
|
if not HAS_FLASH_ATTENTION_SOFTCAPPING:
|
|
print(
|
|
"Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\
|
|
"Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\
|
|
"To update flash-attn, do the below:\n"\
|
|
'\npip install --no-deps --no-build-isolation --upgrade "flash-attn>=2.6.3"'
|
|
)
|
|
except:
|
|
print(
|
|
"Unsloth: Your Flash Attention 2 installation seems to be broken?\n"\
|
|
"A possible explanation is you have a new CUDA version which isn't\n"\
|
|
"yet compatible with FA2? Please file a ticket to Unsloth or FA2.\n"\
|
|
"We shall now use Xformers instead, which does not have any performance hits!\n"\
|
|
"We found this negligible impact by benchmarking on 1x A100."
|
|
)
|
|
|
|
# Stop Flash Attention from importing!
|
|
import transformers.utils.import_utils
|
|
transformers.utils.import_utils.is_flash_attn_2_available = lambda *args, **kwargs: False
|
|
import transformers.utils
|
|
transformers.utils.is_flash_attn_2_available = lambda *args, **kwargs: False
|
|
|
|
HAS_FLASH_ATTENTION = False
|
|
pass
|
|
else:
|
|
HAS_FLASH_ATTENTION = False
|
|
else:
|
|
# Tri Dao's benchmark shows xformers is faster for now.
|
|
HAS_FLASH_ATTENTION = False
|
|
pass
|
|
elif DEVICE_TYPE == "hip":
|
|
SUPPORTS_BFLOAT16 = True
|
|
if _is_package_available("flash_attn"):
|
|
# Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl"
|
|
try:
|
|
try:
|
|
# See https://github.com/unslothai/unsloth/issues/1437
|
|
from flash_attn.flash_attn_interface import flash_attn_gpu
|
|
except:
|
|
from flash_attn.flash_attn_interface import flash_attn_cuda
|
|
HAS_FLASH_ATTENTION = True
|
|
|
|
# Also check for softcapping
|
|
from flash_attn import __version__ as flash_attn_version
|
|
HAS_FLASH_ATTENTION_SOFTCAPPING = Version(flash_attn_version) >= Version("2.6.3")
|
|
if not HAS_FLASH_ATTENTION_SOFTCAPPING:
|
|
print(
|
|
"Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\
|
|
"Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\
|
|
"To update flash-attn, do the below:\n"\
|
|
'\npip install --no-deps --no-build-isolation --upgrade "flash-attn>=2.6.3"'
|
|
)
|
|
except:
|
|
print(
|
|
"Unsloth: Your Flash Attention 2 installation seems to be broken?\n"\
|
|
"A possible explanation is you have a new CUDA version which isn't\n"\
|
|
"yet compatible with FA2? Please file a ticket to Unsloth or FA2.\n"\
|
|
"We shall now use Xformers instead, which does not have any performance hits!\n"\
|
|
"We found this negligible impact by benchmarking on 1x A100."
|
|
)
|
|
|
|
# Stop Flash Attention from importing!
|
|
import transformers.utils.import_utils
|
|
transformers.utils.import_utils.is_flash_attn_2_available = lambda *args, **kwargs: False
|
|
import transformers.utils
|
|
transformers.utils.is_flash_attn_2_available = lambda *args, **kwargs: False
|
|
|
|
HAS_FLASH_ATTENTION = False
|
|
elif DEVICE_TYPE == "xpu":
|
|
SUPPORTS_BFLOAT16 = True
|
|
|
|
# =============================================
|
|
# Get Xformers
|
|
try:
|
|
from xformers import __version__ as xformers_version
|
|
# [TODO] Xformers does NOT work on RTX 50x (12), B200 (10), Jetson (11)
|
|
# See https://github.com/facebookresearch/xformers/issues/1329
|
|
# CUDA error (/workspace/xfrm2/third_party/flash-attention/hopper/flash_fwd_launch_template.h:188)
|
|
major_version, minor_version = torch.cuda.get_device_capability()
|
|
if (
|
|
(f"{major_version}.{minor_version}" in ("10.0", "11.0", "12.0")) and \
|
|
(Version(xformers_version) in (Version("0.0.32.post2"),))
|
|
):
|
|
raise NotImplementedError(
|
|
"Unsloth: Xformers does not work in RTX 50X, Blackwell GPUs as of yet. Please build from source via\n"\
|
|
"```\n"\
|
|
"pip install ninja\n"\
|
|
"pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers\n"\
|
|
"```\n"
|
|
)
|
|
# Temporarily disable 0.0.27 and higher - inference issues
|
|
if False: #Version(xformers_version) >= Version("0.0.27"):
|
|
raise ImportError(
|
|
"Unsloth: If you are in Colab, we updated the top cell install instructions - please change it to below "\
|
|
"then press Disconnect Runtime and then Restart it.\n"\
|
|
"\n"\
|
|
"%%capture\n"
|
|
"# Installs Unsloth, Xformers (Flash Attention) and all other packages!\n"
|
|
'!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n'
|
|
'!pip install --no-deps "xformers<=0.0.27" trl peft accelerate bitsandbytes\n'\
|
|
'\n'\
|
|
f"Otherwise in local machines, your xformers version of {xformers_version} is too new.\n"\
|
|
'Please downgrade xformers via `pip install --force-reinstall "xformers<=0.0.27"'
|
|
)
|
|
pass
|
|
|
|
if Version(torch_version) < Version("2.2.0") and Version(xformers_version) >= Version("0.0.24"):
|
|
raise ImportError(
|
|
f"Unsloth: You have torch = {torch_version} but xformers = {xformers_version}.\n"\
|
|
f"Please install xformers < 0.0.24 for torch = {torch_version}."
|
|
)
|
|
elif Version(torch_version) < Version("2.3.0") and Version(xformers_version) >= Version("0.0.26"):
|
|
raise ImportError(
|
|
f"Unsloth: You have torch = {torch_version} but xformers = {xformers_version}.\n"\
|
|
f"Please install xformers < 0.0.26 for torch = {torch_version}."
|
|
)
|
|
elif Version(torch_version) < Version("2.4.0") and Version(xformers_version) > Version("0.0.27"):
|
|
raise ImportError(
|
|
f"Unsloth: You have torch = {torch_version} but xformers = {xformers_version}.\n"\
|
|
f"Please install xformers <= 0.0.27 for torch = {torch_version}."
|
|
)
|
|
pass
|
|
|
|
from xformers._cpp_lib import _register_extensions
|
|
try:
|
|
_register_extensions() # Check if C++ modules are loaded correctly
|
|
except Exception as error:
|
|
raise ImportError(
|
|
"Unsloth: Xformers was not installed correctly.\n"\
|
|
"Please install xformers separately first.\n"\
|
|
"Then confirm if it's correctly installed by running:\n"\
|
|
"python -m xformers.info\n\n"
|
|
"Longer error message:\n" + str(error)
|
|
)
|
|
pass
|
|
import xformers.ops.fmha as xformers
|
|
xformers_attention = xformers.memory_efficient_attention
|
|
except ModuleNotFoundError:
|
|
xformers = None
|
|
xformers_attention = None
|
|
xformers_version = None
|
|
except Exception as e:
|
|
print("========\nSwitching to PyTorch attention since your Xformers is broken.\n========\n")
|
|
print(str(e))
|
|
xformers = None
|
|
xformers_attention = None
|
|
xformers_version = None
|
|
pass
|
|
|
|
# Check TRL version
|
|
from trl import __version__ as trl_version
|
|
# Unsloth now supports all TRL versions!
|
|
if False:#Version(trl_version) >= Version("0.9.0"):
|
|
raise ImportError(
|
|
"Unsloth: If you are in Colab, we updated the top cell install instructions - please change it to below "\
|
|
"then press Disconnect Runtime and then Restart it.\n"\
|
|
"\n"\
|
|
"%%capture\n"
|
|
"# Installs Unsloth, Xformers (Flash Attention) and all other packages!\n"
|
|
'!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n'
|
|
'!pip install --no-deps "xformers<=0.0.27" trl peft accelerate bitsandbytes\n'\
|
|
'\n'\
|
|
f"Otherwise in local machines, your TRL version of {trl_version} is too new.\n"\
|
|
'Please downgrade TRL via `pip install --force-reinstall trl'
|
|
)
|
|
pass
|
|
|
|
# =============================================
|
|
# Fix new Xformers versions TypeError: Multiple dispatch failed for 'torch._ops.aten.to.dtype_layout'
|
|
# accelerate_old_send_to_device = None
|
|
# accelerate_new_send_to_device = None
|
|
# if xformers_version is not None and Version(xformers_version) >= Version("0.0.27"):
|
|
# import accelerate.utils.operations
|
|
# if hasattr(accelerate.utils.operations, "send_to_device") and \
|
|
# accelerate.utils.operations.send_to_device.__name__ != "_fixed_send_to_device":
|
|
# accelerate_old_send_to_device = accelerate.utils.operations.send_to_device
|
|
# from accelerate.utils.operations import *
|
|
# send_to_device = inspect.getsource(accelerate.utils.operations.send_to_device)
|
|
# send_to_device = re.sub(
|
|
# r"([ ]{4,})return tensor\.to\(device\)",
|
|
# r"\1try: return tensor.to(device)\n\1except: return tensor",
|
|
# send_to_device,
|
|
# ).replace("def send_to_device", "def _fixed_send_to_device")
|
|
# exec(send_to_device)
|
|
# # accelerate.utils.operations.send_to_device = _fixed_send_to_device
|
|
# accelerate_new_send_to_device = _fixed_send_to_device
|
|
# pass
|
|
# pass
|
|
|
|
# Transformers 4.46 breaks dynamic caching. This is a hack
|
|
import transformers.generation.configuration_utils
|
|
if hasattr(transformers.generation.configuration_utils, "ALL_CACHE_IMPLEMENTATIONS"):
|
|
if type(transformers.generation.configuration_utils.ALL_CACHE_IMPLEMENTATIONS) is list:
|
|
if "dynamic" not in transformers.generation.configuration_utils.ALL_CACHE_IMPLEMENTATIONS:
|
|
transformers.generation.configuration_utils.ALL_CACHE_IMPLEMENTATIONS.append("dynamic")
|
|
pass
|
|
pass
|
|
# =============================================
|
|
|
|
# =============================================
|
|
# Torch compile settings
|
|
UNSLOTH_COMPILE_DEBUG = os.environ.get("UNSLOTH_COMPILE_DEBUG", "0") == "1"
|
|
UNSLOTH_COMPILE_MAXIMUM = os.environ.get("UNSLOTH_COMPILE_MAXIMUM", "0") == "1"
|
|
UNSLOTH_COMPILE_IGNORE_ERRORS = os.environ.get("UNSLOTH_COMPILE_IGNORE_ERRORS", "1") == "1"
|
|
# Just remove max_autotune_gemm warning
|
|
from torch._inductor.runtime.hints import DeviceProperties
|
|
|
|
@functools.lru_cache(None)
|
|
def is_big_gpu(index) -> bool:
|
|
|
|
if DEVICE_TYPE == "xpu":
|
|
prop = DeviceProperties.create(torch.device("xpu", index) if type(index) is int else index)
|
|
min_sms = 16
|
|
else:
|
|
prop = DeviceProperties.create(torch.device("cuda", index) if type(index) is int else index)
|
|
min_sms = 80
|
|
|
|
avail_sms = prop.multi_processor_count
|
|
if avail_sms < min_sms:
|
|
return False
|
|
return True
|
|
|
|
import torch._inductor.utils
|
|
torch._inductor.utils.is_big_gpu = is_big_gpu
|
|
patch_torch_compile(
|
|
debug = UNSLOTH_COMPILE_DEBUG,
|
|
O3 = UNSLOTH_COMPILE_MAXIMUM,
|
|
ignore_errors = UNSLOTH_COMPILE_IGNORE_ERRORS,
|
|
)
|
|
|
|
torch_compile_options = {
|
|
"epilogue_fusion" : True,
|
|
"max_autotune" : True,
|
|
"shape_padding" : True,
|
|
"trace.enabled" : UNSLOTH_COMPILE_DEBUG,
|
|
"triton.cudagraphs" : False,
|
|
}
|
|
|
|
import accelerate
|
|
def torch_compile_kwargs(*args, **kwargs):
|
|
print("Unsloth: Enabled auto compiling")
|
|
return {"dynamic" : True, "fullgraph" : False, "options" : torch_compile_options,}
|
|
pass
|
|
|
|
accelerate.utils.dataclasses.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs
|
|
accelerate.utils.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs
|
|
accelerate.accelerator.TorchDynamoPlugin.to_kwargs = torch_compile_kwargs
|
|
del accelerate
|
|
|
|
def patch_regional_compilation():
|
|
# Regional torch 2.5 Recompilation - weirdly very slow??
|
|
if torch.nn.ModuleList.__name__ == "UnslothModuleList": return
|
|
# Only works for torch 2.5
|
|
if Version(torch.__version__) < Version("2.5.0"): return
|
|
|
|
old_module_list = torch.nn.ModuleList
|
|
os.environ["UNSLOTH_PATCHED"] = "1"
|
|
|
|
def UnslothModuleList(*args, **kwargs):
|
|
if len(args) == 1 and len(kwargs) == 0 and type(args[0]) is list:
|
|
args = [old_module_list([torch.compile(x, dynamic = True, options = torch_compile_options, fullgraph = False) for x in args[0]])]
|
|
return old_module_list(*args, **kwargs)
|
|
pass
|
|
UnslothModuleList.__doc__ = old_module_list.__doc__
|
|
|
|
torch.nn.ModuleList = UnslothModuleList
|
|
return
|
|
pass
|
|
|
|
# =============================================
|
|
|
|
def prepare_model_for_kbit_training(
|
|
model : Any,
|
|
use_gradient_checkpointing : Optional = True,
|
|
use_reentrant : Optional[bool] = True,
|
|
) -> Any:
|
|
return prepare_model_for_training(
|
|
model = model,
|
|
use_gradient_checkpointing = use_gradient_checkpointing,
|
|
use_reentrant = use_reentrant,
|
|
full_finetuning = False,
|
|
train_layernorms = False,
|
|
train_embedding = False,
|
|
train_lm_head = False,
|
|
float32_mixed_precision = True,
|
|
)
|
|
pass
|
|
|
|
# =============================================
|
|
# Weirdly LoraLayer.update_layer downcasts PEFT layers to float16??
|
|
# For mixed precision, we need it to be in float32 not float16.
|
|
from peft import __version__ as peft_version
|
|
from peft.utils.integrations import dequantize_module_weight
|
|
if Version(peft_version) < Version("0.12.0"):
|
|
from peft.tuners.lora.layer import LoraLayer
|
|
try:
|
|
source = inspect.getsource(LoraLayer.update_layer)
|
|
text = "if weight is not None:\n"
|
|
start = source.find(text) + len(text)
|
|
end = source.find("self.to(weight.device)", start)
|
|
spaces = re.findall(r"^([ ]{1,})break", source, flags = re.MULTILINE)[0]
|
|
source = source.replace(source[start : end], spaces)
|
|
spaces = len(re.match(r"[\s]{1,}", source).group(0))
|
|
lines = source.split("\n")
|
|
source = "\n".join(x[spaces:] for x in lines)
|
|
source = re.sub(r"([^\.])nn\.", r"\1torch.nn.", source)
|
|
source = source.replace("def update_layer", "def LoraLayer_update_layer")
|
|
exec(source, globals())
|
|
|
|
# Fix up incorrect downcasting of LoRA weights
|
|
from peft.tuners.lora.layer import LoraLayer
|
|
LoraLayer.update_layer = LoraLayer_update_layer
|
|
from peft.tuners.lora import LoraLayer
|
|
LoraLayer.update_layer = LoraLayer_update_layer
|
|
except:
|
|
logger.warning_once(
|
|
"Unsloth unsuccessfully patched LoraLayer.update_layer. Please file a bug report.\n"\
|
|
"Luckily, your training run will still work in the meantime!"
|
|
)
|
|
pass
|
|
pass
|
|
|
|
# =============================================
|
|
|
|
import psutil
|
|
def _get_statistics(statistics = None, force_download = True):
|
|
# We log some basic stats about which environment is being used.
|
|
# We simply download a README.md file from HF - all data is made public.
|
|
# This is simply so we can check if some envs are broken or not.
|
|
# You can disable this by commenting the below out
|
|
try:
|
|
n_cpus = psutil.cpu_count(logical = False)
|
|
keynames = "\n" + "\n".join(os.environ.keys())
|
|
if statistics is not None: pass
|
|
elif "\nCOLAB_" in keynames and n_cpus == 1: statistics = "colab"
|
|
elif "\nCOLAB_" in keynames: statistics = "colabpro"
|
|
elif "\nKAGGLE_" in keynames: statistics = "kaggle"
|
|
elif "\nRUNPOD_" in keynames: statistics = "runpod"
|
|
elif "\nAWS_" in keynames: statistics = "aws"
|
|
elif "\nAZURE_" in keynames: statistics = "azure"
|
|
# elif "\nK_" in keynames or "\nFUNCTION_" in keynames: statistics = "gcp"
|
|
elif "\nINVOCATION_ID" in keynames: statistics = "lambda"
|
|
# else: statistics = "other"
|
|
else:
|
|
def try_vllm_check():
|
|
vendor_files = (
|
|
"/sys/class/dmi/id/product_version",
|
|
"/sys/class/dmi/id/bios_vendor",
|
|
"/sys/class/dmi/id/product_name",
|
|
"/sys/class/dmi/id/chassis_asset_tag",
|
|
"/sys/class/dmi/id/sys_vendor",
|
|
)
|
|
from pathlib import Path
|
|
for vendor_file in vendor_files:
|
|
path = Path(vendor_file)
|
|
if path.is_file():
|
|
file_content = path.read_text().lower()
|
|
if "amazon" in file_content: return "aws"
|
|
elif "microsoft corporation" in file_content: return "azure"
|
|
elif "google" in file_content: return "gcp"
|
|
return "other"
|
|
pass
|
|
try: statistics = try_vllm_check()
|
|
except: statistics = "other"
|
|
pass
|
|
if statistics is not None:
|
|
from transformers import AutoModelForCausalLM
|
|
stats_model = AutoModelForCausalLM.from_pretrained(
|
|
f"unslothai/{statistics}",
|
|
force_download = force_download,
|
|
)
|
|
del stats_model
|
|
pass
|
|
except:
|
|
pass
|
|
pass
|
|
|
|
|
|
def get_statistics():
|
|
# We log some basic stats about which environment is being used.
|
|
# We simply download a README.md file from HF - all data is made public.
|
|
# This is simply so we can check if some envs are broken or not.
|
|
# You can disable this by setting UNSLOTH_DISABLE_STATISTICS
|
|
import os
|
|
if "UNSLOTH_DISABLE_STATISTICS" in os.environ: return
|
|
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars, are_progress_bars_disabled
|
|
disabled = False
|
|
if not are_progress_bars_disabled():
|
|
disable_progress_bars()
|
|
disabled = True
|
|
pass
|
|
_get_statistics(None)
|
|
_get_statistics("repeat", force_download = False)
|
|
try:
|
|
vram = torch.cuda.get_device_properties(0).total_memory / 1024 / 1024 / 1024
|
|
if vram <= 8 : vram = 8
|
|
elif vram <= 16: vram = 16
|
|
elif vram <= 20: vram = 20
|
|
elif vram <= 24: vram = 24
|
|
elif vram <= 40: vram = 40
|
|
elif vram <= 48: vram = 48
|
|
elif vram <= 80: vram = 80
|
|
else: vram = 96
|
|
_get_statistics(f"vram-{vram}")
|
|
except:
|
|
pass
|
|
pass
|
|
try:
|
|
_get_statistics(f"{DEVICE_COUNT if DEVICE_COUNT <= 8 else 9}")
|
|
except:
|
|
pass
|
|
if disabled: enable_progress_bars()
|
|
pass
|
|
|
|
|
|
# =============================================
|
|
# Fixes Bitsandbytes to remove missing warnings
|
|
from transformers.utils.quantization_config import BitsAndBytesConfig, QuantizationMethod
|
|
BitsAndBytesConfig__init__ = inspect.getsource(BitsAndBytesConfig.__init__)
|
|
BitsAndBytesConfig__init__ = re.sub(
|
|
r"if[\s]{1,}kwargs\:[\s]{1,}.+?\n",
|
|
"",
|
|
BitsAndBytesConfig__init__,
|
|
flags = re.MULTILINE,
|
|
)
|
|
BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.split("\n")
|
|
length_spaces = len(re.match(r"[\s]{1,}", BitsAndBytesConfig__init__[0]).group(0))
|
|
BitsAndBytesConfig__init__ = "\n".join(x[length_spaces:] for x in BitsAndBytesConfig__init__)
|
|
BitsAndBytesConfig__init__ = BitsAndBytesConfig__init__.replace(
|
|
"__init__",
|
|
"_BitsAndBytesConfig__init__",
|
|
)
|
|
exec(BitsAndBytesConfig__init__, globals())
|
|
|
|
if DEVICE_COUNT == 1:
|
|
from accelerate.utils.dataclasses import DistributedType
|
|
def _prepare_backend(self, *args, **kwargs): return None, DistributedType.NO
|
|
import accelerate.state
|
|
accelerate.state.PartialState._prepare_backend = _prepare_backend
|
|
accelerate.accelerator.Accelerator.distributed_type = lambda *args, **kwargs: DistributedType.NO
|
|
pass
|
|
|
|
# to move multiple tensors to the same device
|
|
def move_to_device(target_device, *tensors):
|
|
"""
|
|
Move multiple tensors to target device if they're not already there.
|
|
|
|
Args:
|
|
target_device: The target device to move tensors to
|
|
*tensors: Variable number of tensors to potentially move
|
|
|
|
Returns:
|
|
tuple: The tensors on the target device (same objects if already on device, new if moved)
|
|
"""
|
|
if isinstance(target_device, int):
|
|
target_device = torch.device(target_device)
|
|
elif isinstance(target_device, str):
|
|
# if string we expect it to be a device name like "cuda:0"
|
|
target_device = torch.device(target_device)
|
|
elif isinstance(target_device, torch.device):
|
|
pass
|
|
else:
|
|
raise ValueError(f"Invalid target device: {target_device}")
|
|
pass
|
|
moved_tensors = []
|
|
for tensor in tensors:
|
|
if tensor.device != target_device:
|
|
moved_tensors.append(tensor.to(target_device))
|
|
else:
|
|
moved_tensors.append(tensor)
|
|
return tuple(moved_tensors) if len(moved_tensors) > 1 else moved_tensors[0]
|
|
|
|
import transformers.utils.quantization_config
|
|
transformers.utils.quantization_config.BitsAndBytesConfig.__init__ = _BitsAndBytesConfig__init__
|
|
# =============================================
|
|
|
|
# Offloading to disk for modules (lm_head, embed_tokens)
|
|
import pickle
|
|
|
|
def offload_to_disk(W, model, name, temporary_location : str = "_unsloth_temporary_saved_buffers"):
|
|
file_location = os.path.join(temporary_location, model.config._name_or_path)
|
|
if not os.path.exists(file_location):
|
|
os.makedirs(file_location)
|
|
pass
|
|
|
|
filename = os.path.join(file_location, f"{name}.pt")
|
|
W = W.weight if hasattr(W, "weight") else W
|
|
torch.save(W, filename, pickle_module = pickle, pickle_protocol = pickle.HIGHEST_PROTOCOL,)
|
|
# We must use weights_only = False due to pickling
|
|
offloaded_W = torch.load(filename, map_location = "cpu", mmap = True, weights_only = False)
|
|
offloaded_W._offloaded_file_location = filename
|
|
return offloaded_W
|
|
pass
|
|
|
|
|
|
def offload_input_embeddings(model, temporary_location : str = "_unsloth_temporary_saved_buffers"):
|
|
offloaded_W = offload_to_disk(model.get_input_embeddings(), model, "input_embeddings", temporary_location)
|
|
new_input_embeddings = torch.nn.Embedding.from_pretrained(offloaded_W)
|
|
new_input_embeddings._offloaded_file_location = offloaded_W._offloaded_file_location
|
|
model.set_input_embeddings(new_input_embeddings)
|
|
return
|
|
pass
|
|
|
|
|
|
def offload_output_embeddings(model, temporary_location : str = "_unsloth_temporary_saved_buffers"):
|
|
offloaded_W = offload_to_disk(model.get_output_embeddings(), model, "output_embeddings", temporary_location)
|
|
|
|
new_output_embeddings = torch.nn.Linear(1, 1, bias = None)
|
|
del new_output_embeddings.weight
|
|
new_output_embeddings.weight = offloaded_W
|
|
new_output_embeddings.in_features = offloaded_W.shape[1]
|
|
new_output_embeddings.out_features = offloaded_W.shape[0]
|
|
|
|
new_output_embeddings._offloaded_file_location = offloaded_W._offloaded_file_location
|
|
model.set_output_embeddings(new_output_embeddings)
|
|
return
|
|
pass
|
|
|
|
|
|
# Fixes a weird Torch 2.3 bug which says T4s have bfloat16
|
|
def is_bfloat16_supported():
|
|
return SUPPORTS_BFLOAT16
|
|
pass
|
|
|
|
def is_vLLM_available():
|
|
return _is_package_available("vllm")
|
|
pass
|
|
|
|
# Patches models to add RoPE Scaling
|
|
def patch_linear_scaling(
|
|
model_name = "gemma2",
|
|
rope_module = None,
|
|
scaled_rope_module = None,
|
|
attention_module = None,
|
|
):
|
|
assert(rope_module is not None and scaled_rope_module is not None)
|
|
assert(attention_module is not None)
|
|
|
|
rope_name = rope_module.__name__
|
|
scaled_rope_name = scaled_rope_module.__name__
|
|
model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
|
|
exec_code = \
|
|
f"import torch.nn as nn\n"\
|
|
f"from typing import Union, Optional, List, Any, Callable, Tuple\n"\
|
|
f"from {model_filepath} import logger, "\
|
|
f"{model_name.title()}Attention, {model_name.title()}Config"
|
|
|
|
try:
|
|
function = inspect.getsource(attention_module.__init__)
|
|
except:
|
|
# Most likely already patched!
|
|
return None, None
|
|
where = function.find("def")
|
|
function = function.split("\n")
|
|
function = "\n".join(x[where:] for x in function)
|
|
init_name = f"{model_name.title()}Attention__init__"
|
|
function = function.replace("def __init__", f"def {init_name}")
|
|
function = function.replace(
|
|
"super().__init__()",
|
|
f"super({model_name.title()}Attention, self).__init__()",
|
|
)
|
|
fix_rope_function = """
|
|
if getattr(self.config, "rope_scaling", None) is None:
|
|
self.rotary_emb = {rope_function}(
|
|
dim = self.head_dim,
|
|
max_position_embeddings=self.max_position_embeddings,
|
|
base=self.rope_theta,
|
|
)
|
|
else:
|
|
scaling_type = self.config.rope_scaling["type"]
|
|
scaling_factor = self.config.rope_scaling["factor"]
|
|
if scaling_type == "linear":
|
|
self.rotary_emb = {scaled_rope_function}(
|
|
dim = self.head_dim,
|
|
max_position_embeddings=self.max_position_embeddings,
|
|
scaling_factor=scaling_factor,
|
|
base=self.rope_theta,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown RoPE scaling type {{scaling_type}}")
|
|
pass
|
|
"""
|
|
fix_rope_function = fix_rope_function.format(
|
|
rope_function = rope_module.__name__,
|
|
scaled_rope_function = scaled_rope_module.__name__,
|
|
)
|
|
rotary_emb = re.findall(
|
|
r"self\.rotary\_emb \= .+?\)", function,
|
|
flags = re.DOTALL | re.MULTILINE,
|
|
)
|
|
if len(rotary_emb) == 0:
|
|
return None, exec_code + "\n\n" + function
|
|
|
|
rotary_emb = rotary_emb[0]
|
|
function = function.replace(rotary_emb, fix_rope_function, 1)
|
|
function = exec_code + "\n\n" + function
|
|
return init_name, function
|
|
pass
|
|
|
|
|
|
# Patches for Llama-3 LlamaExtendedRotaryEmbedding
|
|
def patch_llama_rope_scaling(
|
|
model_name = "llama",
|
|
rope_module = None,
|
|
scaled_rope_module = None,
|
|
extended_rope_module = None,
|
|
attention_module = None,
|
|
longrope_module = None,
|
|
):
|
|
assert(\
|
|
rope_module is not None and \
|
|
scaled_rope_module is not None and \
|
|
extended_rope_module is not None
|
|
)
|
|
assert(attention_module is not None)
|
|
|
|
rope_name = rope_module.__name__
|
|
scaled_rope_name = scaled_rope_module.__name__
|
|
model_filepath = f"transformers.models.{model_name}.modeling_{model_name}"
|
|
exec_code = \
|
|
f"import torch.nn as nn\n"\
|
|
f"from typing import Union, Optional, List, Any, Callable, Tuple\n"\
|
|
f"from {model_filepath} import logger, "\
|
|
f"{model_name.title()}Attention, {model_name.title()}Config"
|
|
|
|
try:
|
|
function = inspect.getsource(attention_module.__init__)
|
|
except:
|
|
# Most likely already patched!
|
|
return None, None
|
|
where = function.find("def")
|
|
function = function.split("\n")
|
|
function = "\n".join(x[where:] for x in function)
|
|
init_name = f"{model_name.title()}Attention__init__"
|
|
function = function.replace("def __init__", f"def {init_name}")
|
|
function = function.replace(
|
|
"super().__init__()",
|
|
f"super({model_name.title()}Attention, self).__init__()",
|
|
)
|
|
fix_rope_function = """
|
|
if getattr(self.config, "rope_scaling", None) is None:
|
|
self.rotary_emb = {rope_function}(
|
|
dim = self.head_dim,
|
|
max_position_embeddings=self.max_position_embeddings,
|
|
base=self.rope_theta,
|
|
)
|
|
else:
|
|
scaling_type1 = self.config.rope_scaling.get("type", None)
|
|
scaling_type2 = self.config.rope_scaling.get("rope_type", None)
|
|
scaling_type = scaling_type1 if scaling_type1 is not None else scaling_type2
|
|
scaling_factor = self.config.rope_scaling.get("factor")
|
|
|
|
if scaling_type == "linear":
|
|
self.rotary_emb = {scaled_rope_function}(
|
|
dim = self.head_dim,
|
|
max_position_embeddings=self.max_position_embeddings,
|
|
scaling_factor=scaling_factor,
|
|
base=self.rope_theta,
|
|
)
|
|
elif scaling_type == "llama3":
|
|
self.rotary_emb = {extended_rope_function}(
|
|
dim = self.head_dim,
|
|
max_position_embeddings=self.max_position_embeddings,
|
|
base=self.rope_theta,
|
|
)
|
|
elif scaling_type == "longrope":
|
|
self.rotary_emb = {longrope_rope_function}(
|
|
dim = self.head_dim,
|
|
max_position_embeddings = self.max_position_embeddings,
|
|
original_max_position_embeddings = self.config.original_max_position_embeddings,
|
|
base = self.rope_theta,
|
|
short_factor = self.config.rope_scaling['short_factor'],
|
|
long_factor = self.config.rope_scaling['long_factor' ],
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown RoPE scaling type {{scaling_type}}")
|
|
pass
|
|
"""
|
|
|
|
fix_rope_function = fix_rope_function.format(
|
|
rope_function = rope_module.__name__,
|
|
scaled_rope_function = scaled_rope_module.__name__,
|
|
extended_rope_function = extended_rope_module.__name__,
|
|
longrope_rope_function = \
|
|
(longrope_module if longrope_module is not None else rope_module).__name__
|
|
)
|
|
rotary_emb = re.findall(
|
|
r"self\.rotary\_emb \= .+?\)", function,
|
|
flags = re.DOTALL | re.MULTILINE,
|
|
)
|
|
if len(rotary_emb) == 0: return None, function
|
|
rotary_emb = rotary_emb[0]
|
|
function = function.replace(rotary_emb, fix_rope_function, 1)
|
|
function = exec_code + "\n\n" + function
|
|
return init_name, function
|
|
pass
|
|
|
|
|
|
def create_boolean_mask(n = 4096, sliding_window = 2048):
|
|
# Creates a boolean mask for attention
|
|
mask = torch.ones(n, n, dtype = torch.bool)
|
|
if sliding_window == 0:
|
|
return torch.triu(mask, diagonal = 1, out = mask)
|
|
pass
|
|
torch.triu(mask, diagonal = 0, out = mask)
|
|
torch.triu(mask.T, diagonal = -sliding_window, out = mask.T)
|
|
mask = mask.T
|
|
torch.logical_not(mask, out = mask)
|
|
return mask
|
|
pass
|
|
|
|
|
|
def test_mask_creation():
|
|
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
|
for n in range(2, 23):
|
|
for s in range(1, 23):
|
|
correct_mask = AttentionMaskConverter(
|
|
is_causal = True,
|
|
sliding_window = s,
|
|
).to_causal_4d(1, n, n, dtype = torch.float16,).squeeze(0).squeeze(0)
|
|
correct_mask = (correct_mask == correct_mask.min())
|
|
our_mask = create_boolean_mask(n = n, sliding_window = s)
|
|
assert(torch.all(correct_mask == our_mask))
|
|
pass
|
|
correct_mask = AttentionMaskConverter(
|
|
is_causal = True,
|
|
sliding_window = None,
|
|
).to_causal_4d(1, n, n, dtype = torch.float16,).squeeze(0).squeeze(0)
|
|
correct_mask = (correct_mask == correct_mask.min())
|
|
our_mask = create_boolean_mask(n = n, sliding_window = 0)
|
|
assert(torch.all(correct_mask == our_mask))
|
|
pass
|
|
pass
|
|
|
|
|
|
def _unsloth_pre_compute_loss(self, model, inputs, *args, **kwargs):
|
|
num_items_in_batch = None
|
|
|
|
if "num_items_in_batch" in kwargs:
|
|
num_items_in_batch = kwargs["num_items_in_batch"]
|
|
if num_items_in_batch is None:
|
|
# Remove it since the model does not support it!
|
|
kwargs.pop("num_items_in_batch")
|
|
elif "num_items_in_batch" not in inputs:
|
|
inputs["num_items_in_batch"] = num_items_in_batch
|
|
pass
|
|
pass
|
|
|
|
# Get gradient accumulation steps if possible
|
|
if num_items_in_batch is None and \
|
|
getattr(getattr(self, "args", self), "gradient_accumulation_steps", 1) != 1:
|
|
|
|
inner_model = model
|
|
if hasattr(inner_model, "base_model"): inner_model = inner_model. base_model
|
|
if hasattr(inner_model, "model"): inner_model = inner_model.model
|
|
name = inner_model.__class__.__name__
|
|
|
|
logger.warning_once(
|
|
f"Unsloth: Not an error, but {name} does not accept `num_items_in_batch`.\n"\
|
|
"Using gradient accumulation will be very slightly less accurate.\n"\
|
|
"Read more on gradient accumulation issues here: https://unsloth.ai/blog/gradient"
|
|
)
|
|
pass
|
|
outputs = self._old_compute_loss(model, inputs, *args, **kwargs)
|
|
return outputs
|
|
pass
|
|
|
|
|
|
def patch_gradient_accumulation_fix(Trainer):
|
|
# Fixes gradient accumulation
|
|
import inspect
|
|
if hasattr(Trainer, "get_batch_samples"):
|
|
if Trainer.get_batch_samples.__name__ == "_unsloth_get_batch_samples": return
|
|
if \
|
|
not inspect.getsource(Trainer.get_batch_samples).strip()\
|
|
.endswith("return batch_samples, num_items_in_batch"):
|
|
|
|
raise NotImplementedError("Unsloth: Please make a Github issue immediately!!")
|
|
else:
|
|
if Trainer.get_batch_samples.__name__ != "_unsloth_get_batch_samples":
|
|
Trainer.get_batch_samples = _unsloth_get_batch_samples
|
|
pass
|
|
|
|
# Also fix passing in num_items_in_batch
|
|
if not hasattr(Trainer, "_old_compute_loss"):
|
|
Trainer._old_compute_loss = Trainer.compute_loss
|
|
Trainer.compute_loss = _unsloth_pre_compute_loss
|
|
pass
|
|
pass
|
|
else:
|
|
logger.warning_once(
|
|
"Unsloth: We fixed a gradient accumulation bug, "\
|
|
"but it seems like you don't have the latest transformers version!\n"\
|
|
"Please update transformers, TRL and unsloth via:\n"\
|
|
'`pip install --upgrade --no-cache-dir --no-deps unsloth transformers git+https://github.com/huggingface/trl.git`'
|
|
)
|
|
pass
|
|
|
|
# Also fix up loss scaling ie negate loss *= self.args.gradient_accumulation_steps
|
|
if Trainer.training_step.__name__ == "_unsloth_training_step": return
|
|
if "num_items_in_batch" not in inspect.signature(Trainer.training_step).parameters: return
|
|
|
|
function = inspect.getsource(Trainer.training_step)
|
|
where = function.find("def")
|
|
function = function.split("\n")
|
|
function = "\n".join(x[where:] for x in function)
|
|
|
|
# Import all variables that need importing
|
|
import transformers.trainer
|
|
items_in_trainer = dir(transformers.trainer)
|
|
good_items = []
|
|
for item in items_in_trainer:
|
|
if item in function: good_items.append(item)
|
|
pass
|
|
exec("from transformers.trainer import (" + ", ".join(x for x in good_items) + ")", globals())
|
|
|
|
# Accelerate does / self.args.gradient_accumulation_steps internally, so if we already
|
|
# summed it up and did the division before hand, we have to negate it.
|
|
function = function.replace(
|
|
"loss *= self.args.gradient_accumulation_steps",
|
|
"if num_items_in_batch is not None: loss *= self.args.gradient_accumulation_steps",
|
|
)
|
|
function = function.replace("def training_step", "def _unsloth_training_step", 1)
|
|
|
|
# Fix 4.47.0 issue where num_items_in_batch was removed
|
|
# See https://github.com/huggingface/transformers/pull/35121
|
|
function = function.replace(
|
|
"if self.model_accepts_loss_kwargs:",
|
|
"if False:",
|
|
)
|
|
|
|
# Fix when num_items_in_batch is nothing
|
|
# https://github.com/huggingface/transformers/pull/35207
|
|
function = re.sub(
|
|
r"else:\n"\
|
|
r"([\s]{4,})self\.accelerator\.backward\(loss, \*\*kwargs\)\n"\
|
|
r"(.+?)if num_items_in_batch is None\:\n"\
|
|
r"(.+?)return loss\.detach\(\) \/ self\.args\.gradient_accumulation_steps",
|
|
|
|
"else:\n"\
|
|
"\2if num_items_in_batch is None:\n"\
|
|
"\3loss = loss / self.args.gradient_accumulation_steps\n"\
|
|
"\1self.accelerator.backward(loss, **kwargs)",
|
|
|
|
function,
|
|
)
|
|
|
|
exec(function, globals())
|
|
Trainer.training_step = _unsloth_training_step
|
|
pass
|
|
|
|
|
|
def patch_tokenizer(model, tokenizer):
|
|
model, tokenizer = _patch_tokenizer(model, tokenizer)
|
|
if model is not None:
|
|
model.config.update({"unsloth_version" : __version__})
|
|
return model, tokenizer
|
|
pass
|
|
|
|
|
|
def patch_fast_lora():
|
|
import peft.tuners.lora.bnb
|
|
peft.tuners.lora.bnb.Linear4bit.forward = fast_lora_forward
|
|
pass
|
|
|
|
|
|
def unsloth_compile_transformers(
|
|
dtype,
|
|
model_name,
|
|
model_types,
|
|
token = None,
|
|
revision = None,
|
|
trust_remote_code = False,
|
|
sdpa_dynamic_mask = True,
|
|
sdpa_bool_masks = True,
|
|
sdpa_gqa_replace = True,
|
|
sdpa_dynamic_compile = True,
|
|
compile_attention = True,
|
|
disable_causal_masks = True,
|
|
compile_torch_modules = True,
|
|
compile_custom_modules = True,
|
|
compile_function_calls = True,
|
|
fuse_lm_head = True,
|
|
gradient_checkpointing = True,
|
|
manual_replacements = True,
|
|
fast_lora_forwards = True,
|
|
fast_residual_stream = True,
|
|
accurate_accumulation = True,
|
|
epilogue_fusion = True,
|
|
max_autotune = False,
|
|
shape_padding = True,
|
|
cudagraphs = False,
|
|
debug = False,
|
|
fullgraph = True,
|
|
import_from_cache = False,
|
|
disable = False,
|
|
return_logits = False,
|
|
unsloth_force_compile = False,
|
|
):
|
|
if Version(torch_version) < Version("2.4.0"):
|
|
print(
|
|
"="*30 + \
|
|
"Unsloth: Unfortunately Unsloth vision and other newer optimized models need Torch 2.4 or later.\n"\
|
|
f"You have Torch version {torch_version}. Please upgrade your Torch version by visiting https://pytorch.org/\n"\
|
|
"For now your models will not get optimized, but will still work for now!"
|
|
)
|
|
return
|
|
pass
|
|
if trust_remote_code and unsloth_force_compile == False:
|
|
print(
|
|
"Unsloth: We can't trace models if `trust_remote_code = True`, "\
|
|
"so turning off some optimizations!"
|
|
)
|
|
return model_types, False
|
|
model_types = list(dict().fromkeys(model_types).keys())
|
|
if disable: return model_types, False
|
|
|
|
supports_sdpa = [True]
|
|
for model_type in model_types:
|
|
_unsloth_compile_transformers(
|
|
model_type,
|
|
sdpa_dynamic_mask = sdpa_dynamic_mask,
|
|
sdpa_bool_masks = sdpa_bool_masks,
|
|
sdpa_gqa_replace = sdpa_gqa_replace,
|
|
sdpa_dynamic_compile = sdpa_dynamic_compile,
|
|
compile_attention = compile_attention,
|
|
disable_causal_masks = disable_causal_masks,
|
|
compile_torch_modules = compile_torch_modules,
|
|
compile_custom_modules = compile_custom_modules,
|
|
compile_function_calls = compile_function_calls,
|
|
fuse_lm_head = fuse_lm_head,
|
|
gradient_checkpointing = gradient_checkpointing,
|
|
manual_replacements = manual_replacements,
|
|
fast_lora_forwards = fast_lora_forwards,
|
|
fast_residual_stream = fast_residual_stream,
|
|
accurate_accumulation = accurate_accumulation,
|
|
epilogue_fusion = epilogue_fusion,
|
|
max_autotune = max_autotune,
|
|
shape_padding = shape_padding,
|
|
cudagraphs = cudagraphs,
|
|
debug = debug,
|
|
fullgraph = fullgraph,
|
|
import_from_cache = import_from_cache,
|
|
disable = disable,
|
|
return_logits = return_logits,
|
|
supports_sdpa = supports_sdpa,
|
|
)
|
|
pass
|
|
# Redo patches which override compiler
|
|
for temporary_patch in TEMPORARY_PATCHES:
|
|
temporary_patch()
|
|
return model_types, supports_sdpa[0]
|
|
pass
|
|
|
|
# We need an empty logits flag to warn people logits will not be returned anymore unless asked ie
|
|
# os.environ['UNSLOTH_RETURN_LOGITS'] = '1'
|
|
LOGITS_ERROR_STRING = \
|
|
"Unsloth: Logits are empty from 2024.11 onwards. To get raw logits again, please "\
|
|
'set the environment variable `UNSLOTH_RETURN_LOGITS` to `"1" BEFORE starting to train ie before `trainer.train()`. For example:\n'\
|
|
"```\nimport os\n"\
|
|
"os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n"\
|
|
"trainer.train()\n```\n"\
|
|
"No need to restart your console - just add `os.environ['UNSLOTH_RETURN_LOGITS'] = '1'` before trainer.train() and re-run the cell!"
|
|
|
|
def raise_logits_error(*args, **kwargs): raise NotImplementedError(LOGITS_ERROR_STRING)
|
|
def return_none(*args, **kwargs): return None
|
|
class EmptyLogits:
|
|
def __init__(self): return
|
|
def raise_getattr_error(self, attr): return return_none if attr == "to" else raise_logits_error
|
|
__getitem__ = raise_logits_error
|
|
__getattr__ = raise_getattr_error
|
|
def __repr__(self): return LOGITS_ERROR_STRING
|
|
def __str__ (self): return LOGITS_ERROR_STRING
|
|
pass
|
|
EMPTY_LOGITS = EmptyLogits()
|
|
functions = dir(torch.Tensor)
|
|
for j, function in enumerate(functions):
|
|
if function.startswith("__") and function.endswith("__"):
|
|
exec(f"def raise_{j}(*args, **kwargs): print('{function}')", globals(), locals())
|
|
try: exec(f"EMPTY_LOGITS.{function} = raise_{j}", globals(), locals())
|
|
except: continue
|
|
pass
|
|
|
|
import importlib
|
|
USE_MODELSCOPE = os.environ.get("UNSLOTH_USE_MODELSCOPE", "0") == "1"
|
|
if USE_MODELSCOPE:
|
|
if importlib.util.find_spec("modelscope") is None:
|
|
raise ImportError(f'You are using the modelscope hub, please install modelscope by `pip install modelscope -U`')
|
|
pass
|
|
pass
|
|
|
|
|
|
def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, model):
|
|
from peft import LoraConfig
|
|
|
|
if loftq_config is None: loftq_config = {}
|
|
|
|
signature = str(inspect.signature(LoraConfig))
|
|
SUPPORTS_LOFTQ = "loftq_config" in signature
|
|
|
|
if lora_dropout != 0:
|
|
logger.warning_once(
|
|
f"Unsloth: Dropout = 0 is supported for fast patching. You are using dropout = {lora_dropout}.\n"\
|
|
f"Unsloth will patch all other layers, except LoRA matrices, causing a performance hit."
|
|
)
|
|
pass
|
|
|
|
if bias != "none":
|
|
logger.warning_once(
|
|
f"Unsloth: bias = `none` is supported for fast patching. You are using bias = {bias}.\n"\
|
|
f"Unsloth will patch all other layers, except LoRA matrices, causing a performance hit."
|
|
)
|
|
pass
|
|
|
|
if not (type(init_lora_weights) is bool or \
|
|
init_lora_weights == "gaussian" or init_lora_weights == "loftq"):
|
|
raise ValueError(
|
|
'Unsloth: `init_lora_weights` must be either [True, False, "gaussian", "loftq"].'
|
|
)
|
|
pass
|
|
|
|
if init_lora_weights == "loftq":
|
|
|
|
if not SUPPORTS_LOFTQ:
|
|
import peft
|
|
raise RuntimeError(
|
|
f"Unsloth: Your PEFT version of {peft.__version__} does not support LoftQ init.\n"\
|
|
"Please install PEFT 0.7.2 or higher.\n"\
|
|
"You can also install from source: `pip install git+https://github.com/huggingface/peft.git"
|
|
)
|
|
pass
|
|
|
|
if loftq_config == {}:
|
|
from peft import LoftQConfig
|
|
logger.warning_once(
|
|
"Unsloth: init_lora_weights = `loftq` is set, but `loftq_config` is None.\n"\
|
|
"We shall use `loftq_config = LoftQConfig(loftq_bits = 4, loftq_iter = 1)`."
|
|
)
|
|
loftq_config = LoftQConfig(loftq_bits = 4, loftq_iter = 1)
|
|
pass
|
|
|
|
if hasattr(model.config, "quantization_config"):
|
|
raise ValueError(
|
|
"Unsloth: You are using `loftq` init, yet `load_in_4bit = True` was set.\n"\
|
|
"Reload your model without any quantization by setting `load_in_4bit = False`."
|
|
)
|
|
pass
|
|
pass
|
|
|
|
return loftq_config
|
|
|
|
def fast_inference_setup(model_name, model_config):
|
|
fast_inference = True
|
|
if not is_vLLM_available():
|
|
logger.warning_once("Unsloth: vLLM is not installed! Will use Unsloth inference!")
|
|
fast_inference = False
|
|
pass
|
|
from unsloth_zoo.vllm_utils import (
|
|
patch_vllm,
|
|
vllm_dynamic_quant_supported,
|
|
)
|
|
patch_vllm()
|
|
if model_name.endswith("unsloth-bnb-4bit"):
|
|
if not vllm_dynamic_quant_supported(model_name, model_config):
|
|
# Instead use -bnb-4bit variant
|
|
logger.warning_once(
|
|
f"Unsloth: Switching from Unsloth dynamic quant to normal quant since\n"\
|
|
f"we do not yet support fast inference for {model_name}"
|
|
)
|
|
model_name = model_name[:-len("unsloth-bnb-4bit")] + "bnb-4bit"
|
|
pass
|
|
pass
|
|
return fast_inference, model_name
|
|
|
|
def patch_peft_fast_inference(model):
|
|
vllm_engine = getattr(model.model, "vllm_engine", None)
|
|
if vllm_engine is not None:
|
|
model.vllm_engine = model.model.vllm_engine
|
|
model.fast_generate = model.model.fast_generate
|
|
model.fast_generate_batches = model.model.fast_generate_batches
|
|
|
|
# Also saving and loading LoRA
|
|
from unsloth_zoo.vllm_utils import save_lora, load_lora
|
|
model.save_lora = functools.partial(save_lora, model)
|
|
model.load_lora = functools.partial(load_lora, model)
|
|
pass
|
|
|
|
def error_out_no_vllm(*args, **kwargs):
|
|
raise NotImplementedError("Unsloth: vLLM is not yet supported for fast inference for this model! Please use `.generate` instead")
|
|
|
|
|
|
def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.Module:
|
|
"""
|
|
Transform a model for Quantization-Aware Training (QAT) during fine-tuning.
|
|
|
|
On a high level, this means fake quantizing the base (frozen) model during training.
|
|
Fake quantization refers to simulating quantization numerics in high precision (e.g. bf16).
|
|
This helps mitigate quantization degradations when the model is quantized after training.
|
|
|
|
QAT can be optionally combined with LoRA fine-tuning to for additional throughput improvement.
|
|
For more details: https://dev-discuss.pytorch.org/t/speeding-up-qat-by-1-89x-with-lora/2700
|
|
"""
|
|
from torchao.quantization import (
|
|
Float8DynamicActivationInt4WeightConfig,
|
|
Float8DynamicActivationFloat8WeightConfig,
|
|
Int8DynamicActivationInt4WeightConfig,
|
|
Int4WeightOnlyConfig,
|
|
PerRow,
|
|
quantize_,
|
|
)
|
|
from torchao.quantization.qat import QATConfig
|
|
filter_fn = None
|
|
if qat_scheme == "fp8-int4":
|
|
group_size = 128
|
|
base_config = Float8DynamicActivationInt4WeightConfig()
|
|
filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
|
|
elif qat_scheme == "fp8-fp8":
|
|
base_config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
|
|
elif qat_scheme == "int8-int4":
|
|
group_size = 32
|
|
base_config = Int8DynamicActivationInt4WeightConfig(group_size=group_size)
|
|
filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
|
|
elif qat_scheme == "int4":
|
|
group_size = 128
|
|
base_config = Int4WeightOnlyConfig(group_size=group_size)
|
|
filter_fn = lambda m, _: isinstance(m, torch.nn.Linear) and m.in_features >= group_size
|
|
else:
|
|
raise ValueError(f"Unexpected QAT scheme {qat_scheme}")
|
|
pass
|
|
quantize_(model, QATConfig(base_config, step="prepare"), filter_fn=filter_fn)
|
|
return model
|
|
pass
|