The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Hi @ahatamiz Ali jan, Thanks for your contribution; I've tested your PR;

For future reference I tested with;

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

# Set required environment variables for VLLM
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "12355"

dataset = load_dataset("trl-lib/tldr", split="train")


def reward_len(completions, **kwargs):
    return [-abs(20 - len(completion)) for completion in completions]


training_args = GRPOConfig(
    output_dir="pr3634",
    use_vllm=True,
    vllm_mode="colocate",
    gradient_checkpointing=True,
    num_generations=4,
    per_device_train_batch_size=4,
    vllm_gpu_memory_utilization=0.10,
    vllm_tensor_parallel_size=1,
    max_prompt_length=512,
    max_completion_length=1024,
    max_steps=2
)

trainer = GRPOTrainer(
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=reward_len,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()

then ran with;

accelerate launch --config_file examples/accelerate_configs/single_gpu.yaml PR_3634.py

Thank you @shirinyamani for the review and also this pointer. Sure, I will keep this in mind.

Kind Regards,
Ali