GRPO

[qgallouedec]

No description provided.

[qgallouedec]

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer

batch_size = 4
gradient_accumulation_steps = 2
output_dir = f"GRPO-bsz{batch_size}-grad_acc{gradient_accumulation_steps}-fixed"

training_args = GRPOConfig(
    output_dir=output_dir,
    per_device_train_batch_size=batch_size,
    gradient_accumulation_steps=gradient_accumulation_steps,
    logging_steps=2,
)

dummy_dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only")

def reward_len(prompts, completions):
    return [len(completion) for completion in completions]

trainer = GRPOTrainer(
    model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
    reward_funcs=reward_len,
    args=training_args,
    train_dataset=dummy_dataset["train"],
)

trainer.train()