Fix TRL 0.25.1+ GRPO vision crash and reward function TypeError by danielhanchen · Pull Request #3975 · unslothai/unsloth

danielhanchen · 2026-02-03T11:43:46Z

Summary

Fix vision GRPO crash on TRL 0.25.1+ when notebooks pre-apply chat templates
Fix reward function TypeError when expecting plain text but receiving conversation format dicts

Changes

Fix 2: Vision GRPO crash (rl.py)

TRL 0.25.1+ calls prepare_multimodal_messages() unconditionally for vision models. When notebooks pre-apply tokenizer.apply_chat_template() (converting prompts to strings), the function crashes iterating over characters.

Solution: Add _patch_prepare_multimodal_messages() that wraps the TRL function with an isinstance(messages, str) guard. String prompts now pass through unchanged.

Fix 6: Reward function TypeError (rl_replacements.py)

TRL 0.25.0+ passes prompts and completions to _calculate_rewards in different formats:

Conversational inputs: list of dicts [{"role": "assistant", "content": "..."}]
Non-conversational inputs: plain strings

This inconsistency causes reward functions to crash when they expect strings but receive dicts (or vice versa).

Solution: Add grpo_trainer__calculate_rewards_text_fix() that makes _calculate_rewards always use prompts_text and completions_text (plain decoded strings) for consistent behavior.

Test plan

Verified Fix 2: prepare_multimodal_messages("test string", []) returns string unchanged
Verified Fix 6: Compiled cache shows _calculate_rewards uses prompts_text, completions_text
Smoke tested nb2_gpt_oss_2048 with TRL 0.25.1 - runs without TypeError
Smoke tested vision model loading with TRL 0.25.1 - works correctly
Re-applied on latest main and re-confirmed all tests pass

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

for more information, see https://pre-commit.ci

gemini-code-assist · 2026-02-03T11:44:05Z

Summary of Changes

Hello @danielhanchen, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces two crucial compatibility fixes for the TRL library versions 0.25.0 and newer. It resolves issues where pre-applied chat templates led to crashes in vision models and eliminates TypeError exceptions in reward functions by standardizing input formats. These changes enhance the robustness and predictability of the system when interacting with updated TRL functionalities.

Highlights

Vision GRPO crash fix for TRL 0.25.1+: Addresses a crash in TRL 0.25.1+ vision models when prepare_multimodal_messages is called unconditionally on prompts that have already had chat templates applied. A new patch adds an isinstance(messages, str) guard to ensure string prompts are passed through unchanged, preventing crashes.
Reward function TypeError fix for TRL 0.25.0+: Resolves a TypeError in reward functions caused by inconsistent input formats (list of dicts for conversational vs. plain strings for non-conversational) for prompts and completions in TRL 0.25.0+. The fix ensures _calculate_rewards always receives plain decoded text strings (prompts_text, completions_text) for consistent behavior.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

unsloth/models/rl.py
- Added _patch_prepare_multimodal_messages function to introduce a string type guard for trl.data_utils.prepare_multimodal_messages, preventing crashes when chat templates are pre-applied.
- Integrated the _patch_prepare_multimodal_messages call into PatchFastRL to ensure the fix is applied during the RL patching process.
unsloth/models/rl_replacements.py
- Introduced grpo_trainer__calculate_rewards_text_fix to modify the _generate_and_score_completions function. This modification ensures that _calculate_rewards consistently uses prompts_text and completions_text (plain decoded strings), thereby preventing TypeError issues caused by mixed input formats.

Activity

Verified Fix 2: prepare_multimodal_messages("test string", []) returns the string unchanged.
Verified Fix 6: Confirmed through compiled cache that _calculate_rewards utilizes prompts_text and completions_text.
Smoke tested nb2_gpt_oss_2048 with TRL 0.25.1, confirming it runs without TypeError.
Smoke tested vision model loading with TRL 0.25.1, confirming correct functionality.
Re-applied changes on the latest main branch and re-confirmed all tests pass.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces two important fixes for compatibility with recent versions of TRL, addressing a crash in vision models with pre-applied chat templates and a TypeError in the reward function due to inconsistent data formats. The changes are well-implemented and follow the existing patching patterns in the codebase. The logic is sound and directly solves the described issues. I have one suggestion to improve the robustness of a string replacement by using a regular expression, which is consistent with other parts of the code.

gemini-code-assist · 2026-02-03T11:45:13Z

unsloth/models/rl_replacements.py

+        function = function.replace(
+            "self._calculate_rewards(inputs, prompts, completions, completion_ids_list)",
+            "self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list)",
+        )


The use of function.replace for patching the _calculate_rewards call is a bit fragile as it's sensitive to whitespace changes in the upstream TRL library. Using re.sub would make this patch more robust against minor formatting variations. This is also consistent with how other patches are implemented in this file.

Suggested change

function = function.replace(

"self._calculate_rewards(inputs, prompts, completions, completion_ids_list)",

"self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list)",

)

function = re.sub(

r"self\._calculate_rewards\(\s*inputs,\s*prompts,\s*completions,\s*completion_ids_list\s*\)",

"self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list)",

function,

)

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7e7184e783

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-02-03T11:45:16Z

unsloth/models/rl.py

+    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


Preserve image placeholders when messages are strings

The new guard returns the original string whenever messages is a str, regardless of any images passed via *args/**kwargs. In TRL, prepare_multimodal_messages is responsible for inserting image placeholders into the prompt; bypassing it means any vision inputs provided alongside a pre-applied chat template are silently ignored, so GRPO will train on text-only prompts. This regression shows up whenever a notebook pre-applies apply_chat_template() but still supplies images to the vision model.

Useful? React with 👍 / 👎.

chatgpt-codex-connector · 2026-02-03T11:45:16Z

unsloth/models/rl_replacements.py

+        function = function.replace(
+            "self._calculate_rewards(inputs, prompts, completions, completion_ids_list)",
+            "self._calculate_rewards(inputs, prompts_text, completions_text, completion_ids_list)",


Keep structured prompts available for reward functions

This replacement forces _calculate_rewards to always receive prompts_text/completions_text even in conversational mode. Reward functions that intentionally inspect message structure (e.g., roles, tool-call fields, or metadata in dicts) will now receive flattened strings and can no longer operate correctly. That’s a behavioral regression for existing conversational reward functions on TRL 0.25.0+.

Useful? React with 👍 / 👎.

danielhanchen and others added 2 commits February 3, 2026 11:43

[pre-commit.ci] auto fixes from pre-commit.com hooks

7e7184e

for more information, see https://pre-commit.ci

gemini-code-assist bot reviewed Feb 3, 2026

View reviewed changes

chatgpt-codex-connector bot reviewed Feb 3, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Comments

Fix TRL 0.25.1+ GRPO vision crash and reward function TypeError#3975

Fix TRL 0.25.1+ GRPO vision crash and reward function TypeError#3975
danielhanchen wants to merge 2 commits intomainfrom
fix/trl-025-grpo-reward-vision-v2

danielhanchen commented Feb 3, 2026

Uh oh!

gemini-code-assist bot commented Feb 3, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Feb 3, 2026

Uh oh!

chatgpt-codex-connector bot left a comment

Uh oh!

chatgpt-codex-connector bot Feb 3, 2026

Uh oh!

chatgpt-codex-connector bot Feb 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Comments

Conversation

danielhanchen commented Feb 3, 2026

Summary

Changes

Fix 2: Vision GRPO crash (rl.py)

Fix 6: Reward function TypeError (rl_replacements.py)

Test plan

Uh oh!

gemini-code-assist bot commented Feb 3, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector bot Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector bot Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant