from datasets import load_dataset
from trl.trainer.grpo_trainer import GRPOTrainer
from trl.trainer.grpo_config import GRPOConfig


dataset = load_dataset("trl-lib/tldr", split="train[:200]")


def reward_func(completions, **kwargs):
    return [len(set(c)) for c in completions]

args = GRPOConfig(
    output_dir="mask_truncated_ture_with_max_completion_length",
    use_vllm=False,
    bf16=True,
    gradient_checkpointing=True,
    logging_steps=50,
    mask_truncated_completions=True, 
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=1,
    num_generations=4,  
    max_completion_length=50,
)


trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B",
    args=args,
    reward_funcs=reward_func,
    train_dataset=dataset,
)


trainer.train() 

accelerate launch 3248.py