Upload train_grpo_code.py with huggingface_hub

train_grpo_code.py

@@ -26,13 +26,12 @@ import os
 import re
 import subprocess
 import tempfile
-import traceback
 from typing import Any
 
 import torch
 import trackio
 from datasets import load_dataset
-from peft import PeftModel
+from peft import LoraConfig, PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from trl import GRPOTrainer, GRPOConfig
 
@@ -44,7 +43,7 @@ print("=" * 60)
 BASE_MODEL = "Qwen/Qwen3-0.6B"
 SFT_ADAPTER = "chaddy81/qwen3-0.6b-multicode-sft"
 OUTPUT_REPO = "chaddy81/qwen3-0.6b-multicode-grpo"
-MAX_EXAMPLES = 2000  # Limit for reasonable training time
+MAX_EXAMPLES = 1000  # Reduced for faster training
 
 print(f"\n📦 Configuration:")
 print(f"   Base model: {BASE_MODEL}")
@@ -68,7 +67,7 @@ def extract_python_code(text: str) -> str:
     for pattern in patterns:
         matches = re.findall(pattern, text, re.DOTALL)
         if matches:
-            return matches[-1].strip()  # Return last code block
+            return matches[-1].strip()
 
     # If no markdown blocks, try to find code after common markers
     markers = ["Solution:", "Answer:", "Code:"]
@@ -82,18 +81,8 @@ def extract_python_code(text: str) -> str:
     return text.strip()
 
 
-def run_python_code(code: str, stdin_input: str, timeout: float = 5.0) -> tuple[bool, str]:
-    """
-    Execute Python code with given input and return (success, output).
-
-    Args:
-        code: Python source code to execute
-        stdin_input: Input to pass via stdin
-        timeout: Maximum execution time in seconds
-
-    Returns:
-        Tuple of (success: bool, output: str)
-    """
+def run_python_code(code: str, stdin_input: str, timeout: float = 3.0) -> tuple[bool, str]:
+    """Execute Python code with given input and return (success, output)."""
     try:
         with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
             f.write(code)
