GRPO training script for SEC model

grpo_training.py

@@ -0,0 +1,200 @@
+# /// script
+# dependencies = ["trl>=0.12.0", "peft>=0.18.0", "transformers>=4.45.0", "torch>=2.0.0", "trackio", "wandb", "accelerate>=0.21.0", "bitsandbytes"]
+# ///
+
+import os
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForCausalLM, 
+    BitsAndBytesConfig,
+    TrainingArguments
+)
+from trl import GRPOTrainer, GRPOConfig
+import trackio
+import wandb
+from huggingface_hub import HfApi
+
+def main():
+    # Initialize tracking
+    trackio.init(project="sec_grpo_training", run_name="llama32_1b_sec_grpo")
+    
+    print("🚀 Starting GRPO training for SEC model...")
+    
+    # Configuration
+    model_name = "ligaments-enterprise/llama3.2-1b-instruct-sec-finetuned"
+    dataset_name = "ligaments-enterprise/sec-data-preferences"
+    output_model = "ligaments-enterprise/llama3.2-1b-sec-grpo"
+    
+    # BitsAndBytesConfig for QLoRA (4-bit quantization)
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_use_double_quant=True
+    )
+    
+    # Load tokenizer and model
+    print(f"📦 Loading model: {model_name}")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load model with QLoRA
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True,
+        attn_implementation="flash_attention_2" if torch.cuda.get_device_capability()[0] >= 8 else None
+    )
+    
+    # Prepare model for k-bit training
+    model = prepare_model_for_kbit_training(model)
+    
+    # LoRA configuration for GRPO
+    lora_config = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        r=16,  # rank
+        lora_alpha=32,  # alpha scaling
+        lora_dropout=0.1,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        bias="none",
+    )
+    
+    # Apply LoRA to model
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+    
+    # Load preference dataset
+    print(f"📊 Loading preference dataset: {dataset_name}")
+    dataset = load_dataset(dataset_name, split="train")
+    print(f"✅ Loaded {len(dataset)} preference pairs")
+    
+    # Create train/eval split
+    train_test_split = dataset.train_test_split(test_size=0.1, seed=42)
+    train_dataset = train_test_split["train"]
+    eval_dataset = train_test_split["test"]
+    
+    print(f"📈 Training samples: {len(train_dataset)}")
+    print(f"📉 Evaluation samples: {len(eval_dataset)}")
+    
+    # GRPO Configuration
+    training_args = GRPOConfig(
+        output_dir="./grpo_sec_model",
+        
+        # Basic training settings
+        num_train_epochs=2,
+        per_device_train_batch_size=1,
+        per_device_eval_batch_size=1,
+        gradient_accumulation_steps=8,  # Effective batch size = 8
+        
+        # Learning rate and optimization
+        learning_rate=5e-6,  # Lower LR for RL fine-tuning
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.03,
+        
+        # Memory and efficiency
+        gradient_checkpointing=True,
+        dataloader_pin_memory=True,
+        bf16=True,
+        remove_unused_columns=False,
+        
+        # GRPO specific parameters
+        beta=0.1,  # KL penalty coefficient
+        grpo_score_clip=5.0,  # Clip scores to prevent instability
+        
+        # Evaluation and logging
+        eval_strategy="steps",
+        eval_steps=50,
+        logging_steps=10,
+        save_strategy="steps",
+        save_steps=100,
+        save_total_limit=3,
+        
+        # Tracking
+        report_to=["trackio"],
+        run_name="sec_grpo_training",
+        
+        # Hub integration
+        push_to_hub=True,
+        hub_model_id=output_model,
+        hub_strategy="every_save",
+        
+        # Length settings
+        max_length=512,
+        max_prompt_length=256,
+    )
+    
+    # Initialize GRPO Trainer
+    print("🎯 Initializing GRPO Trainer...")
+    trainer = GRPOTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        tokenizer=tokenizer,
+        peft_config=lora_config,
+    )
+    
+    # Log initial metrics
+    trackio.log({
+        "model_name": model_name,
+        "dataset_name": dataset_name,
+        "output_model": output_model,
+        "train_samples": len(train_dataset),
+        "eval_samples": len(eval_dataset),
+        "lora_rank": lora_config.r,
+        "lora_alpha": lora_config.lora_alpha,
+        "beta": training_args.beta,
+        "learning_rate": training_args.learning_rate,
+    })
+    
+    # Start training
+    print("🚀 Starting GRPO training...")
+    try:
+        trainer.train()
+        
+        # Log final metrics
+        trainer_state = trainer.state
+        trackio.log({
+            "final_train_loss": trainer_state.log_history[-1].get("train_loss", 0),
+            "final_eval_loss": trainer_state.log_history[-1].get("eval_loss", 0),
+            "training_completed": True
+        })
+        
+        # Save final model
+        print("💾 Saving final model...")
+        trainer.save_model()
+        
+        # Push to hub
+        print("📤 Pushing to Hub...")
+        trainer.push_to_hub(commit_message="GRPO training completed")
+        
+        print(f"✅ GRPO training completed successfully!")
+        print(f"📦 Model saved to: {output_model}")
+        
+        # Create evaluation summary
+        eval_summary = {
+            "total_steps": trainer_state.global_step,
+            "total_epochs": trainer_state.epoch,
+            "final_train_loss": trainer_state.log_history[-1].get("train_loss", "N/A"),
+            "final_eval_loss": trainer_state.log_history[-1].get("eval_loss", "N/A"),
+            "model_parameters": sum(p.numel() for p in model.parameters() if p.requires_grad),
+        }
+        
+        print("📊 Training Summary:")
+        for key, value in eval_summary.items():
+            print(f"   {key}: {value}")
+            
+        trackio.log(eval_summary)
+        
+    except Exception as e:
+        print(f"❌ Training failed: {e}")
+        trackio.log({"error": str(e), "training_completed": False})
+        raise e
+
+if __name__ == "__main__":
+    main()