trl-fine-tuning

1"""2Basic GRPO Training Template3=============================4 5A minimal, production-ready template for GRPO training with TRL.6Adapt this for your specific task by modifying:71. Dataset loading (get_dataset function)82. Reward functions (reward_*_func)93. System prompt (SYSTEM_PROMPT)104. Hyperparameters (GRPOConfig)11"""12 13import torch14import re15from datasets import load_dataset16from transformers import AutoModelForCausalLM, AutoTokenizer17from peft import LoraConfig18from trl import GRPOTrainer, GRPOConfig19 20# ==================== CONFIGURATION ====================21 22MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"23OUTPUT_DIR = "outputs/grpo-model"24MAX_PROMPT_LENGTH = 25625MAX_COMPLETION_LENGTH = 51226 27SYSTEM_PROMPT = """28Respond in the following format:29<reasoning>30[Your step-by-step thinking]31</reasoning>32<answer>33[Final answer]34</answer>35"""36 37# ==================== DATASET ====================38 39def get_dataset(split="train"):40    """41    Load and prepare your dataset.42 43    Returns: Dataset with columns:44    - 'prompt': List[Dict] with role/content45    - 'answer': str (ground truth, optional)46    """47    # Example: GSM8K math dataset48    data = load_dataset('openai/gsm8k', 'main')[split]49 50    def process_example(x):51        # Extract ground truth answer52        answer = x['answer'].split('####')[1].strip() if '####' in x['answer'] else None53 54        return {55            'prompt': [56                {'role': 'system', 'content': SYSTEM_PROMPT},57                {'role': 'user', 'content': x['question']}58            ],59            'answer': answer60        }61 62    return data.map(process_example)63 64# ==================== HELPER FUNCTIONS ====================65 66def extract_xml_tag(text: str, tag: str) -> str:67    """Extract content between XML tags."""68    pattern = f'<{tag}>(.*?)</{tag}>'69    match = re.search(pattern, text, re.DOTALL)70    return match.group(1).strip() if match else ""71 72def extract_answer(text: str) -> str:73    """Extract the final answer from structured output."""74    return extract_xml_tag(text, 'answer')75 76# ==================== REWARD FUNCTIONS ====================77 78def correctness_reward_func(prompts, completions, answer, **kwargs):79    """80    Reward correct answers.81    Weight: 2.0 (highest priority)82    """83    responses = [comp[0]['content'] for comp in completions]84    extracted = [extract_answer(r) for r in responses]85    return [2.0 if ans == gt else 0.0 for ans, gt in zip(extracted, answer)]86 87def format_reward_func(completions, **kwargs):88    """89    Reward proper XML format.90    Weight: 0.591    """92    pattern = r'<reasoning>.*?</reasoning>\s*<answer>.*?</answer>'93    responses = [comp[0]['content'] for comp in completions]94    return [0.5 if re.search(pattern, r, re.DOTALL) else 0.0 for r in responses]95 96def incremental_format_reward_func(completions, **kwargs):97    """98    Incremental reward for partial format compliance.99    Weight: up to 0.5100    """101    responses = [comp[0]['content'] for comp in completions]102    rewards = []103 104    for r in responses:105        score = 0.0106        if '<reasoning>' in r:107            score += 0.125108        if '</reasoning>' in r:109            score += 0.125110        if '<answer>' in r:111            score += 0.125112        if '</answer>' in r:113            score += 0.125114 115        # Penalize extra content after closing tag116        if '</answer>' in r:117            extra = r.split('</answer>')[-1].strip()118            score -= len(extra) * 0.001119 120        rewards.append(score)121 122    return rewards123 124# ==================== MODEL SETUP ====================125 126def setup_model_and_tokenizer():127    """Load model and tokenizer with optimizations."""128    model = AutoModelForCausalLM.from_pretrained(129        MODEL_NAME,130        torch_dtype=torch.bfloat16,131        attn_implementation="flash_attention_2",132        device_map="auto"133    )134 135    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)136    tokenizer.pad_token = tokenizer.eos_token137 138    return model, tokenizer139 140def get_peft_config():141    """LoRA configuration for parameter-efficient training."""142    return LoraConfig(143        r=16,144        lora_alpha=32,145        target_modules=[146            "q_proj", "k_proj", "v_proj", "o_proj",147            "gate_proj", "up_proj", "down_proj"148        ],149        task_type="CAUSAL_LM",150        lora_dropout=0.05,151    )152 153# ==================== TRAINING ====================154 155def main():156    """Main training function."""157 158    # Load data159    print("Loading dataset...")160    dataset = get_dataset()161    print(f"Dataset size: {len(dataset)}")162 163    # Setup model164    print("Loading model...")165    model, tokenizer = setup_model_and_tokenizer()166 167    # Training configuration168    training_args = GRPOConfig(169        output_dir=OUTPUT_DIR,170        run_name="grpo-training",171 172        # Learning rate173        learning_rate=5e-6,174        adam_beta1=0.9,175        adam_beta2=0.99,176        weight_decay=0.1,177        warmup_ratio=0.1,178        lr_scheduler_type='cosine',179 180        # Batch settings181        per_device_train_batch_size=1,182        gradient_accumulation_steps=4,183 184        # GRPO specific185        num_generations=8,186        max_prompt_length=MAX_PROMPT_LENGTH,187        max_completion_length=MAX_COMPLETION_LENGTH,188 189        # Training duration190        num_train_epochs=1,191 192        # Optimization193        bf16=True,194        optim="adamw_8bit",195        max_grad_norm=0.1,196 197        # Logging198        logging_steps=1,199        save_steps=100,200        report_to="wandb",  # Change to "none" to disable logging201    )202 203    # Initialize trainer204    trainer = GRPOTrainer(205        model=model,206        processing_class=tokenizer,207        reward_funcs=[208            incremental_format_reward_func,209            format_reward_func,210            correctness_reward_func,211        ],212        args=training_args,213        train_dataset=dataset,214        peft_config=get_peft_config(),215    )216 217    # Train218    print("Starting training...")219    trainer.train()220 221    # Save final model222    print(f"Saving model to {OUTPUT_DIR}/final")223    trainer.save_model(f"{OUTPUT_DIR}/final")224 225    print("Training complete!")226 227if __name__ == "__main__":228    main()229