From b9ea947f2e426b6136b88d3fc02a302f1fc0c7a6 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 3 Feb 2026 11:35:25 +0000 Subject: [PATCH] Add Fix 2 and Fix 6 for TRL 0.25.1+ compatibility Fix 2 (rl.py): Add _patch_prepare_multimodal_messages() - Wraps prepare_multimodal_messages with isinstance(messages, str) guard - Fixes vision GRPO crash when notebooks pre-apply chat templates - String prompts now pass through unchanged Fix 6 (rl_replacements.py): Add grpo_trainer__calculate_rewards_text_fix() - Makes _calculate_rewards use prompts_text/completions_text for TRL 0.25.0+ - Ensures reward functions receive consistent plain text format - Fixes TypeError when reward functions expect strings but get dicts Also adds fallback for device_synchronize import for older unsloth_zoo versions. --- unsloth/models/rl.py | 37 ++++++++++++++++++++++++++++ unsloth/models/rl_replacements.py | 40 ++++++++++++++++++++++++++++++- 2 files changed, 76 insertions(+), 1 deletion(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index 6fa07bdd18..32621ed31f 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1550,10 +1550,47 @@ def patch_trl_openenv(): return +def _patch_prepare_multimodal_messages(): + """Fix 2: TRL >= 0.25.1 calls prepare_multimodal_messages unconditionally for + vision models. When notebooks pre-apply chat templates (converting prompts to strings), + the function crashes iterating over characters. This patch adds isinstance(messages, str) + guard to return strings unchanged.""" + try: + import trl.data_utils as _du + except ImportError: + return + + _original = getattr(_du, "prepare_multimodal_messages", None) + if _original is None: + return + if getattr(_original, "_unsloth_patched", False): + return + + def _safe_prepare_multimodal_messages(messages, *args, **kwargs): + # If messages is already a string (pre-applied chat template), return as-is + if isinstance(messages, str): + return messages + return _original(messages, *args, **kwargs) + + _safe_prepare_multimodal_messages._unsloth_patched = True + _du.prepare_multimodal_messages = _safe_prepare_multimodal_messages + + # Also patch in grpo_trainer module if imported + try: + import trl.trainer.grpo_trainer as _gt + if hasattr(_gt, "prepare_multimodal_messages"): + _gt.prepare_multimodal_messages = _safe_prepare_multimodal_messages + except ImportError: + pass + + logger.info("Unsloth: Patched prepare_multimodal_messages with string guard") + + def PatchFastRL(algorithm = None, FastLanguageModel = None): if FastLanguageModel is not None: PatchRL(FastLanguageModel) patch_trl_rl_trainers() patch_trl_openenv() + _patch_prepare_multimodal_messages() if type(algorithm) is str and algorithm.islower(): PatchRLStatistics(algorithm) diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index ce83396960..06969227b4 100755 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -29,7 +29,16 @@ from unsloth_zoo.rl_replacements import RL_REPLACEMENTS, left_pack_padding from unsloth_zoo.utils import Version from importlib.metadata import version as importlib_version from unsloth_zoo.log import logger -from unsloth_zoo.device_type import device_synchronize +# device_synchronize may not exist in older unsloth_zoo versions +try: + from unsloth_zoo.device_type import device_synchronize +except ImportError: + import torch + def device_synchronize(): + if torch.cuda.is_available(): + torch.cuda.synchronize() + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + torch.xpu.synchronize() import importlib.util from ..device_type import ( is_hip, @@ -577,6 +586,35 @@ def grpo_trainer__generate_and_score_completions(function_name, function): RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__generate_and_score_completions) +# Fix 6: TRL 0.25.0+ _calculate_rewards text arguments +# +# TRL 0.25.0+ passes `prompts` and `completions` to _calculate_rewards in different formats: +# - For conversational inputs: list of dicts [{"role": "assistant", "content": "..."}] +# - For non-conversational inputs: plain text strings +# +# This inconsistency causes reward functions to fail when they expect one format but get the other. +# The variables `prompts_text` and `completions_text` always contain plain decoded text strings. +# +# Fix: Always pass plain text (prompts_text, completions_text) to _calculate_rewards for consistency. +# This ensures reward functions receive predictable string format regardless of conversational mode. +def grpo_trainer__calculate_rewards_text_fix(function_name, function): + if function_name != "_generate_and_score_completions": + return function + + # Only apply if prompts_text and completions_text exist (TRL 0.25.0+) + if "prompts_text" in function and "completions_text" in function: + # Replace the _calculate_rewards call to use text versions + function = function.replace( + "self._calculate_rewards(inputs, prompts, completions, completion_ids_list)", + "self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list)", + ) + + return function + + +RL_FUNCTIONS["grpo_trainer"].append(grpo_trainer__calculate_rewards_text_fix) + + # Fix {"reasoning_effort" : "high"} not applied def grpo_trainer_fix_maybe_apply_chat_template(function_name, function): spaces = function.find("def ")