Add Fix 2 and Fix 6 for TRL 0.25.1+ GRPO 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 (cherry picked from commit 7a8ce2ef1e5ef3ce5524ce489f710e47c31c98b4)
This commit is contained in:
parent
bced78373f
commit
85bcfe166d
2 changed files with 63 additions and 10 deletions
|
|
@ -1888,16 +1888,39 @@ def patch_trl_vllm_generation():
|
|||
return
|
||||
|
||||
|
||||
def patch_trl_vllm_generation():
|
||||
# trl moved vllm stuff to trl/generation/vllm_generation.py
|
||||
# We need to min_p patch it to not instantiate another vLLM instance if we already have one with fast_inference
|
||||
# Find the instance of self.llm = LLM(..) (multiline) and wrap it around an if clause
|
||||
for function in RL_ADDITIONAL_FUNCTIONS["vllm_generation"]:
|
||||
logger.info(
|
||||
f"Unsloth: Patching trl VLLMGeneration with function: {function.__name__}"
|
||||
)
|
||||
function()
|
||||
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):
|
||||
|
|
@ -1906,5 +1929,6 @@ def PatchFastRL(algorithm = None, FastLanguageModel = None):
|
|||
patch_trl_rl_trainers()
|
||||
patch_trl_openenv()
|
||||
patch_trl_vllm_generation()
|
||||
_patch_prepare_multimodal_messages()
|
||||
if type(algorithm) is str and algorithm.islower():
|
||||
PatchRLStatistics(algorithm)
|
||||
|
|
|
|||
|
|
@ -537,6 +537,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 ")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue