Upload grpo_training.py with huggingface_hub

grpo_training.py

@@ -1,200 +1,97 @@
 # /// 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"]
+# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "torch", "transformers"]
 # ///
 
-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
+from peft import LoraConfig
+from trl.trainer.grpo_trainer import GRPOTrainer, GRPOConfig
+from transformers import AutoTokenizer
 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"],
+import torch
+
+# Load your fine-tuned model and preference dataset
+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"
+
+# Load dataset
+dataset = load_dataset(dataset_name, split="train")
+print(f"Loaded {len(dataset)} preference pairs from {dataset_name}")
+
+# Create train/eval split for monitoring
+dataset_split = dataset.train_test_split(test_size=0.1, seed=42)
+train_dataset = dataset_split["train"]
+eval_dataset = dataset_split["test"]
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Configure GRPO training
+config = GRPOConfig(
+    output_dir=output_model,
+    num_train_epochs=3,
+    per_device_train_batch_size=1,
+    per_device_eval_batch_size=1,
+    gradient_accumulation_steps=8,  # Effective batch size = 8
+    learning_rate=1e-6,
+    max_length=1024,
+
+    # Evaluation and logging
+    eval_strategy="steps",
+    eval_steps=50,
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=100,
+
+    # Hub integration
+    push_to_hub=True,
+    hub_model_id=output_model,
+    hub_strategy="every_save",
+
+    # Optimization
+    gradient_checkpointing=True,
+    bf16=True if torch.cuda.is_bf16_supported() else False,
+    fp16=False if torch.cuda.is_bf16_supported() else True,
+
+    # Trackio monitoring
+    report_to="trackio",
+    run_name="llama3.2-1b-sec-grpo-training",
+    project="ligaments-sec-alignment",
+
+    # GRPO specific parameters
+    kl_penalty="kl",  # KL penalty for policy regularization
+    temperature=0.7,
+)
+
+# Initialize GRPO trainer
+trainer = GRPOTrainer(
+    model=model_name,
+    tokenizer=tokenizer,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    peft_config=LoraConfig(
+        r=16,
+        lora_alpha=32,
+        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+        lora_dropout=0.05,
         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()
+        task_type="CAUSAL_LM"
+    ),
+    args=config,
+)
+
+print("Starting GRPO training...")
+print(f"Training on {len(train_dataset)} preference pairs")
+print(f"Evaluating on {len(eval_dataset)} preference pairs")
+print(f"Output model will be saved to: {output_model}")
+
+# Train the model
+trainer.train()
+
+# Push final model to Hub
+trainer.push_to_hub()
+
+print("GRPO training completed successfully!")
+print(f"Final model available at: https://huggingface.co/{output_model}")