# 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 as _UnslothVisionDataCollatorBase, ) from unsloth.models.vision import check_dataset_for_missing_videos 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", "check_dataset_for_missing_videos", ] logger = logging.getLogger(__name__) class UnslothVisionDataCollator(_UnslothVisionDataCollatorBase): """ Drop-in zoo collator that validates local video paths on every batch (deduped across batches), applying formatting_func first so formatter-made paths are checked too. Raises FileNotFoundError on missing files instead of silently training on empty video tensors (issue #5085). """ __slots__ = ("_checked_video_paths",) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._checked_video_paths = set() def __call__(self, examples): formatting_func = self.formatting_func if formatting_func is not None: examples = [formatting_func(example) for example in examples] check_dataset_for_missing_videos( examples, raise_error = True, checked = self._checked_video_paths, ) if formatting_func is None: return super().__call__(examples) # why: base __call__ would reapply formatting_func; applied above. self.formatting_func = None try: return super().__call__(examples) finally: self.formatting_func = formatting_func _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: # 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 (no embeddings here) new_groups.append(group) continue # Non-GaLore group: split out embedding params embed_params = [] other_params = [] for p in group["params"]: 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, ) if config.weight_quant: QGaLoreAdamW8bit.init_weight_quantization( self.model, param_groups, group_size = config.weight_group_size, stochastic = config.stochastic_round, ) # Pre-hooks dequantize INT8 weights to float before each forward, # letting the optimizer 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 are thin ``*args, **kwargs`` wrappers; for those, walk the MRO and return the first parent with 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): # tokenizer is now 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) trainer_params.remove("self") trainer_params.remove("args") # Fields that should be passed to Config init config_fields = { field.name: field for field in dataclasses.fields(config_class) if field.init } config_dict = { name: getattr(training_args, name) for name in config_fields if hasattr(training_args, name) } # Params 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 config_dict.update(additional_config_kwargs) # Only build the config if the previous init wasn't TrainingArguments: # reinitialising it would re-trigger mutually-exclusive param checks. # See https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_config.py#L499-L502 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") 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) is_unsupported_model = any(x in PADDING_FREE_BLOCKLIST for x in model_types) 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") # Disable padding-free for VLMs / custom collators / blocklisted models 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 = f"Unsloth: Sample packing skipped ({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) # get_peft_model installs a pre-train forward detector for plain LoRA/vision models, # but only RL trainers run the reset via prepare_for_training_mode. Wire it into the # SFT train() path too, else a grad-enabled probe before train() leaves the poisoned # Dynamo cache in place and the detector hook installed on every training forward. # (For UnslothSFTTrainer the later prepare_for_training_mode assignment supersedes this.) if not getattr(self, "_unsloth_train_reset_wrapped", False): try: from unsloth.models._utils import _unsloth_reset_stray_compile_cache _orig_train = self.train @wraps(_orig_train) def _train_with_reset(*train_args, **train_kwargs): try: _unsloth_reset_stray_compile_cache(self) except Exception: pass return _orig_train(*train_args, **train_kwargs) self.train = _train_with_reset self._unsloth_train_reset_wrapped = True except Exception: pass 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