From 586a5b046de3b7745fac639903f9ae1a1c84e9ce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 3 Feb 2026 02:03:46 -0800 Subject: [PATCH] Fix Vision GRPO string prompts and OpenEnv async compatibility (#3964) * [fix] Vision GRPO string prompts and OpenEnv async compatibility - Guard prepare_multimodal_messages in GRPO trainer to skip processing when prompts are pre-templated strings. Notebooks that pre-apply apply_chat_template() produce strings with image tokens already embedded; calling prepare_multimodal_messages on those crashes with TypeError. - Apply nest_asyncio when OpenEnv EnvClient exposes async reset/step, so scripts using run_until_complete() wrappers work in all contexts. - Add wrapper to call patch_torchcodec_audio_decoder() from unsloth_zoo for AudioDecoder dict-compatibility. * Add apply_chat_template guard for pre-templated string prompts in Vision GRPO When notebooks pre-apply apply_chat_template, prompts become strings. The existing guard skips prepare_multimodal_messages for strings. This adds a second guard to skip apply_chat_template in the forward_kwargs block, using prompts directly as prompts_text instead. Covers both TRL 0.25.x (no tools param) and TRL 0.26.2+ (with tools=self.tools). Non-matching replacements silently pass for older TRL versions. * Add TRL 0.25.1 single-line variant for apply_chat_template guard TRL 0.25.1 uses single-line formatting for apply_chat_template: apply_chat_template({"prompt": prompt}, ...)["prompt"] While TRL 0.26.2+ uses multi-line formatting: apply_chat_template( {"prompt": prompt}, ... )["prompt"] Add both variants to ensure full backwards compatibility. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/__init__.py | 6 ++ unsloth/import_fixes.py | 42 ++++++++++++++ unsloth/models/rl_replacements.py | 93 +++++++++++++++++++++++++++++++ 3 files changed, 141 insertions(+) mode change 100644 => 100755 unsloth/models/rl_replacements.py diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 0b819a546e..18de06757b 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -134,8 +134,10 @@ from .import_fixes import ( patch_datasets, patch_enable_input_require_grads, fix_openenv_no_vllm, + patch_openspiel_env_async, fix_executorch, patch_vllm_for_notebooks, + patch_torchcodec_audio_decoder, ) fix_xformers_performance_issue() @@ -149,8 +151,10 @@ patch_trackio() patch_datasets() patch_enable_input_require_grads() fix_openenv_no_vllm() +patch_openspiel_env_async() fix_executorch() patch_vllm_for_notebooks() +patch_torchcodec_audio_decoder() del fix_xformers_performance_issue del fix_vllm_aimv2_issue @@ -163,8 +167,10 @@ del patch_trackio del patch_datasets del patch_enable_input_require_grads del fix_openenv_no_vllm +del patch_openspiel_env_async del fix_executorch del patch_vllm_for_notebooks +del patch_torchcodec_audio_decoder # Torch 2.4 has including_emulation if DEVICE_TYPE == "cuda": diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index 89fd152857..715ada16e5 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -812,3 +812,45 @@ def fix_vllm_pdl_blackwell(): else: # Just set the env var - vLLM might be an older version without supports_pdl logger.info(f"Unsloth: Set TRITON_DISABLE_PDL=1 for SM100 ({sm100_gpu_name})") + + +def patch_openspiel_env_async(): + """Apply nest_asyncio for OpenEnv EnvClient async compatibility. + + OpenEnv's EnvClient uses async methods (reset/step). In Jupyter notebooks + these work via top-level await, but converted scripts need + asyncio.get_event_loop().run_until_complete() wrappers. Applying nest_asyncio + ensures nested event loop calls work in all contexts without replacing the + original async methods (which would break scripts that already have their own + sync wrappers). + """ + try: + import inspect + from openenv.core.env_client import EnvClient + + if not inspect.iscoroutinefunction(EnvClient.reset): + return # Already sync, nothing to do + + try: + import nest_asyncio + + nest_asyncio.apply() + logger.info( + "Unsloth: Applied nest_asyncio for OpenEnv EnvClient async compatibility" + ) + except ImportError: + logger.info( + "Unsloth: nest_asyncio not installed, OpenEnv async methods may need manual wrapping" + ) + except (ImportError, AttributeError): + pass # openenv not installed + + +def patch_torchcodec_audio_decoder(): + """Call unsloth_zoo's AudioDecoder patch.""" + try: + from unsloth_zoo.dataset_utils import patch_torchcodec_audio_decoder as _patch + + _patch() + except (ImportError, AttributeError): + pass diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py old mode 100644 new mode 100755 index 67c3f9ded9..2f67c2ebd2 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -385,6 +385,99 @@ def grpo_trainer__generate_and_score_completions(function_name, function): function = function.replace(string_to_find, replacement_string) + # 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. + # Skipping it keeps prompts as strings so TRL uses the non-conversational path, which + # ensures completions are strings and reward functions work correctly. + string_to_find_vision = """ if images is not None: + prompts = [ + prepare_multimodal_messages(prompt, image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ]""" + + replacement_string_vision = """ if images is not None: + # Unsloth: skip prepare_multimodal_messages for pre-templated string prompts + if not prompts or not isinstance(prompts[0], str): + prompts = [ + prepare_multimodal_messages(prompt, image_list) + for prompt, image_list in zip(prompts, images, strict=True) + ]""" + + function = function.replace(string_to_find_vision, replacement_string_vision) + + # Unsloth: Skip apply_chat_template in the forward_kwargs block for pre-templated + # string prompts. When prompts are already strings (from notebooks that pre-applied + # apply_chat_template), calling it again crashes because strings aren't dicts. + # We use prompts directly as prompts_text instead. + + # TRL 0.26.2+ variant (has tools=self.tools) + string_to_find_fwd = """ if images is not None: + prompts_text = [ + apply_chat_template( + {"prompt": prompt}, self.processing_class, tools=self.tools, **self.chat_template_kwargs + )["prompt"] + for prompt in prompts + ]""" + + replacement_string_fwd = """ if images is not None: + # Unsloth: skip apply_chat_template for pre-templated string prompts + if prompts and isinstance(prompts[0], str): + prompts_text = prompts + else: + prompts_text = [ + apply_chat_template( + {"prompt": prompt}, self.processing_class, tools=self.tools, **self.chat_template_kwargs + )["prompt"] + for prompt in prompts + ]""" + + function = function.replace(string_to_find_fwd, replacement_string_fwd) + + # TRL 0.25.x variant (no tools parameter) + string_to_find_fwd_old = """ if images is not None: + prompts_text = [ + apply_chat_template( + {"prompt": prompt}, self.processing_class, **self.chat_template_kwargs + )["prompt"] + for prompt in prompts + ]""" + + replacement_string_fwd_old = """ if images is not None: + # Unsloth: skip apply_chat_template for pre-templated string prompts + if prompts and isinstance(prompts[0], str): + prompts_text = prompts + else: + prompts_text = [ + apply_chat_template( + {"prompt": prompt}, self.processing_class, **self.chat_template_kwargs + )["prompt"] + for prompt in prompts + ]""" + + function = function.replace(string_to_find_fwd_old, replacement_string_fwd_old) + + # TRL 0.25.1 single-line variant (no tools, single-line apply_chat_template call) + string_to_find_fwd_single = """ if images is not None: + prompts_text = [ + apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] + for prompt in prompts + ]""" + + replacement_string_fwd_single = """ if images is not None: + # Unsloth: skip apply_chat_template for pre-templated string prompts + if prompts and isinstance(prompts[0], str): + prompts_text = prompts + else: + prompts_text = [ + apply_chat_template({"prompt": prompt}, self.processing_class, **self.chat_template_kwargs)["prompt"] + for prompt in prompts + ]""" + + function = function.replace( + string_to_find_fwd_single, replacement_string_fwd_single + ) + # This path is for TRL 0.24.0 images is a variable exclusive to this version string_to_find = """ if images is not None: output["num_images"] = num_images"""