From 8e941a64226b050b3737e942ae340473c0ab6bed Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 5 Jan 2026 07:02:36 +0000 Subject: [PATCH 1/5] Improve TRL compatibility and GRPO state restore --- unsloth/kernels/cross_entropy_loss.py | 2 +- unsloth/models/cohere.py | 4 +- unsloth/models/gemma.py | 4 +- unsloth/models/gemma2.py | 4 +- unsloth/models/granite.py | 4 +- unsloth/models/loader_utils.py | 1 - unsloth/models/rl.py | 65 ++++++++++++++++++++++++--- unsloth/trainer.py | 4 +- 8 files changed, 67 insertions(+), 21 deletions(-) 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 From defbf038b22edc9ce35693b2f632f5298887188c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 5 Jan 2026 07:03:34 +0000 Subject: [PATCH 2/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/rl.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 35a15d03a6..23cbbf0256 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -46,10 +46,13 @@ torch_compile_options = { # 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 @@ -355,6 +358,7 @@ 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 @@ -367,7 +371,11 @@ def _wrap_grpo_generate_and_score(trainer_cls): try: return original(self, *args, **kwargs) finally: - if was_training is False and hasattr(self, "model") and hasattr(self.model, "for_inference"): + if ( + was_training is False + and hasattr(self, "model") + and hasattr(self.model, "for_inference") + ): try: self.model.for_inference() except Exception: From 9d4ccdbff58f0db94a1c5677ae04807897c6916e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 5 Jan 2026 13:29:58 +0000 Subject: [PATCH 3/5] Drop rl.py GRPO changes from this branch --- unsloth/models/rl.py | 86 ++++++++++---------------------------------- 1 file changed, 18 insertions(+), 68 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 23cbbf0256..fd0c69bb0e 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -43,31 +43,10 @@ torch_compile_options = { "triton.cudagraphs": False, } -# 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 trl import __version__ as trl_version from unsloth_zoo.utils import 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") +trl_version = Version(trl_version) def vLLMSamplingParams(**kwargs): @@ -241,7 +220,7 @@ RLTrainer_replacement = ''' import os from typing import * from dataclasses import dataclass, field -from unsloth_zoo.utils import Version +from packaging.version import Version import torch import numpy as np from contextlib import nullcontext @@ -263,18 +242,12 @@ 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) - # Restore previous mode when possible + # Return inference mode if hasattr(self, 'model') and hasattr(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() + self.model.for_inference() # Reset gradient checkpointing buffers to free memory while staying ready for next run try: reset_unsloth_gradient_checkpointing_buffers() @@ -359,32 +332,6 @@ 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 import trl @@ -747,6 +694,19 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): ) RLTrainer_post += training_check + # Sync chat_template from processing_class to vLLM's tokenizer + # This fixes base models that have custom chat templates applied after loading + if "model" in call_args: + vllm_chat_template_sync = ( + "if hasattr(self, 'llm') and self.llm is not None and hasattr(self.llm, 'get_tokenizer'):\n" + " _vllm_tok = self.llm.get_tokenizer()\n" + " _pc = getattr(self, 'processing_class', None) or getattr(self, 'tokenizer', None)\n" + " if _vllm_tok is not None and _pc is not None and getattr(_pc, 'chat_template', None) is not None and getattr(_vllm_tok, 'chat_template', None) is None:\n" + " _vllm_tok.chat_template = _pc.chat_template\n" + "pass\n" + ) + RLTrainer_post += vllm_chat_template_sync + # Edit optional metrics other_metrics_processor = "" if trainer_file in RL_METRICS_CHANGES: @@ -1112,16 +1072,6 @@ 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__) From d5fcc7ddded527c03d642c39a68c7dac93788881 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 5 Jan 2026 13:39:03 +0000 Subject: [PATCH 4/5] Restore TRL version fallback in rl.py --- unsloth/models/rl.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index fd0c69bb0e..11a5215c99 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): From 18f335cec524ab42949c53b0d6659d9e3b85a2c2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 5 Jan 2026 13:39:15 +0000 Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- unsloth/models/rl.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 11a5215c99..2f6aef2709 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -46,10 +46,13 @@ torch_compile_options = { # 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