@@ -120,7 +109,7 @@ def run_python_code(code: str, stdin_input: str, timeout: float = 5.0) -> tuple[
 
 
 def normalize_output(output: str) -> str:
-    """Normalize output for comparison (strip whitespace, normalize newlines)."""
+    """Normalize output for comparison."""
     return '\n'.join(line.strip() for line in output.strip().split('\n'))
 
 
@@ -134,29 +123,27 @@ def code_execution_reward(
     Reward function that executes generated code against test cases.
 
     Returns:
-        - 1.0 if code passes all available tests
-        - 0.5 if code passes some tests
-        - 0.0 if code fails all tests or has errors
+        - 1.0 if code passes all tests
+        - Partial credit for some tests
+        - 0.0 if fails all tests
     """
     rewards = []
 
     for completion, tests, exs in zip(completions, official_tests, examples):
-        # Extract code from completion
         code = extract_python_code(completion)
 
         if not code or len(code) < 10:
             rewards.append(0.0)
             continue
 
-        # Combine official tests and examples
+        # Combine tests (limit to avoid long execution)
         all_tests = []
         if tests:
-            all_tests.extend(tests[:3])  # Limit to first 3 official tests
+            all_tests.extend(tests[:2])
         if exs:
-            all_tests.extend(exs[:2])    # Add up to 2 examples
+            all_tests.extend(exs[:2])
 
         if not all_tests:
-            # No tests available, give neutral reward
             rewards.append(0.0)
             continue
 
@@ -168,10 +155,9 @@ def code_execution_reward(
             test_input = test.get('input', '')
             expected_output = test.get('output', '')
 
-            success, actual_output = run_python_code(code, test_input, timeout=3.0)
+            success, actual_output = run_python_code(code, test_input, timeout=2.0)
 
             if success:
-                # Compare outputs (normalized)
                 if normalize_output(actual_output) == normalize_output(expected_output):
                     passed += 1
 
@@ -200,18 +186,14 @@ dataset = load_dataset(
 )
 print(f"   Total examples: {len(dataset)}")
 
-# Filter for Python problems only
-print("   Filtering for Python problems...")
-dataset = dataset.filter(lambda x: x.get('language') == 'python')
-print(f"   Python problems: {len(dataset)}")
-
-# Filter for problems with test cases
-print("   Filtering for problems with tests...")
+# Filter for Python problems with tests
+print("   Filtering for Python problems with tests...")
 dataset = dataset.filter(
-    lambda x: (x.get('official_tests') and len(x['official_tests']) > 0) or
-              (x.get('examples') and len(x['examples']) > 0)
+    lambda x: x.get('language') == 'python' and
+              ((x.get('official_tests') and len(x['official_tests']) > 0) or
+               (x.get('examples') and len(x['examples']) > 0))
 )
-print(f"   Problems with tests: {len(dataset)}")
+print(f"   Filtered: {len(dataset)}")
 
 # Limit dataset size
 if len(dataset) > MAX_EXAMPLES:
@@ -222,44 +204,63 @@ print(f"\n✅ Final dataset: {len(dataset)} examples")
 
 
 # ============================================================================
-# Model Loading
+# Model Loading - Merge SFT then save for GRPO
 # ============================================================================
 
-print("\n🔧 Loading model...")
+print("\n🔧 Loading and preparing model...")
 
-# Load base model
+# Step 1: Load base model and SFT adapter
 print("   Loading base model...")
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
+    device_map="cpu",
     trust_remote_code=True,
 )
 
-# Load SFT adapter
 print("   Loading SFT adapter...")
 model = PeftModel.from_pretrained(base_model, SFT_ADAPTER)
 
-# Merge for GRPO (GRPO works better with merged models)
-print("   Merging adapter...")
+print("   Merging SFT adapter into base model...")
 model = model.merge_and_unload()
 
+# Step 2: Save merged model temporarily
+merged_path = "/tmp/merged_sft_model"
+print(f"   Saving merged model to {merged_path}...")
+model.save_pretrained(merged_path, safe_serialization=True)
+
 # Load tokenizer
 print("   Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(SFT_ADAPTER, trust_remote_code=True)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 tokenizer.padding_side = "left"
+tokenizer.save_pretrained(merged_path)
+
+# Free memory
+del model
+del base_model
+torch.cuda.empty_cache()
 
-print("   ✅ Model ready")
+print("   ✅ Merged model saved")
 
 
 # ============================================================================
-# GRPO Training
+# GRPO Training with fresh LoRA
 # ============================================================================
 
 print("\n⚙️  Configuring GRPO trainer...")
 
+# LoRA config for GRPO (smaller rank for efficiency)
+peft_config = LoraConfig(
+    r=8,
+    lora_alpha=16,
+    lora_dropout=0.05,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+    bias="none",
+    task_type="CAUSAL_LM",
+)
+
 config = GRPOConfig(
     # Output & Hub
     output_dir="qwen3-grpo-code",
@@ -269,14 +270,14 @@ config = GRPOConfig(
     hub_private_repo=False,
 
     # GRPO parameters
-    num_generations=4,           # Generate 4 completions per prompt
-    max_completion_length=512,   # Max tokens for generated code
+    num_generations=4,
+    max_completion_length=256,  # Shorter for faster training
 
     # Training parameters
     num_train_epochs=1,
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps=4,
-    learning_rate=1e-6,          # Very low LR for GRPO
+    per_device_train_batch_size=1,  # Small batch for memory
+    gradient_accumulation_steps=8,
+    learning_rate=5e-7,
 
     # Optimization
     warmup_ratio=0.1,
@@ -287,26 +288,27 @@ config = GRPOConfig(
     # Logging & checkpoints
     logging_steps=10,
     save_strategy="steps",
-    save_steps=100,
+    save_steps=50,
     save_total_limit=2,
 
     # Monitoring
     report_to="trackio",
     project="qwen3-grpo-code",
-    run_name="grpo-codeforces-v1",
+    run_name="grpo-codeforces-v2",
 )
 
-print("   Initializing trainer...")
+print("   Initializing trainer with merged SFT model + new LoRA...")
 trainer = GRPOTrainer(
-    model=model,
+    model=merged_path,  # Pass path - trainer loads with proper gradients
     processing_class=tokenizer,
     reward_funcs=code_execution_reward,
     train_dataset=dataset,
     args=config,
+    peft_config=peft_config,  # New LoRA for GRPO
 )
 
 print("\n🚀 Starting GRPO training...")
-print("   This will generate code, execute it, and learn from results.")
+print("   Training will generate code, execute it, and learn from results.")
 print("=" * 60)
 
 trainer.train()