Improve TRL compatibility and GRPO state restore

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,

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(

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(

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(

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(

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

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__)

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