def build_config(run_name="codeforge-grpo", output_dir="checkpoints"):
    from trl import GRPOConfig
    return GRPOConfig(
        learning_rate=5e-6,
        adam_beta1=0.9, adam_beta2=0.99, weight_decay=0.01,
        warmup_ratio=0.1,
        lr_scheduler_type="cosine",
        optim="adamw_torch_fused",

        per_device_train_batch_size=1,
        gradient_accumulation_steps=2,
        num_generations=8,                  # G=8 for parallel A100s
        max_prompt_length=256,              
        max_completion_length=512,          # Keep short to speed up MBPP training

        beta=0.04,                          # KL penalty
        temperature=0.9, top_p=0.95,

        bf16=True,                          # Target A100s
        gradient_checkpointing=True,

        logging_steps=1, save_steps=100,
        save_total_limit=3,
        output_dir=output_dir,
        report_to="wandb",
        run_name=run_name,
    )