Add TRL truncation regression and metadata loss fixes (Fixes 1 and 3) (#3971)

* Add TRL truncation regression and metadata loss fixes

Fix 1: TRL 0.24.0-0.25.1 right-truncation regression
- These versions pass max_length=self.max_prompt_length and truncation=True
  to the tokenizer, which right-truncates prompts and strips the assistant
  turn suffix
- Use regex to remove these kwargs from the generated code

Fix 3: Metadata loss for chat_template_kwargs
- TRL 0.24.0+ extracts prompts = [x["prompt"] for x in inputs], losing metadata
  like reasoning_effort
- Inject code to store per-sample chat_template_kwargs on self before extraction
- Preserve these kwargs in prompts_text generation for all TRL versions

Tested with TRL versions 0.22.2, 0.23.1, 0.24.0, 0.25.1, 0.26.2, and 0.27.1.

* Update Fix 1 comment with detailed TRL version behavior explanation

Expand the comment for the TRL 0.24.0-0.25.1 truncation regression fix
to clarify what each TRL version does:

- TRL 0.22.2-0.23.1: Uses truncate_with_protected_tokens() for smart
  truncation that preserves rightmost tokens and protects special tokens
- TRL 0.24.0-0.25.1: Removed smart truncation, passes kwargs directly
  to tokenizer (max_length, truncation=True, add_special_tokens=False)
- TRL 0.26.2+: Removed these kwargs entirely

The fix removes these problematic kwargs so 0.24.0-0.25.1 behaves like
0.26.2+ (no tokenizer-level truncation).

---------

Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com>

unsloth/models/rl_replacements.py

@@ -247,6 +247,30 @@ def grpo_trainer__generate_single_turn(function_name, function):
         "",
         function,
     )
+
+    # TRL 0.24.0-0.25.1 truncation regression fix
+    #
+    # TRL 0.22.2-0.23.1 used smart truncation via truncate_with_protected_tokens():
+    #   - Tokenizes first without truncation
+    #   - Then truncates keeping the RIGHTMOST tokens (preserves assistant turn)
+    #   - Protects special tokens (image_token, vision_start/end) from removal
+    #
+    # TRL 0.24.0-0.25.1 removed this and passed kwargs directly to the tokenizer:
+    #   max_length=self.max_prompt_length, truncation=True, add_special_tokens=False
+    # This causes issues because tokenizer truncation doesn't protect special tokens
+    # and may not preserve the end of the prompt properly.
+    #
+    # TRL 0.26.2+ removed these kwargs entirely (no tokenizer-level truncation).
+    #
+    # Fix: Remove these kwargs so TRL 0.24.0-0.25.1 behaves like 0.26.2+ (no truncation).
+    # This is a no-op for versions that don't have these kwargs (0.22.2-0.23.1, 0.26.2+).
+    for pattern in [
+        r'["\']?max_length["\']?\s*[:=]\s*self\.max_prompt_length\s*,\s*\n?',
+        r'["\']?truncation["\']?\s*[:=]\s*True\s*,\s*\n?',
+        r'["\']?add_special_tokens["\']?\s*[:=]\s*False\s*,\s*\n?',
+    ]:
+        function = re.sub(pattern, "", function)
+
     return function
 
 
@@ -385,6 +409,31 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
 
     function = function.replace(string_to_find, replacement_string)
 
+    # TRL 0.24.0+ extracts prompts = [x["prompt"] for x in inputs], losing metadata
+    # like reasoning_effort. Inject code to store per-sample chat_template_kwargs on self.
+    _metadata_extraction = (
+        "\n"
+        "        # Unsloth: Extract per-sample chat_template_kwargs before metadata is lost\n"
+        "        _ct_ = getattr(self.processing_class, 'chat_template', None) or ''\n"
+        "        _sk_ = {'prompt', 'chosen', 'rejected', 'completion', 'messages', 'label',\n"
+        "                'images', 'image', 'videos', 'video', 'audios', 'audio'}\n"
+        "        self._unsloth_batch_chat_kwargs = []\n"
+        "        for _inp_ in inputs:\n"
+        "            _kw_ = {}\n"
+        "            if isinstance(_inp_, dict):\n"
+        "                for _k_ in _inp_.keys() - _sk_:\n"
+        "                    if _k_ in _ct_ and isinstance(_inp_[_k_], str):\n"
+        "                        _kw_[_k_] = _inp_[_k_]\n"
+        "            self._unsloth_batch_chat_kwargs.append(_kw_)\n"
+    )
+    # Insert after: prompts = [x["prompt"] for x in inputs]
+    _target_line = 'prompts = [x["prompt"] for x in inputs]'
+    if _target_line in function:
+        function = function.replace(
+            _target_line,
+            _target_line + _metadata_extraction,
+        )
+
     # Unsloth: Skip prepare_multimodal_messages when prompts are pre-templated strings.
     # When notebooks pre-apply apply_chat_template(), prompts become strings with image tokens
     # already embedded. Calling prepare_multimodal_messages on strings crashes with TypeError.
@@ -538,9 +587,10 @@ def grpo_trainer_fix_maybe_apply_chat_template(function_name, function):
         _chat_template_ = getattr(self.processing_class, "chat_template", None)
         if _chat_template_ is None: _chat_template_ = ""
         _supported_keys_ = set(("prompt", "chosen", "rejected", "completion", "messages", "label"))
+        _batch_chat_kwargs_ = getattr(self, "_unsloth_batch_chat_kwargs", None)
 
         prompts_text = []
-        for _example_ in __INPUTS__REPLACEMENT__:
+        for _idx_, _example_ in enumerate(__INPUTS__REPLACEMENT__):
             _tokenizer_kwargs_ = {}
             if type(_example_) is not dict:
                 _example_ = {"prompt": _example_}
@@ -550,6 +600,10 @@ def grpo_trainer_fix_maybe_apply_chat_template(function_name, function):
                     v = _example_[k]
                     if type(v) is str:
                         _tokenizer_kwargs_[k] = v
+            if _batch_chat_kwargs_ is not None and _idx_ < len(_batch_chat_kwargs_):
+                for _bk_, _bv_ in _batch_chat_kwargs_[_idx_].items():
+                    if _bk_ not in _tokenizer_kwargs_:
+                        _tokenizer_kwargs_[_bk_] = _bv_
             _x_ = maybe_apply_chat_template(_example_, self.processing_class, **_tokenizer_kwargs_)["prompt"]
             prompts_text.append(_x_)
     """