ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix Vision GRPO string prompts and OpenEnv async compatibility (#3964)

Browse source 
* [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 <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-02-03 02:03:46 -08:00 • committed by GitHub
commit 586a5b046d
3 changed files with 141 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/__init__.py
View file

@ -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":

42
unsloth/import_fixes.py
View file

@ -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

93
unsloth/models/rl_replacements.py Normal file → Executable file
View file

@ -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"""