🧺 [2/N] Refactor _generate in GRPO/RLOO: Use prompt_ids from generation

[qgallouedec]

This PR belongs to a sequence of PR that aims to refactor the generation part of GRPO/RLOO to allow for easier customization and ultimately tool calling

Previous:

• 🟩 Drop image_split_sizes in favour of image_grid_thw #4111
• 📽 Multi image support for GRPO/RLOO #4113
• 😷 Refactor GRPO/RLOO to isolate _generate #4114
• 🧺 [1/N] Refactor _generate in GRPO/RLOO: list of ints instead of tensors #4146

Next:

• 🧺 [3/N] Refactor _generate in GRPO/RLOO: Rely on generator for prompt truncation #4153
• 🧺 [4/N] Refactor _generate in GRPO/RLOO: Move forward_kwargs outside generation method #4154
• 🧺 [5/N] Refactor _generate in GRPO/RLOO: Insert images in the prompt #4155

Instead of getting prompt_ids for all generation methods like this

prompt_inputs = super()._prepare_inputs(prompt_inputs)
prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
forward_kwargs = {k: v for k, v in prompt_inputs.items() if k not in ["input_ids", "attention_mask"]}
prompt_ids = [p[m].tolist() for p, m in zip(prompt_ids, prompt_mask.bool())]

we rely on each method to provide the prompt_ids.

Benchmark

I ran a quick benchmark with two configs, results are here:

I can't remember which run correspond to before, and which one correspond to after 😅

from datasets import load_dataset
from trl import GRPOTrainer, GRPOConfig
import os

os.environ["TRACKIO_PROJECT"] = "4152"
os.environ["TRACKIO_SPACE_ID"] = "qgallouedec/trackio"

# Dummy reward function: count the number of unique characters in the completions
def reward_num_unique_chars(completions, **kwargs):
    return [len(set(c)) for c in completions]

dataset = load_dataset("trl-lib/tldr", split="train[:500]")
dataset.select_columns(["prompt"])

training_args = GRPOConfig(
    output_dir="tmp",
    per_device_train_batch_size=4,  # reduce the batch size to reduce memory usage
    gradient_accumulation_steps=8,
    num_generations=4,  # reduce the number of generations to reduce memory usage
    max_completion_length=256,  # reduce the completion length to reduce memory usage
    steps_per_generation=4,
    logging_steps=1,
    num_train_epochs=1,
    vllm_mode="colocate",
    use_vllm=True,
    report_to="trackio"
)
trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B-Base",
    reward_funcs=[reward_num_unique_chars],
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

https://huggingface.co/spaces/qgallouedec/trackio?project=4152

[lewtun]

Thanks for the clean refactor!

Overall LGTM with a question about whether the difference across benchmark runs is within the variance of a repeated run (i.e. if you run one of these again on main does it also fluctuate from the original run on main)?