623 lines
22 KiB
Python
623 lines
22 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 logging
|
|
import os
|
|
import psutil
|
|
import warnings
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional, List
|
|
from functools import wraps
|
|
|
|
import trl
|
|
import inspect
|
|
from trl import SFTTrainer
|
|
|
|
# why: bypass partially-initialised unsloth ns during _gpu_init load
|
|
from .models._utils import is_bfloat16_supported
|
|
from unsloth.utils import (
|
|
configure_padding_free,
|
|
configure_sample_packing,
|
|
enable_padding_free_metadata,
|
|
enable_sample_packing,
|
|
)
|
|
from unsloth_zoo.training_utils import (
|
|
unsloth_train as _unsloth_train,
|
|
)
|
|
from unsloth_zoo.vision_utils import (
|
|
UnslothVisionDataCollator,
|
|
)
|
|
from unsloth_zoo.hf_utils import get_transformers_model_type
|
|
from unsloth_zoo.utils import Version
|
|
import dataclasses
|
|
|
|
__all__ = [
|
|
"UnslothTrainingArguments",
|
|
"UnslothTrainer",
|
|
"unsloth_train",
|
|
"_patch_trl_trainer",
|
|
"UnslothVisionDataCollator",
|
|
"QGaloreConfig",
|
|
]
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_AUTO_PADDING_FREE_ENV_DISABLED = os.environ.get(
|
|
"UNSLOTH_DISABLE_AUTO_PADDING_FREE", ""
|
|
).strip().lower() in {"1", "true", "yes", "on"}
|
|
|
|
PADDING_FREE_BLOCKLIST = {
|
|
"gemma2", # - gemma2: Uses slow_attention_softcapping which has torch.compile issues
|
|
"gpt_oss", # - gpt_oss: Uses Flex Attention which doesn't handle padding_free correctly
|
|
}
|
|
|
|
|
|
def _should_pack(config) -> bool:
|
|
if config is None or not getattr(config, "packing", False):
|
|
return False
|
|
return not getattr(config, "_unsloth_disable_auto_packing", False)
|
|
|
|
|
|
def _should_auto_padding_free(config) -> bool:
|
|
if (
|
|
config is None
|
|
or _AUTO_PADDING_FREE_ENV_DISABLED
|
|
or getattr(config, "packing", False)
|
|
):
|
|
return False
|
|
return getattr(config, "padding_free", None) is None
|
|
|
|
|
|
def _disable_sample_packing(config):
|
|
if config is None:
|
|
return
|
|
for attr, value in (("packing", False), ("padding_free", False)):
|
|
if hasattr(config, attr):
|
|
setattr(config, attr, value)
|
|
if hasattr(config, "remove_unused_columns"):
|
|
setattr(config, "remove_unused_columns", True)
|
|
setattr(config, "_unsloth_disable_auto_packing", True)
|
|
|
|
|
|
_AUTO_PACK_SKIP_MESSAGES = (
|
|
"packing is not supported",
|
|
"padding-free training",
|
|
"passing a custom data collator",
|
|
)
|
|
|
|
|
|
def _should_skip_auto_packing_error(exc: Exception) -> bool:
|
|
message = str(exc).lower()
|
|
return any(msg in message for msg in _AUTO_PACK_SKIP_MESSAGES)
|
|
|
|
|
|
# Unsloth gradient accumulation fix:
|
|
from transformers import __version__ as transformers_version, ProcessorMixin
|
|
|
|
if Version(transformers_version) > Version("4.45.2"):
|
|
|
|
def unsloth_train(trainer, *args, **kwargs):
|
|
return trainer.train(*args, **kwargs)
|
|
|
|
else:
|
|
|
|
def unsloth_train(trainer, *args, **kwargs):
|
|
if len(args) != 0 or len(kwargs) != 0:
|
|
raise RuntimeError(
|
|
"Unsloth: Our custom gradient accumulation fixed trainer does not support other arguments.\n"
|
|
"If you want to use our fix inside of HF, please update `transformers` to the latest version via:\n"
|
|
"`pip uninstall transformers -y && pip install --upgrade --no-cache-dir transformers`"
|
|
)
|
|
print(
|
|
"Unsloth: Using our custom gradient accumulation fixed trainer, which is not feature complete.\n"
|
|
"If you want to use our fix inside of HF, please update `transformers` to the latest version via:\n"
|
|
"`pip uninstall transformers -y && pip install --upgrade --no-cache-dir transformers`"
|
|
)
|
|
return _unsloth_train(trainer)
|
|
|
|
|
|
try:
|
|
from trl import SFTConfig as TrainingArguments
|
|
except:
|
|
from transformers import TrainingArguments
|
|
|
|
|
|
@dataclass
|
|
class QGaloreConfig:
|
|
"""Configuration for Q-GaLore optimizer integration.
|
|
|
|
Pass an instance of this class to ``UnslothTrainingArguments`` (via
|
|
``q_galore_config``) to enable Q-GaLore training.
|
|
"""
|
|
|
|
rank: int = 256
|
|
update_proj_gap: int = 200
|
|
scale: float = 0.25
|
|
proj_quant: bool = True
|
|
proj_quant_group_size: int = -1
|
|
proj_quant_n_bit: int = 4
|
|
weight_quant: bool = False
|
|
stochastic_round: bool = True
|
|
weight_group_size: int = 128
|
|
cos_threshold: float = 0.4
|
|
gamma_proj: float = 2.0
|
|
queue_size: int = 5
|
|
target_modules: Optional[List[str]] = None
|
|
|
|
|
|
class UnslothTrainingArguments(TrainingArguments):
|
|
def __init__(
|
|
self,
|
|
embedding_learning_rate: float = None,
|
|
q_galore_config: Optional[QGaloreConfig] = None,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
self.q_galore_config = q_galore_config
|
|
self.embedding_learning_rate = embedding_learning_rate
|
|
super().__init__(*args, **kwargs)
|
|
self.embedding_learning_rate = embedding_learning_rate
|
|
|
|
|
|
def _create_unsloth_optimizer(
|
|
model,
|
|
optimizer_cls,
|
|
optimizer_kwargs,
|
|
embedding_lr = 5e-5,
|
|
):
|
|
lr = optimizer_kwargs["lr"]
|
|
weight_decay = optimizer_kwargs.get("weight_decay", 0.0)
|
|
|
|
param_groups = {
|
|
"non_embeddings": {},
|
|
"embeddings": {},
|
|
}
|
|
|
|
for name, param in model.named_parameters():
|
|
if not param.requires_grad:
|
|
continue
|
|
if name.endswith("modules_to_save.default.weight"):
|
|
partial_name = name[: -len(".modules_to_save.default.weight")]
|
|
partial_name = partial_name[partial_name.rfind(".") + 1 :]
|
|
print(
|
|
f"Unsloth: Setting lr = {embedding_lr:.2e} instead of {lr:.2e} for {partial_name}."
|
|
)
|
|
param_groups["embeddings"][name] = param
|
|
else:
|
|
param_groups["non_embeddings"][name] = param
|
|
|
|
optimizer_grouped_parameters = [
|
|
{
|
|
"params": list(param_groups["non_embeddings"].values()),
|
|
"weight_decay": weight_decay,
|
|
"lr": lr,
|
|
},
|
|
{
|
|
"params": list(param_groups["embeddings"].values()),
|
|
"weight_decay": weight_decay,
|
|
"lr": embedding_lr,
|
|
},
|
|
]
|
|
optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
|
|
return optimizer
|
|
|
|
|
|
class UnslothTrainer(SFTTrainer):
|
|
def create_optimizer(self):
|
|
# --- Q-GaLore optimizer ---
|
|
q_galore_config = getattr(self.args, "q_galore_config", None)
|
|
if q_galore_config is not None and self.optimizer is None:
|
|
embedding_lr = getattr(self.args, "embedding_learning_rate", None)
|
|
return self._create_q_galore_optimizer(q_galore_config, embedding_lr)
|
|
|
|
# --- Embedding-LR optimizer ---
|
|
embedding_learning_rate = getattr(self.args, "embedding_learning_rate", None)
|
|
if embedding_learning_rate is None:
|
|
return super().create_optimizer()
|
|
|
|
if self.optimizer is None:
|
|
optimizer_cls, optimizer_kwargs = SFTTrainer.get_optimizer_cls_and_kwargs(
|
|
self.args
|
|
)
|
|
self.optimizer = _create_unsloth_optimizer(
|
|
self.model,
|
|
optimizer_cls,
|
|
optimizer_kwargs,
|
|
embedding_learning_rate,
|
|
)
|
|
return self.optimizer
|
|
|
|
def _create_q_galore_optimizer(self, config: "QGaloreConfig", embedding_lr = None):
|
|
"""Build the Q-GaLore optimizer from a QGaloreConfig."""
|
|
from unsloth.optimizers.q_galore_adamw import (
|
|
QGaLoreAdamW8bit,
|
|
make_q_galore_param_groups,
|
|
install_weight_quant_hooks,
|
|
)
|
|
|
|
lr = self.args.learning_rate
|
|
weight_decay = self.args.weight_decay
|
|
|
|
param_groups = make_q_galore_param_groups(
|
|
self.model,
|
|
lr = lr,
|
|
weight_decay = weight_decay,
|
|
rank = config.rank,
|
|
update_proj_gap = config.update_proj_gap,
|
|
scale = config.scale,
|
|
proj_quant = config.proj_quant,
|
|
proj_quant_group_size = config.proj_quant_group_size,
|
|
proj_quant_n_bit = config.proj_quant_n_bit,
|
|
weight_quant = config.weight_quant,
|
|
stochastic_round = config.stochastic_round,
|
|
weight_group_size = config.weight_group_size,
|
|
cos_threshold = config.cos_threshold,
|
|
gamma_proj = config.gamma_proj,
|
|
queue_size = config.queue_size,
|
|
target_modules = config.target_modules,
|
|
)
|
|
|
|
# --- Split embedding params with custom LR (Fix #2) ---
|
|
if embedding_lr is not None:
|
|
# Build a fast param->name lookup (O(N) instead of O(N*M))
|
|
param_to_name = {id(p): name for name, p in self.model.named_parameters()}
|
|
|
|
new_groups = []
|
|
for group in param_groups:
|
|
if "rank" in group:
|
|
# GaLore group — keep as-is (embeddings are never in here)
|
|
new_groups.append(group)
|
|
continue
|
|
# Non-GaLore group: split out embedding params
|
|
embed_params = []
|
|
other_params = []
|
|
for p in group["params"]:
|
|
# Check if this param belongs to a modules_to_save embedding
|
|
name = param_to_name.get(id(p))
|
|
if name and name.endswith("modules_to_save.default.weight"):
|
|
partial_name = name[: -len(".modules_to_save.default.weight")]
|
|
partial_name = partial_name[partial_name.rfind(".") + 1 :]
|
|
print(
|
|
f"Unsloth: Setting lr = {embedding_lr:.2e} instead of {lr:.2e} for {partial_name}."
|
|
)
|
|
embed_params.append(p)
|
|
else:
|
|
other_params.append(p)
|
|
if other_params:
|
|
other_group = dict(group)
|
|
other_group["params"] = other_params
|
|
new_groups.append(other_group)
|
|
if embed_params:
|
|
embed_group = dict(group)
|
|
embed_group["params"] = embed_params
|
|
embed_group["lr"] = embedding_lr
|
|
new_groups.append(embed_group)
|
|
param_groups = new_groups
|
|
|
|
# --- Forward optimizer hyperparameters (Fix #3) ---
|
|
self.optimizer = QGaLoreAdamW8bit(
|
|
param_groups,
|
|
lr = lr,
|
|
weight_decay = weight_decay,
|
|
betas = (self.args.adam_beta1, self.args.adam_beta2),
|
|
eps = self.args.adam_epsilon,
|
|
)
|
|
|
|
# Initialize INT8 weight quantization if enabled
|
|
if config.weight_quant:
|
|
QGaLoreAdamW8bit.init_weight_quantization(
|
|
self.model,
|
|
param_groups,
|
|
group_size = config.weight_group_size,
|
|
stochastic = config.stochastic_round,
|
|
)
|
|
# Forward pre-hooks dequantize INT8 weights to float before each
|
|
# forward pass, allowing the optimizer to free float weight memory
|
|
# between steps.
|
|
install_weight_quant_hooks(self.model)
|
|
|
|
n_galore = sum(len(g["params"]) for g in param_groups if "rank" in g)
|
|
n_other = sum(len(g["params"]) for g in param_groups if "rank" not in g)
|
|
print(
|
|
f"🦥 Unsloth: Q-GaLore enabled — "
|
|
f"{n_galore} GaLore params (rank={config.rank}), "
|
|
f"{n_other} standard params."
|
|
)
|
|
|
|
return self.optimizer
|
|
|
|
|
|
# From `trl>=0.13.0`, they changed how to pass several params to the trainer
|
|
# We need to patch to make the transition smooth
|
|
def _resolve_trainer_params(trainer_class, init_fn):
|
|
"""Resolve the real named parameters for a trainer __init__.
|
|
|
|
Some TRL trainers (e.g., ORPOTrainer in TRL 0.27.1) are thin wrappers
|
|
with only ``def __init__(self, *args, **kwargs)``. For those, walk the
|
|
MRO and return the first parent class that has real named parameters.
|
|
"""
|
|
params = inspect.signature(init_fn).parameters
|
|
named = {
|
|
k
|
|
for k, v in params.items()
|
|
if v.kind
|
|
in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
|
|
and k != "self"
|
|
}
|
|
if named:
|
|
return set(params.keys())
|
|
|
|
# Thin wrapper detected - walk MRO for real signature
|
|
for cls in trainer_class.__mro__[1:]:
|
|
if cls is object:
|
|
continue
|
|
parent_init = cls.__dict__.get("__init__")
|
|
if parent_init is None:
|
|
continue
|
|
try:
|
|
parent_params = inspect.signature(parent_init).parameters
|
|
parent_named = {
|
|
k
|
|
for k, v in parent_params.items()
|
|
if v.kind
|
|
in (
|
|
inspect.Parameter.POSITIONAL_OR_KEYWORD,
|
|
inspect.Parameter.KEYWORD_ONLY,
|
|
)
|
|
and k != "self"
|
|
}
|
|
if parent_named:
|
|
return set(parent_params.keys())
|
|
except (ValueError, TypeError):
|
|
continue
|
|
return set(params.keys())
|
|
|
|
|
|
def _backwards_compatible_trainer(trainer_class, config_class):
|
|
original_init = trainer_class.__init__
|
|
|
|
@wraps(original_init)
|
|
def new_init(self, *args, **kwargs):
|
|
# All Trainer tokenizer are now called processing_class
|
|
trainer_params = _resolve_trainer_params(trainer_class, original_init)
|
|
|
|
if "processing_class" in trainer_params and "tokenizer" in kwargs:
|
|
kwargs["processing_class"] = kwargs.pop("tokenizer")
|
|
|
|
if ("args" in kwargs) and (Version(trl) >= Version("0.13.0.dev0")):
|
|
training_args = kwargs.pop("args", None)
|
|
|
|
# Get parameters that Trainer.__init__ actually expects
|
|
trainer_params.remove("self")
|
|
trainer_params.remove("args")
|
|
|
|
# Get fields that should be passed to Config init
|
|
config_fields = {
|
|
field.name: field
|
|
for field in dataclasses.fields(config_class)
|
|
if field.init
|
|
}
|
|
|
|
# Create config dict with valid fields from training_args
|
|
config_dict = {
|
|
name: getattr(training_args, name)
|
|
for name in config_fields
|
|
if hasattr(training_args, name)
|
|
}
|
|
|
|
# Get parameters that exist in Config but not in TrainingArguments
|
|
from transformers import TrainingArguments
|
|
|
|
moved_params = set(inspect.signature(config_class).parameters.keys()) - set(
|
|
inspect.signature(TrainingArguments).parameters.keys()
|
|
)
|
|
|
|
# Separate kwargs into trainer kwargs and config kwargs
|
|
trainer_kwargs = {}
|
|
additional_config_kwargs = {}
|
|
|
|
for key, value in kwargs.items():
|
|
if key in trainer_params:
|
|
trainer_kwargs[key] = value
|
|
elif key in moved_params or key in config_fields:
|
|
additional_config_kwargs[key] = value
|
|
else:
|
|
additional_config_kwargs[key] = value
|
|
|
|
# Update config_dict with additional kwargs
|
|
config_dict.update(additional_config_kwargs)
|
|
|
|
# Create Config with all the collected parameters
|
|
# Reinitialising config class with parameters (that were none initially but populated on first init)
|
|
# causes the 2nd init to fail as there are mutual exclusive checks on pairs of parameters.
|
|
# Refer: https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_config.py#L499-L502 for example
|
|
# So we only create config class if the previous init was not TrainingArguments
|
|
if not isinstance(training_args, TrainingArguments):
|
|
config = config_class(**config_dict)
|
|
else:
|
|
config = training_args
|
|
|
|
# Reconstruct kwargs for Trainer
|
|
kwargs = trainer_kwargs
|
|
kwargs["args"] = config
|
|
original_init(self, *args, **kwargs)
|
|
|
|
return new_init
|
|
|
|
|
|
def _patch_sft_trainer_auto_packing(trl_module):
|
|
sft_trainer = getattr(trl_module, "SFTTrainer", None)
|
|
if sft_trainer is None:
|
|
return
|
|
if getattr(sft_trainer, "_unsloth_auto_packing_wrapped", False):
|
|
return
|
|
|
|
original_init = sft_trainer.__init__
|
|
|
|
@wraps(original_init)
|
|
def new_init(self, *args, **kwargs):
|
|
config_arg = None
|
|
if len(args) >= 2:
|
|
config_arg = args[1]
|
|
else:
|
|
config_arg = kwargs.get("args")
|
|
|
|
# Check if model type is unsupported for padding_free
|
|
model = kwargs.get("model")
|
|
is_unsupported_model = False
|
|
is_vlm = False
|
|
if model is not None:
|
|
model_config = getattr(model, "config", None)
|
|
if model_config is not None:
|
|
model_types = get_transformers_model_type(model_config)
|
|
# Blocklist: models that don't work correctly with padding_free
|
|
is_unsupported_model = any(
|
|
x in PADDING_FREE_BLOCKLIST for x in model_types
|
|
)
|
|
|
|
# Check if VLM
|
|
architectures = getattr(model_config, "architectures", None)
|
|
if architectures is None:
|
|
architectures = []
|
|
is_vlm = any(
|
|
x.endswith("ForConditionalGeneration") for x in architectures
|
|
)
|
|
is_vlm = is_vlm or hasattr(model_config, "vision_config")
|
|
|
|
processing_class = kwargs.get("processing_class") or kwargs.get("tokenizer")
|
|
data_collator = kwargs.get("data_collator")
|
|
|
|
# We also disable vision language models for padding free collators
|
|
blocked = (
|
|
(data_collator is not None)
|
|
or isinstance(processing_class, ProcessorMixin)
|
|
or is_vlm
|
|
or is_unsupported_model
|
|
or (
|
|
os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
|
|
) # Disable padding free on forced logits
|
|
)
|
|
requested_pack = bool(getattr(config_arg, "packing", False))
|
|
if blocked:
|
|
if hasattr(config_arg, "packing"):
|
|
setattr(config_arg, "packing", False)
|
|
if hasattr(config_arg, "padding_free"):
|
|
setattr(config_arg, "padding_free", False)
|
|
|
|
if blocked and requested_pack:
|
|
reason = "custom data collator"
|
|
if data_collator is None and isinstance(processing_class, ProcessorMixin):
|
|
reason = "processor-based model"
|
|
elif is_vlm:
|
|
reason = "vision-language model"
|
|
elif is_unsupported_model:
|
|
reason = f"unsupported model type(s): {', '.join(model_types)}"
|
|
message = "Unsloth: Sample packing skipped " f"({reason} detected)."
|
|
print(message)
|
|
|
|
packing_active = False
|
|
if _should_pack(config_arg) and not blocked:
|
|
configure_sample_packing(config_arg)
|
|
packing_active = True
|
|
logger.info("Unsloth: Sample packing enabled for SFTTrainer instance.")
|
|
|
|
# Resolve padding_free: None (default) = auto-enable unless env-disabled or packing
|
|
auto_padding_free_active = False
|
|
padding_free_requested = getattr(config_arg, "padding_free", None) is True
|
|
if not blocked:
|
|
if padding_free_requested:
|
|
configure_padding_free(config_arg)
|
|
elif _should_auto_padding_free(config_arg):
|
|
configure_padding_free(config_arg)
|
|
auto_padding_free_active = True
|
|
logger.info(
|
|
"Unsloth: Padding-free batching auto-enabled for SFTTrainer instance."
|
|
)
|
|
|
|
try:
|
|
original_init(self, *args, **kwargs)
|
|
except ValueError as exc:
|
|
if packing_active and _should_skip_auto_packing_error(exc):
|
|
logger.info(
|
|
"Unsloth: Auto sample packing failed because trainer reported an incompatible setup (%s).",
|
|
exc,
|
|
)
|
|
_disable_sample_packing(config_arg)
|
|
packing_active = False
|
|
original_init(self, *args, **kwargs)
|
|
else:
|
|
raise
|
|
|
|
trainer_args = getattr(self, "args", None)
|
|
trainer_packing = bool(trainer_args and getattr(trainer_args, "packing", False))
|
|
trainer_padding_free = bool(
|
|
trainer_args and getattr(trainer_args, "padding_free", False)
|
|
)
|
|
|
|
if blocked and trainer_args is not None:
|
|
# Mirror the block on the trainer args to avoid re-enabling later
|
|
setattr(trainer_args, "packing", False)
|
|
setattr(trainer_args, "padding_free", False)
|
|
|
|
if (
|
|
not blocked
|
|
and trainer_packing
|
|
and (packing_active or _should_pack(trainer_args))
|
|
):
|
|
enable_sample_packing(self.model, self)
|
|
print(
|
|
"🦥 Unsloth: Packing enabled - training is >2x faster and uses less VRAM!"
|
|
)
|
|
elif not blocked and trainer_padding_free:
|
|
enable_padding_free_metadata(self.model, self)
|
|
message = (
|
|
"🦥 Unsloth: Padding-free auto-enabled, enabling faster training."
|
|
if auto_padding_free_active
|
|
else "🦥 Unsloth: Padding-free enabled, enabling faster training."
|
|
)
|
|
print(message)
|
|
|
|
sft_trainer.__init__ = new_init
|
|
sft_trainer._unsloth_auto_packing_wrapped = True
|
|
|
|
|
|
def _patch_trl_trainer():
|
|
import trl
|
|
|
|
if hasattr(trl, "__UNSLOTH_BACKWARDS_COMPATIBLE__"):
|
|
return
|
|
if Version(trl) <= Version("0.11.0"):
|
|
return
|
|
|
|
import trl.trainer
|
|
|
|
trl_classes = dir(trl.trainer)
|
|
trl_trainers = set(
|
|
x[: -len("Trainer")] for x in trl_classes if x.endswith("Trainer")
|
|
)
|
|
trl_configs = set(x[: -len("Config")] for x in trl_classes if x.endswith("Config"))
|
|
trl_classes = list(trl_trainers & trl_configs)
|
|
|
|
for x in trl_classes:
|
|
try:
|
|
exec(
|
|
f"trl.{x}Trainer.__init__ = _backwards_compatible_trainer(trl.{x}Trainer, trl.{x}Config)",
|
|
globals(),
|
|
)
|
|
except:
|
|
continue
|
|
|
|
_patch_sft_trainer_auto_packing(trl)
|
|
|
|
trl.__UNSLOTH_BACKWARDS_COMPATIBLE__ = True
|