From e478faef18442c1467d7452bb8bc1bada967756d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Fri, 12 Dec 2025 04:58:43 -0800
Subject: [PATCH] Fixes

---
 unsloth/import_fixes.py           | 2 +-
 unsloth/models/rl_replacements.py | 8 +++++++-
 unsloth/trainer.py                | 5 +++--
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py
index a43be23194..4cbe57cd95 100644
--- a/unsloth/import_fixes.py
+++ b/unsloth/import_fixes.py
@@ -282,7 +282,7 @@ def patch_enable_input_require_grads():
     # Ref: https://github.com/huggingface/transformers/pull/41993/files#diff-6b72b98c4c2dcfc6cc606843917733f5d858374fbc22a735ff483bbc0c1e63eaL1979-R1996
     try:
         original_source = inspect.getsource(PreTrainedModel.enable_input_require_grads)
-    except (OSError, TypeError):
+    except:
         return
 
     # Only patch if the new pattern exists (iterating over self.modules())
diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py
index 2cf3527c9b..7dab0d7307 100644
--- a/unsloth/models/rl_replacements.py
+++ b/unsloth/models/rl_replacements.py
@@ -26,7 +26,9 @@ import torch
 import inspect
 from collections import defaultdict
 from unsloth_zoo.rl_replacements import RL_REPLACEMENTS, left_pack_padding
+from unsloth_zoo.utils import Version
 from unsloth_zoo.log import logger
+import importlib.util
 from ..device_type import (
     is_hip,
     get_device_type,
@@ -942,11 +944,15 @@ def openenv_vllm_reload_weights():
     #
     # The fix: Use wake_up() with no tags, which wakes everything. Unsloth's patched
     # CuMemAllocator.wake_up skips weights anyway, so this is safe.
+    if importlib.util.find_spec("trl") is None:
+        return
+    if Version(importlib_version("trl")) < Version("0.26.0"):
+        return
     try:
         import trl.experimental.openenv.utils as openenv_utils
         import trl.experimental.openenv as openenv
     except ImportError as e:
-        logger.warning(f"Unsloth: Failed to import trl openenv: {e}")
+        logger.info(f"Unsloth: Failed to import trl openenv: {e}")
         return
 
     src = inspect.getsource(openenv_utils.generate_rollout_completions)
diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index 339af63f33..5cd1bfd08d 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -36,7 +36,7 @@ from unsloth_zoo.vision_utils import (
     UnslothVisionDataCollator,
 )
 from unsloth_zoo.hf_utils import get_transformers_model_type
-from packaging.version import Version
+from unsloth_zoo.utils import Version
 import dataclasses
 
 __all__ = [
@@ -315,10 +315,11 @@ def _patch_sft_trainer_auto_packing(trl_module):
 
         # We also disable vision language models for padding free collators
         blocked = (
-            data_collator is not None
+            (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: