vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-11-09 07:39:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35

GRPO: train on video prompts instead of scoring them as text (#13144)

Browse source 
* GRPO: train on video prompts instead of scoring them as text

TRL builds the logprob forward kwargs from image columns only, so a
{"type": "video"} prompt was generated against its video and then scored
without it. Rebuild the video tensors with the processor's chat template
(once per distinct prompt), merge them into forward_kwargs, keep them with
their samples across the shuffle, and strip the None keys Arrow adds to
content parts, which made Qwen2-VL templates render extra video
placeholders and crash with StopIteration.

* Run the video helper tests against any installed unsloth_zoo

* Strip Arrow video keys from text rows too; never dedup in-memory videos by repr

* Render the video rebuild with GRPOConfig.chat_template like generation does

* Refuse vLLM video before generation and inline image parts beside video

* Tighten video helper docstrings

* Probe prompts at the left-pad anchor like images, for synthetic and floor sources
This commit is contained in:
Daniel Han 2026-10-09 12:29:48 -07:00 • committed by GitHub
parent 371da3d3bb
commit 6c723f7477
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 536 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
tests/test_grpo_legacy_trl_vision.py
View file

@ -140,7 +140,9 @@ def _run_legacy(processor_output, inputs, completion_ids):
import torch
from unsloth.models.rl_replacements import (
_unsloth_grpo_clean_video_prompts,
_unsloth_grpo_image_cell,
_unsloth_grpo_video_inputs,
_unsloth_grpo_vision_inputs,
_unsloth_reject_grpo_image_list,
grpo_trainer__generate_and_score_completions,
@ -153,6 +155,8 @@ def _run_legacy(processor_output, inputs, completion_ids):
"torch": torch,
"_unsloth_grpo_vision_inputs": _unsloth_grpo_vision_inputs,
"_unsloth_grpo_image_cell": _unsloth_grpo_image_cell,
"_unsloth_grpo_clean_video_prompts": _unsloth_grpo_clean_video_prompts,
"_unsloth_grpo_video_inputs": _unsloth_grpo_video_inputs,
# Injected on this TRL too now: its vLLM server path cannot carry a multi image row,
# so the guard is handed the trainer and decides at runtime.
"_unsloth_reject_grpo_image_list": _unsloth_reject_grpo_image_list,

279
tests/test_grpo_video_inputs.py Normal file
View file

@ -0,0 +1,279 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Video prompts in GRPO reach the old, reference and policy logprob forwards.
TRL builds those forward kwargs from `image` / `images` only, so a `{"type": "video"}` prompt
was generated against its video and then scored as text. unslothai/unsloth#3357."""
import inspect
import types
import pytest
import torch
def _video_prompt(
path,
text = "What happens?",
arrow = False,
):
video = {"type": "video", "video": path}
words = {"type": "text", "text": text}
if arrow:
# Arrow gives every part the union of the keys.
video["text"] = None
words["video"] = None
return [{"role": "user", "content": [video, words]}]
def test_arrow_none_keys_are_stripped_from_video_prompts_only():
from unsloth.models.rl_replacements import _unsloth_grpo_clean_video_prompts
text_only = [{"role": "user", "content": [{"type": "text", "text": "hi", "image": None}]}]
video = _video_prompt("a.mp4", arrow = True)
cleaned = _unsloth_grpo_clean_video_prompts([text_only, video])
assert cleaned[0] is text_only
assert cleaned[1][0]["content"] == [
{"type": "video", "video": "a.mp4"},
{"type": "text", "text": "What happens?"},
]
assert video[0]["content"][1]["video"] is None, "the dataset row was mutated"
# a text row of a column that also holds videos carries the key too
text_row = [{"role": "user", "content": [{"type": "text", "text": "hi", "video": None}]}]
assert _unsloth_grpo_clean_video_prompts([text_row])[0][0]["content"] == [
{"type": "text", "text": "hi"}
]
batch = [text_only]
assert _unsloth_grpo_clean_video_prompts(batch) is batch
class _Processor:
"""apply_chat_template stand in: video v of a prompt gets grid (1, 2, 2) * (v + 1) rows
filled with the prompt's id, so the expansion order can be read back."""
def __init__(self):
self.calls = []
def apply_chat_template(self, conversation, **kwargs):
self.calls.append(conversation)
self.kwargs = kwargs
grids, rows, seconds = [], [], []
for prompt in conversation:
pid = int(prompt[0]["content"][-1]["text"])
n = sum(1 for p in prompt[0]["content"] if p.get("type") == "video")
for v in range(n):
grid = torch.tensor([1, 2, 2 * (v + 1)])
grids.append(grid)
rows.append(torch.full((int(grid.prod()), 3), float(pid)))
seconds.append(float(pid) + v / 10)
return {
"input_ids": torch.zeros(len(conversation), 4, dtype = torch.long),
"pixel_values_videos": torch.cat(rows),
"video_grid_thw": torch.stack(grids),
"second_per_grid_ts": torch.tensor(seconds),
}
def _trainer(processor, use_vllm = False):
return types.SimpleNamespace(
processing_class = processor,
use_vllm = use_vllm,
chat_template_kwargs = {},
accelerator = types.SimpleNamespace(device = torch.device("cpu")),
)
@pytest.fixture
def zoo_with_video_keys(monkeypatch):
"""The helper refuses a zoo whose key tuple lacks video; present one that has it."""
zoo = pytest.importorskip("unsloth_zoo.rl_replacements")
keys = tuple(getattr(zoo, "GRPO_VISION_KEYS", ()))
if "pixel_values_videos" not in keys:
keys = keys + ("pixel_values_videos", "video_grid_thw", "second_per_grid_ts", "num_videos")
monkeypatch.setattr(zoo, "GRPO_VISION_KEYS", keys, raising = False)
def _two_video_prompt(pid):
return [
{
"role": "user",
"content": [
{"type": "video", "video": "a.mp4"},
{"type": "video", "video": "b.mp4"},
{"type": "text", "text": str(pid)},
],
}
]
def test_each_distinct_prompt_is_decoded_once_and_repeated_per_generation(zoo_with_video_keys):
from unsloth.models.rl_replacements import _unsloth_grpo_video_inputs
processor = _Processor()
one = _video_prompt("a.mp4", text = "1")
two = _two_video_prompt(2)
prompts = [one, one, two, two] # num_generations = 2
out = _unsloth_grpo_video_inputs(_trainer(processor), prompts)
assert len(processor.calls) == 1 and len(processor.calls[0]) == 2
assert out["num_videos"] == [1, 1, 2, 2]
assert out["video_grid_thw"].tolist() == [
[1, 2, 2],
[1, 2, 2],
[1, 2, 2],
[1, 2, 4],
[1, 2, 2],
[1, 2, 4],
]
owners = out["pixel_values_videos"][:, 0].tolist()
assert owners == [1.0] * 8 + [2.0] * 12 + [2.0] * 12
assert out["second_per_grid_ts"].tolist() == pytest.approx([1.0, 1.0, 2.0, 2.1, 2.0, 2.1])
def test_in_memory_videos_are_never_shared_between_prompts(zoo_with_video_keys):
import numpy as np
from unsloth.models.rl_replacements import _unsloth_grpo_video_inputs
processor = _Processor()
a = np.zeros((64, 64, 3))
b = np.zeros((64, 64, 3))
b[32, 32, 0] = 1.0 # same repr() as a
assert repr(a) == repr(b)
prompts = [_video_prompt(a, text = "1"), _video_prompt(b, text = "1")]
_unsloth_grpo_video_inputs(_trainer(processor), prompts)
assert len(processor.calls[0]) == 2
def test_the_trainers_chat_template_is_used(zoo_with_video_keys):
from unsloth.models.rl_replacements import _unsloth_grpo_video_inputs
processor = _Processor()
trainer = _trainer(processor)
trainer.chat_template = "{{ messages }}"
_unsloth_grpo_video_inputs(trainer, [_video_prompt("a.mp4", text = "1")])
assert processor.kwargs["chat_template"] == "{{ messages }}"
def test_a_batch_without_video_costs_nothing():
from unsloth.models.rl_replacements import _unsloth_grpo_video_inputs
processor = _Processor()
text = [{"role": "user", "content": [{"type": "text", "text": "1"}]}]
assert _unsloth_grpo_video_inputs(_trainer(processor), [text, text]) is None
assert _unsloth_grpo_video_inputs(_trainer(processor), ["plain text prompt"]) is None
assert processor.calls == []
def test_vllm_and_mixed_image_video_batches_are_refused(zoo_with_video_keys):
from unsloth.models.rl_replacements import _unsloth_grpo_video_inputs
prompts = [_video_prompt("a.mp4", text = "1")]
with pytest.raises(NotImplementedError, match = "fast_inference"):
_unsloth_grpo_video_inputs(_trainer(_Processor(), use_vllm = True), prompts)
with pytest.raises(NotImplementedError, match = "image and video"):
_unsloth_grpo_video_inputs(_trainer(_Processor()), prompts, images = [["img"]])
inline = _video_prompt("a.mp4", text = "1")
inline[0]["content"].insert(0, {"type": "image", "image": "x.png"})
with pytest.raises(NotImplementedError, match = "image and video"):
_unsloth_grpo_video_inputs(_trainer(_Processor()), [inline])
def test_vllm_is_refused_before_generation():
from unsloth.models.rl_replacements import _unsloth_grpo_clean_video_prompts
trainer = types.SimpleNamespace(use_vllm = True)
with pytest.raises(NotImplementedError, match = "fast_inference"):
_unsloth_grpo_clean_video_prompts([_video_prompt("a.mp4")], trainer)
text = [[{"role": "user", "content": [{"type": "text", "text": "hi"}]}]]
assert _unsloth_grpo_clean_video_prompts(text, trainer) is text
def test_an_old_zoo_is_refused_rather_than_training_on_text(monkeypatch):
zoo = pytest.importorskip("unsloth_zoo.rl_replacements")
from unsloth.models.rl_replacements import _unsloth_grpo_video_inputs
monkeypatch.setattr(zoo, "GRPO_VISION_KEYS", ("pixel_values", "image_grid_thw"), raising = False)
with pytest.raises(RuntimeError, match = "upgrade"):
_unsloth_grpo_video_inputs(_trainer(_Processor()), [_video_prompt("a.mp4", text = "1")])
def test_video_rows_survive_the_shuffle_and_the_slice():
trl_utils = pytest.importorskip("trl.trainer.utils")
from unsloth.models.rl_replacements import (
_unsloth_grpo_split_vision_by_sample,
_unsloth_grpo_unsplit_vision,
)
# sample 0: one 4 row video, sample 1: 4 + 8 rows, sample 2: one 4 row video
grid = torch.tensor([[1, 2, 2], [1, 2, 2], [1, 2, 4], [1, 2, 2]])
owners = torch.tensor([0] * 4 + [1] * 12 + [2] * 4)
batch = {
"prompt_ids": torch.arange(3).unsqueeze(-1),
"advantages": torch.zeros(3),
"num_videos": [1, 2, 1],
"pixel_values_videos": owners.reshape(-1, 1).float(),
"video_grid_thw": grid,
"second_per_grid_ts": torch.tensor([0.0, 1.0, 1.5, 2.0]),
}
torch.manual_seed(0)
split = _unsloth_grpo_split_vision_by_sample(batch)
assert isinstance(split["pixel_values_videos"], list)
shuffled = trl_utils.shuffle_sequence_dict(split)
for half in trl_utils.split_tensor_dict(shuffled, 3):
restored = _unsloth_grpo_unsplit_vision(trl_utils.unsplit_pixel_values_by_grid(half))
sample = restored["prompt_ids"].item()
n = restored["num_videos"][0]
assert n == [1, 2, 1][sample]
assert restored["video_grid_thw"].shape[0] == n
rows = int(restored["video_grid_thw"].prod(-1).sum())
assert restored["pixel_values_videos"].flatten().tolist() == [sample] * rows
expected = {0: [0.0], 1: [1.0, 1.5], 2: [2.0]}[sample]
assert restored["second_per_grid_ts"].tolist() == expected
def test_the_trl_trainer_gets_the_videos_merged_into_its_forward_kwargs():
import ast
import importlib.util
spec = importlib.util.find_spec("trl.trainer.grpo_trainer")
if spec is None or spec.origin is None:
pytest.skip("trl is not installed")
from unsloth.models.rl_replacements import grpo_trainer__generate_and_score_completions
# From the file: once unsloth is imported, trl.GRPOTrainer is the patched class.
with open(spec.origin, "r", encoding = "utf-8") as fh:
module_source = fh.read()
node = next(
n
for n in ast.walk(ast.parse(module_source))
if isinstance(n, ast.FunctionDef) and n.name == "_generate_and_score_completions"
)
source = " " + ast.get_source_segment(module_source, node)
if "forward_kwargs = {}" not in source:
pytest.skip("this TRL has no forward_kwargs (< 0.24.0)")
patched = grpo_trainer__generate_and_score_completions(
"_generate_and_score_completions", source
)
assert "prompts = _unsloth_grpo_clean_video_prompts(prompts, self)" in patched
assert "_unsloth_video_kwargs = _unsloth_grpo_video_inputs(" in patched
merge = patched.index("forward_kwargs = {**forward_kwargs, **_unsloth_video_kwargs}")
old_logps = patched.index("_get_per_token_logps_and_entropies(")
assert merge < old_logps, "the videos arrive after the old and reference logprobs"
def test_both_logprob_passes_treat_video_rows_as_vision_rows():
# Left packing and sequence packing move tokens, which breaks M-RoPE video positions.
from unsloth.models import rl_replacements
source = inspect.getsource(rl_replacements.grpo_trainer__get_per_token_logps_and_entropies)
sentinel = source.index('pixel_values = vision_inputs.get("pixel_values_videos", None)')
assert sentinel < source.index(
"if pixel_values is None:\n left_pad_tokens_per_prompt"
)
for fn in (
rl_replacements.grpo_trainer__get_per_token_logps_and_entropies,
rl_replacements.grpo_trainer_compute_loss,
):
assert "video_grid_thw is not None" in inspect.getsource(fn), fn.__name__

1
tests/version_compat/test_grpo_left_pad_probe_at_the_floor.py
View file

@ -93,6 +93,7 @@ def _run(source: str, inputs: dict):
"_unsloth_reject_grpo_image_list": lambda *a, **k: None,
"_unsloth_grpo_image_cell": lambda cell: cell if isinstance(cell, list) else [cell],
"_unsloth_grpo_vision_inputs": lambda inputs: {},
"_unsloth_grpo_video_inputs": lambda *args, **kwargs: None,
"sanitize_logprob": lambda x: x,
}
exec(

256
unsloth/models/rl_replacements.py
View file

@ -1098,6 +1098,10 @@ def _unsloth_grpo_vision_inputs(source):
"num_images",
"token_type_ids",
"mm_token_type_ids",
"pixel_values_videos",
"video_grid_thw",
"second_per_grid_ts",
"num_videos",
)
}
@ -1120,6 +1124,7 @@ def _unsloth_grpo_split_vision_by_sample(batch):
except TypeError:
return None
batch = _unsloth_grpo_split_videos_by_sample(batch)
pixel_values = batch.get("pixel_values", None)
if pixel_values is None:
return batch
@ -1218,6 +1223,188 @@ def _unsloth_grpo_split_vision_by_sample(batch):
return split
def _unsloth_grpo_split_videos_by_sample(batch):
"""Group the flat video rows, grid and second_per_grid_ts by sample (num_videos) so the
_prepare_inputs shuffle keeps them with their sample; _unsloth_grpo_unsplit_vision merges back."""
pixel_values_videos = batch.get("pixel_values_videos", None)
video_grid_thw = batch.get("video_grid_thw", None)
num_videos = batch.get("num_videos", None)
if (
pixel_values_videos is None
or video_grid_thw is None
or num_videos is None
or isinstance(pixel_values_videos, list)
):
return batch
num_videos = [int(n) for n in num_videos]
if sum(num_videos) != video_grid_thw.shape[0]:
return batch
rows_per_sample = [
int(group.prod(dim = -1).sum()) for group in torch.split(video_grid_thw, num_videos)
]
split = dict(batch)
split["pixel_values_videos"] = list(torch.split(pixel_values_videos, rows_per_sample))
split["video_grid_thw"] = list(torch.split(video_grid_thw, num_videos))
second_per_grid_ts = batch.get("second_per_grid_ts", None)
if second_per_grid_ts is not None:
second_per_grid_ts = torch.as_tensor(second_per_grid_ts)
split["second_per_grid_ts"] = list(torch.split(second_per_grid_ts, num_videos))
return split
def _unsloth_grpo_prompt_videos(prompt, kind = "video"):
"""Number of `kind` parts in a conversational prompt."""
if not isinstance(prompt, list):
return 0
count = 0
for message in prompt:
content = message.get("content", None) if isinstance(message, dict) else None
if isinstance(content, list):
count += sum(
1 for part in content if isinstance(part, dict) and part.get("type") == kind
)
return count
def _unsloth_grpo_has_video_key(prompt):
if not isinstance(prompt, list):
return False
for message in prompt:
content = message.get("content", None) if isinstance(message, dict) else None
if isinstance(content, list) and any(
isinstance(part, dict) and "video" in part for part in content
):
return True
return False
def _unsloth_grpo_clean_video_prompts(prompts, trainer = None):
"""Arrow adds `video: None` to every content part of a column holding videos; Qwen2-VL
templates test `'video' in content` and render a placeholder per part (StopIteration)."""
if not any(_unsloth_grpo_has_video_key(prompt) for prompt in prompts):
return prompts
if getattr(trainer, "use_vllm", False) and any(
_unsloth_grpo_prompt_videos(prompt) for prompt in prompts
):
# Before generation, which vLLM would otherwise run without the videos.
raise NotImplementedError(
"Unsloth: video GRPO runs on the transformers generation path only. "
"Load the model with fast_inference = False (use_vllm = False)."
)
cleaned = []
for prompt in prompts:
if not _unsloth_grpo_has_video_key(prompt):
cleaned.append(prompt)
continue
messages = []
for message in prompt:
content = message.get("content", None)
if isinstance(content, list):
message = dict(message)
message["content"] = [
{k: v for k, v in part.items() if v is not None}
if isinstance(part, dict)
else part
for part in content
]
messages.append(message)
cleaned.append(messages)
return cleaned
def _unsloth_grpo_prompt_key(value):
"""Dedup key: equal for the repeated copies of a prompt, never shared by different media.
repr() truncates arrays, so in-memory frames are keyed by identity, not by content."""
if isinstance(value, dict):
return tuple((k, _unsloth_grpo_prompt_key(v)) for k, v in value.items())
if isinstance(value, (list, tuple)):
return tuple(_unsloth_grpo_prompt_key(v) for v in value)
if value is None or isinstance(value, (str, int, float, bool)):
return value
return ("id", id(value))
def _unsloth_grpo_video_inputs(
trainer,
prompts,
images = None,
):
"""Video tensors for the logprob forwards, which TRL builds from images only. Rendered with
the same chat template call generation makes, once per distinct prompt. None without video."""
num_videos = [_unsloth_grpo_prompt_videos(prompt) for prompt in prompts]
if not any(num_videos):
return None
if getattr(trainer, "use_vllm", False):
raise NotImplementedError(
"Unsloth: video GRPO runs on the transformers generation path only. "
"Load the model with fast_inference = False (use_vllm = False)."
)
if (images is not None and any(images)) or any(
_unsloth_grpo_prompt_videos(prompt, "image") for prompt in prompts
):
raise NotImplementedError(
"Unsloth: GRPO does not support image and video inputs in the same batch yet."
)
try:
from unsloth_zoo.rl_replacements import GRPO_VISION_KEYS
except Exception:
GRPO_VISION_KEYS = ()
if "pixel_values_videos" not in GRPO_VISION_KEYS:
raise RuntimeError(
"Unsloth: video GRPO needs a newer unsloth_zoo that forwards video inputs to "
"both GRPO logprob passes. Please upgrade: pip install -U unsloth_zoo"
)
unique, order = {}, []
for prompt in prompts:
order.append(unique.setdefault(_unsloth_grpo_prompt_key(prompt), len(unique)))
first = {}
for i, u in enumerate(order):
first.setdefault(u, i)
unique_prompts = [prompts[first[u]] for u in range(len(unique))]
template_kwargs = dict(getattr(trainer, "chat_template_kwargs", None) or {})
if getattr(trainer, "chat_template", None) is not None:
# GRPOConfig.chat_template, which generation renders with too
template_kwargs["chat_template"] = trainer.chat_template
processed = trainer.processing_class.apply_chat_template(
conversation = unique_prompts,
add_generation_prompt = True,
tokenize = True,
return_dict = True,
padding = True,
return_tensors = "pt",
**template_kwargs,
)
pixel_values_videos = processed.get("pixel_values_videos", None)
video_grid_thw = processed.get("video_grid_thw", None)
if pixel_values_videos is None or video_grid_thw is None:
raise NotImplementedError(
f"Unsloth: video GRPO supports processors that return pixel_values_videos and "
f"video_grid_thw (Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Qwen3.5); "
f"{type(trainer.processing_class).__name__} does not."
)
second_per_grid_ts = processed.get("second_per_grid_ts", None)
if second_per_grid_ts is not None:
second_per_grid_ts = torch.as_tensor(second_per_grid_ts)
unique_counts = [num_videos[first[u]] for u in range(len(unique))]
video_index = torch.split(torch.arange(video_grid_thw.shape[0]), unique_counts)
row_ends = video_grid_thw.prod(dim = -1).cumsum(0)
row_starts = row_ends - video_grid_thw.prod(dim = -1)
videos, rows = [], []
for u in order:
for v in video_index[u].tolist():
videos.append(v)
rows.append(pixel_values_videos[int(row_starts[v]) : int(row_ends[v])])
device = trainer.accelerator.device
out = {
"pixel_values_videos": torch.cat(rows, dim = 0).to(device),
"video_grid_thw": video_grid_thw[videos].to(device),
"num_videos": num_videos,
}
if second_per_grid_ts is not None:
out["second_per_grid_ts"] = second_per_grid_ts[videos].to(device)
return out
def _unsloth_grpo_unsplit_vision(batch):
"""Undo _unsloth_grpo_split_vision_by_sample once this step's slice has been taken, so
the forward sees the layout the processor produced. TRL's own unsplit merges only
@ -1231,6 +1418,9 @@ def _unsloth_grpo_unsplit_vision(batch):
"image_sizes",
"image_position_ids",
"pixel_position_ids",
"pixel_values_videos",
"video_grid_thw",
"second_per_grid_ts",
):
value = batch.get(key, None)
if not isinstance(value, list) or len(value) == 0:
@ -1440,6 +1630,19 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
_unsloth_text_only = images is None
except NameError:
_unsloth_text_only = True
try:
_unsloth_video_images = images
except NameError:
_unsloth_video_images = None
try:
_unsloth_video_prompts = prompts
except NameError:
_unsloth_video_prompts = []
_unsloth_video_kwargs = _unsloth_grpo_video_inputs(
self, _unsloth_video_prompts, _unsloth_video_images
)
if _unsloth_video_kwargs is not None:
_unsloth_text_only = False
if _unsloth_text_only:
# Left pad prompt before calculation old and ref hidden states
left_pad_tokens_per_prompt = calculate_pad_tokens_in_prompt(prompt_completion_ids, logits_to_keep, self.processing_class.pad_token_id)
@ -1674,7 +1877,28 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
if _target_line in function:
function = function.replace(
_target_line,
_target_line + _metadata_extraction,
_target_line
+ _metadata_extraction
+ " prompts = _unsloth_grpo_clean_video_prompts(prompts, self)\n",
)
# TRL builds the logprob forward kwargs from images only; add the batch's videos.
_video_merge_anchor = "\n else:\n forward_kwargs = {}\n"
if _video_merge_anchor in function:
function = function.replace(
_video_merge_anchor,
_video_merge_anchor
+ " if _unsloth_video_kwargs is not None:\n"
+ " forward_kwargs = {**forward_kwargs, **_unsloth_video_kwargs}\n",
1,
)
elif "_unsloth_grpo_video_inputs(" in function:
function = function.replace(
" if _unsloth_video_kwargs is not None:\n _unsloth_text_only = False\n",
" if _unsloth_video_kwargs is not None:\n"
" raise NotImplementedError(\n"
" 'Unsloth: video GRPO needs TRL 0.24.0 or newer: pip install -U trl'\n"
" )\n",
)
_output_extras = """
@ -1763,7 +1987,11 @@ def grpo_trainer__generate_and_score_completions(function_name, function):
function = patched
_mm_alignment = """
if "mm_token_type_ids" in forward_kwargs or "image_grid_thw" in forward_kwargs:
if (
"mm_token_type_ids" in forward_kwargs
or "image_grid_thw" in forward_kwargs
or "video_grid_thw" in forward_kwargs
):
_mm_token_type_ids = _unsloth_fix_mm_token_type_ids(
self.processing_class,
prompt_completion_ids,
@ -2054,11 +2282,20 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function):
)
pixel_values = vision_inputs.get("pixel_values", None)
image_grid_thw = vision_inputs.get("image_grid_thw", None)
video_grid_thw = vision_inputs.get("video_grid_thw", None)
if pixel_values is None:
# Only a sentinel below (the chunks carry the tensors): video rows must take the
# vision path too, since left and sequence packing move M-RoPE video positions.
pixel_values = vision_inputs.get("pixel_values_videos", None)
num_images = vision_inputs.get("num_images", None)
# Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models.
token_type_ids = vision_inputs.get("token_type_ids", None)
mm_token_type_ids = vision_inputs.get("mm_token_type_ids", None)
if mm_token_type_ids is not None or image_grid_thw is not None:
if (
mm_token_type_ids is not None
or image_grid_thw is not None
or video_grid_thw is not None
):
mm_token_type_ids = _unsloth_fix_mm_token_type_ids(
self.processing_class, input_ids, mm_token_type_ids
)
@ -2940,6 +3177,12 @@ RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_accumulation
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_vision_inputs))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_split_vision_by_sample))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_unsplit_vision))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_split_videos_by_sample))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_prompt_videos))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_has_video_key))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_clean_video_prompts))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_prompt_key))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_video_inputs))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_grpo_image_cell))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_reject_grpo_image_list))
RL_PRE_ITEMS["grpo_trainer"].append(inspect.getsource(_unsloth_clear_stateful_mrope))
@ -3013,6 +3256,7 @@ def grpo_trainer_compute_loss(function_name, function):
# Transformers 5.x needs token_type_ids/mm_token_type_ids for some vision models.
token_type_ids = _vision_inputs.get("token_type_ids", None)
mm_token_type_ids = _vision_inputs.get("mm_token_type_ids", None)
video_grid_thw = _vision_inputs.get("video_grid_thw", None)
num_items_in_batch = inputs.get("num_items_in_batch", None)
sampling_per_token_logps = inputs.get("sampling_per_token_logps", None)
tool_mask = inputs.get("tool_mask", None)
@ -3022,7 +3266,11 @@ def grpo_trainer_compute_loss(function_name, function):
input_ids = torch.cat([prompt_ids, completion_ids], dim = 1)
bsz, qlen = input_ids.shape
attention_mask = torch.cat([prompt_mask, completion_mask], dim = 1)
if mm_token_type_ids is not None or image_grid_thw is not None:
if (
mm_token_type_ids is not None
or image_grid_thw is not None
or video_grid_thw is not None
):
mm_token_type_ids = _unsloth_fix_mm_token_type_ids(
self.processing_class,
input_ids,