From 92899dbf38b8187b779eb0fc59da3fb8a76d5a7b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 3 Feb 2026 03:00:12 -0800 Subject: [PATCH] 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 --- unsloth/models/rl_replacements.py | 56 ++++++++++++++++++++++++++++++- 1 file changed, 55 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 2f67c2ebd2..ce83396960 100755 --- a/unsloth/models/rl_replacements.py +++ b/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_) """