from trl import GRPOTrainer, GRPOConfig
from datasets import load_dataset

train_dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")

def main():
    training_args = GRPOConfig(
        max_completion_length=8,
        per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
        num_generations=3,  # reduce the number of generations to reduce memory usage
        use_vllm=True,
        vllm_mode="colocate",
        report_to="none",
    )

    trainer = GRPOTrainer(
        model="Qwen/Qwen3-0.6B",
        reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
        args=training_args,
        train_dataset=train_dataset,
    )

    trainer.train()

if __name__ == "__main__":
    main()