diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 912e6f7e3f..fbb14013ff 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -24,7 +24,7 @@ from .utils import ( is_cdna, ) from transformers.models.llama.modeling_llama import logger -from packaging.version import Version +from unsloth_zoo.utils import Version from unsloth_zoo.loss_utils import ( patch_loss_functions as _patch_loss_functions, diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py index e9f56763d6..c33317ee02 100644 --- a/unsloth/models/cohere.py +++ b/unsloth/models/cohere.py @@ -15,7 +15,7 @@ from .llama import * from ._utils import __version__ from unsloth_zoo.hf_utils import dtype_from_config -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( AttentionConfig, @@ -35,8 +35,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.42"): raise ImportError( diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 291d442673..1789a9cd92 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -14,7 +14,7 @@ from .llama import * from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import ( build_sdpa_packed_attention_mask, @@ -34,8 +34,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.38"): raise ImportError( diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 4b2503b8a1..16d04955d3 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -14,7 +14,7 @@ from .llama import * from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( @@ -41,8 +41,6 @@ try: repeat_kv, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.42"): raise ImportError( diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py index f85f1b641f..aae746aed1 100644 --- a/unsloth/models/granite.py +++ b/unsloth/models/granite.py @@ -15,7 +15,7 @@ from .llama import * import os from ._utils import __version__ -from unsloth_zoo.utils import _get_dtype +from unsloth_zoo.utils import _get_dtype, Version from unsloth_zoo.hf_utils import dtype_from_config from ..utils.packing import get_packed_info_from_kwargs from ..utils.attention_dispatch import ( @@ -41,8 +41,6 @@ try: GraniteForCausalLM, ) except: - from packaging.version import Version - transformers_version = Version(transformers_version) if not transformers_version >= Version("4.45.0"): raise ImportError( diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index 85332e1116..fe2a89d893 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -28,7 +28,6 @@ from .mapper import ( ) # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading! -from packaging.version import Version from transformers import __version__ as transformers_version from unsloth.models._utils import TorchAOConfig from unsloth_zoo.utils import Version diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 88aeeda8a1..35a15d03a6 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -43,10 +43,28 @@ torch_compile_options = { "triton.cudagraphs": False, } -from trl import __version__ as trl_version +# vLLM compatibility shim (TRL expects GuidedDecodingParams even if vLLM doesn't provide it) +try: + import vllm.sampling_params as _unsloth_vllm_sp + if not hasattr(_unsloth_vllm_sp, "GuidedDecodingParams"): + class GuidedDecodingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + _unsloth_vllm_sp.GuidedDecodingParams = GuidedDecodingParams +except Exception: + pass + +from trl import __version__ as trl_version_raw +from importlib.metadata import version as importlib_version from unsloth_zoo.utils import Version -trl_version = Version(trl_version) +try: + trl_version = Version(trl_version_raw) +except Exception: + try: + trl_version = Version(importlib_version("trl")) + except Exception: + trl_version = Version("0.0.0") def vLLMSamplingParams(**kwargs): @@ -220,7 +238,7 @@ RLTrainer_replacement = ''' import os from typing import * from dataclasses import dataclass, field -from packaging.version import Version +from unsloth_zoo.utils import Version import torch import numpy as np from contextlib import nullcontext @@ -242,12 +260,18 @@ def prepare_for_training_mode(f): @functools.wraps(f) def wrapper(self, *args, **kwargs): # Enable training mode + _was_training = None + if hasattr(self, 'model') and hasattr(self.model, "training"): + _was_training = self.model.training if hasattr(self, 'model') and hasattr(self.model, "for_training"): self.model.for_training() output = f(self, *args, **kwargs) - # Return inference mode + # Restore previous mode when possible if hasattr(self, 'model') and hasattr(self.model, "for_inference"): - self.model.for_inference() + if _was_training is False: + self.model.for_inference() + elif _was_training is True and hasattr(self.model, "for_training"): + self.model.for_training() # Reset gradient checkpointing buffers to free memory while staying ready for next run try: reset_unsloth_gradient_checkpointing_buffers() @@ -331,6 +355,27 @@ class Unsloth{RLTrainer_name}(_Unsloth{RLTrainer_name}): pass ''' +def _wrap_grpo_generate_and_score(trainer_cls): + if not hasattr(trainer_cls, "_generate_and_score_completions"): + return + original = trainer_cls._generate_and_score_completions + if getattr(original, "_unsloth_restore_training_wrapped", False): + return + + def wrapped(self, *args, **kwargs): + was_training = getattr(getattr(self, "model", None), "training", None) + try: + return original(self, *args, **kwargs) + finally: + if was_training is False and hasattr(self, "model") and hasattr(self.model, "for_inference"): + try: + self.model.for_inference() + except Exception: + pass + + wrapped._unsloth_restore_training_wrapped = True + trainer_cls._generate_and_score_completions = wrapped + def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): # Patch for vLLM and Unsloth PEFT @@ -1059,6 +1104,16 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): globals(), ) + if trainer_file == "grpo_trainer": + try: + _wrap_grpo_generate_and_score( + getattr(created_module, f"Unsloth{RLTrainer_name}") + ) + except Exception as e: + logger.info( + f"Unsloth: Could not wrap _generate_and_score_completions for {RLTrainer_name}: {e}" + ) + def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, imports): init = inspect.getsource(RLTrainer.__init__) diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 0d98cff305..858dcf2cd3 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -211,7 +211,7 @@ def _backwards_compatible_trainer(trainer_class, config_class): if "processing_class" in trainer_params and "tokenizer" in kwargs: kwargs["processing_class"] = kwargs.pop("tokenizer") - if ("args" in kwargs) and (Version(trl.__version__) >= Version("0.13.0.dev0")): + 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 @@ -412,7 +412,7 @@ def _patch_trl_trainer(): if hasattr(trl, "__UNSLOTH_BACKWARDS_COMPATIBLE__"): return - if Version(trl.__version__) <= Version("0.11.0"): + if Version(trl) <= Version("0.11.0"): return import trl.trainer