Add GRPO training with proper state management

app.py

@@ -3,6 +3,7 @@ HF Spaces app for VLIW kernel optimization via RL.
 """
 import gradio as gr
 import threading
+import time
 
 # Check imports at startup
 startup_log = []
@@ -37,33 +38,64 @@ try:
 except Exception as e:
     startup_log.append(f"✗ CUDA check: {e}")
 
-# Global state
-training_log = []
-is_training = False
+# Training state
+training_state = {
+    "is_training": False,
+    "should_stop": False,
+    "log": [],
+}
+state_lock = threading.Lock()
 
 
 def get_status():
     return "\n".join(startup_log)
 
 
-def test_model_load(model_name):
-    """Test loading the model."""
-    global training_log
-    training_log = []
+def simple_reward_fn(completions, **kwargs):
+    """Simple reward: prefer longer, code-like outputs."""
+    rewards = []
+    for c in completions:
+        text = c[0]["content"] if isinstance(c, list) else str(c)
+        score = min(len(text) / 200.0, 1.0)
+        if any(kw in text for kw in ["def ", "for ", "if ", "while ", "return "]):
+            score += 0.3
+        rewards.append(score)
+    return rewards
 
-    try:
-        import torch
-        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-        training_log.append(f"Testing model: {model_name}")
-        training_log.append("Loading tokenizer...")
+def run_training(model_name, num_steps):
+    """Run RL training with GRPO."""
+    import torch
+    from datasets import Dataset
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import LoraConfig
+    from trl import GRPOConfig, GRPOTrainer
+    from transformers import TrainerCallback
+
+    log = []
+
+    def add_log(msg):
+        log.append(f"[{time.strftime('%H:%M:%S')}] {msg}")
+        with state_lock:
+            training_state["log"] = log.copy()
+
+    with state_lock:
+        training_state["is_training"] = True
+        training_state["should_stop"] = False
+        training_state["log"] = []
+
+    try:
+        add_log(f"Starting training: {model_name}, {num_steps} steps")
 
+        # Load tokenizer
+        add_log("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
-        training_log.append("✓ Tokenizer loaded")
+        add_log("✓ Tokenizer ready")
 
-        training_log.append("Loading model with 4-bit quantization...")
+        # Load model with 4-bit quantization
+        add_log("Loading model (4-bit quantization)...")
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -75,47 +107,147 @@ def test_model_load(model_name):
             device_map="auto",
             trust_remote_code=True,
         )
-        training_log.append("✓ Model loaded")
+        add_log(f"✓ Model loaded on {next(model.parameters()).device}")
 
-        # Quick test
-        training_log.append("Testing generation...")
-        inputs = tokenizer("def hello():", return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, max_new_tokens=20)
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        training_log.append(f"✓ Generation test: {result[:50]}...")
+        # Create dataset
+        add_log("Creating training dataset...")
+        prompts = [
+            "Write optimized VLIW assembly code for matrix multiplication using SIMD instructions",
+            "Generate efficient parallel code for vector dot product",
+            "Create VLIW code for memory-bound reduction operation",
+            "Write pipelined code for element-wise array operations",
+        ] * 8  # 32 prompts total
+        dataset = Dataset.from_dict({"prompt": prompts})
+        add_log(f"✓ Dataset: {len(prompts)} prompts")
+
+        # LoRA config
+        add_log("Setting up LoRA...")
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
 
-        training_log.append("\n✓ All tests passed!")
+        # Stop callback
+        class StopCallback(TrainerCallback):
+            def on_step_end(self, args, state, control, **kwargs):
+                with state_lock:
+                    if training_state["should_stop"]:
+                        control.should_training_stop = True
+                return control
 
-        # Cleanup
-        del model
-        torch.cuda.empty_cache()
+        # GRPO config
+        add_log("Creating GRPO trainer...")
+        config = GRPOConfig(
+            output_dir="./grpo_output",
+            num_train_epochs=1,
+            max_steps=num_steps,
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=2,
+            learning_rate=5e-6,
+            logging_steps=1,
+            save_steps=999999,  # Don't save checkpoints
+            report_to="none",
+            remove_unused_columns=False,
+            max_completion_length=128,
+            num_generations=4,
+        )
+
+        trainer = GRPOTrainer(
+            model=model,
+            args=config,
+            train_dataset=dataset,
+            reward_funcs=simple_reward_fn,
+            peft_config=lora_config,
+            processing_class=tokenizer,
+            callbacks=[StopCallback()],
+        )
+        add_log("✓ Trainer ready")
+
+        # Train
+        add_log("Starting training loop...")
+        train_result = trainer.train()
+
+        metrics = train_result.metrics
+        add_log(f"✓ Training complete!")
+        add_log(f"  Steps: {metrics.get('train_steps', 'N/A')}")
+        add_log(f"  Loss: {metrics.get('train_loss', 'N/A'):.4f}" if 'train_loss' in metrics else "  Loss: N/A")
+
+        # Test generation
+        add_log("Testing trained model...")
+        test_prompt = "Write efficient VLIW code for:"
+        inputs = tokenizer(test_prompt, return_tensors="pt").to(model.device)
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_new_tokens=50, do_sample=True, temperature=0.7)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        add_log(f"Sample output: {result[:100]}...")
+
+        add_log("\n✓ All done!")
 
     except Exception as e:
         import traceback
-        training_log.append(f"✗ Error: {e}")
-        training_log.append(traceback.format_exc())
+        add_log(f"✗ Error: {e}")
+        add_log(traceback.format_exc()[:500])
+    finally:
+        with state_lock:
+            training_state["is_training"] = False
+        # Cleanup
+        try:
+            del model
+            torch.cuda.empty_cache()
+        except:
+            pass
+
+    return "\n".join(log)
 
-    return "\n".join(training_log)
 
+def start_training(model_name, num_steps):
+    """Start training (blocking for simplicity)."""
+    with state_lock:
+        if training_state["is_training"]:
+            return "Training already in progress. Please wait."
 
-def get_log():
-    """Return current log."""
-    if not training_log:
-        return "No operations run yet. Click 'Test Model Loading' to start."
-    return "\n".join(training_log)
+    return run_training(model_name, int(num_steps))
+
+
+def stop_training():
+    """Request stop."""
+    with state_lock:
+        if not training_state["is_training"]:
+            return "No training in progress"
+        training_state["should_stop"] = True
+    return "Stop requested. Training will stop after current step."
+
+
+def get_progress():
+    """Get current log."""
+    with state_lock:
+        if not training_state["log"]:
+            return "No training started yet"
+        return "\n".join(training_state["log"])
 
 
 # Gradio UI
 with gr.Blocks(title="VLIW Optimizer") as demo:
-    gr.Markdown("# VLIW Kernel Optimizer - Test Mode")
-    gr.Markdown("Testing model loading and generation before full training.")
+    gr.Markdown("# VLIW Kernel Optimizer - RL Training")
+    gr.Markdown("""
+    Train a language model with reinforcement learning to generate optimized VLIW/SIMD code.
+
+    **Instructions:**
+    1. Select a model (1.5B is faster, 3B may produce better results)
+    2. Set training steps (10-50 recommended for testing)
+    3. Click 'Start Training' and wait for completion
+    """)
 
     with gr.Row():
         with gr.Column(scale=1):
             status_box = gr.Textbox(
                 label="System Status",
                 value=get_status(),
-                lines=10,
+                lines=9,
                 interactive=False,
             )
 
@@ -128,17 +260,31 @@ with gr.Blocks(title="VLIW Optimizer") as demo:
                 value="Qwen/Qwen2.5-Coder-1.5B-Instruct",
                 label="Model",
             )
+            steps_slider = gr.Slider(
+                minimum=1,
+                maximum=100,
+                value=10,
+                step=1,
+                label="Training Steps",
+            )
 
-            test_btn = gr.Button("Test Model Loading", variant="primary")
+            with gr.Row():
+                start_btn = gr.Button("Start Training", variant="primary")
+                stop_btn = gr.Button("Stop", variant="stop")
 
     output_box = gr.Textbox(
-        label="Output",
-        lines=15,
+        label="Training Log",
+        lines=20,
         interactive=False,
-        value="Click 'Test Model Loading' to verify the setup.",
+        value="Click 'Start Training' to begin.",
     )
 
-    test_btn.click(test_model_load, [model_dropdown], [output_box])
+    start_btn.click(
+        start_training,
+        [model_dropdown, steps_slider],
+        [output_box],
+    )
+    stop_btn.click(stop_training, [], [output_box])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)