Add training functionality with GRPO

app.py

@@ -1,10 +1,10 @@
 """
 HF Spaces app for VLIW kernel optimization via RL.
-Minimal working version.
 """
 import gradio as gr
+import threading
 
-# Check all imports at startup
+# Check imports at startup
 startup_log = []
 
 def check_import(name, import_fn):
@@ -37,65 +37,224 @@ try:
 except Exception as e:
     startup_log.append(f"✗ CUDA check: {e}")
 
+# Training state
+training_state = {
+    "is_training": False,
+    "should_stop": False,
+    "current_step": 0,
+    "best_cycles": float("inf"),
+    "log": [],
+}
+training_lock = threading.Lock()
+
 
 def get_status():
     return "\n".join(startup_log)
 
 
-def test_train(model_name, steps):
-    """Test training function."""
-    try:
-        import torch
-        from datasets import Dataset
-        from transformers import AutoTokenizer
-        from peft import LoraConfig
-        from trl import GRPOConfig, GRPOTrainer
+def reward_fn(completions, **kwargs):
+    """Simple reward function for testing."""
+    rewards = []
+    for completion in completions:
+        # Reward longer, code-like completions
+        text = completion[0]["content"] if isinstance(completion, list) else str(completion)
+        score = min(len(text) / 100.0, 1.0)  # Simple length-based reward
+        if "def " in text or "for " in text or "if " in text:
+            score += 0.5
+        rewards.append(score)
+    return rewards
+
+
+def run_training(model_name, num_steps, progress_callback):
+    """Run RL training."""
+    import torch
+    from datasets import Dataset
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import LoraConfig
+    from trl import GRPOConfig, GRPOTrainer
 
-        log = [f"Testing with {model_name}, {steps} steps"]
-        log.append(f"CUDA available: {torch.cuda.is_available()}")
+    with training_lock:
+        training_state["is_training"] = True
+        training_state["should_stop"] = False
+        training_state["current_step"] = 0
+        training_state["log"] = ["Starting training..."]
 
+    try:
+        progress_callback("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
-        log.append("Tokenizer loaded")
 
-        dataset = Dataset.from_dict({"prompt": ["Write hello world"] * 4})
-        log.append("Dataset created")
+        progress_callback("Loading model with 4-bit quantization...")
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+
+        progress_callback("Creating dataset...")
+        prompts = [
+            "Write optimized VLIW assembly for matrix multiplication",
+            "Generate SIMD code for vector addition",
+            "Create parallel code for reduction operation",
+            "Write efficient loop for memory copy",
+        ] * 4  # 16 prompts
+
+        dataset = Dataset.from_dict({"prompt": prompts})
+
+        progress_callback("Setting up LoRA config...")
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["q_proj", "v_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+        )
 
-        # Just check config creation
+        progress_callback("Creating trainer...")
         config = GRPOConfig(
-            output_dir="./test",
+            output_dir="./grpo_output",
             num_train_epochs=1,
+            max_steps=num_steps,
             per_device_train_batch_size=1,
+            gradient_accumulation_steps=4,
+            learning_rate=1e-5,
+            logging_steps=1,
             report_to="none",
+            remove_unused_columns=False,
+        )
+
+        trainer = GRPOTrainer(
+            model=model,
+            args=config,
+            train_dataset=dataset,
+            reward_funcs=reward_fn,
+            peft_config=lora_config,
+            processing_class=tokenizer,
         )
-        log.append("Config created")
 
-        return "\n".join(log)
+        progress_callback("Starting training loop...")
+        for step in range(num_steps):
+            with training_lock:
+                if training_state["should_stop"]:
+                    progress_callback("Training stopped by user")
+                    break
+                training_state["current_step"] = step + 1
+
+            # Run one step
+            try:
+                trainer.train()
+                progress_callback(f"Step {step + 1}/{num_steps} completed")
+            except Exception as e:
+                progress_callback(f"Step {step + 1} error: {str(e)[:100]}")
+                break
+
+        progress_callback("Training complete!")
 
     except Exception as e:
         import traceback
-        return f"Error: {e}\n\n{traceback.format_exc()}"
+        progress_callback(f"Error: {e}\n{traceback.format_exc()}")
+    finally:
+        with training_lock:
+            training_state["is_training"] = False
 
 
-with gr.Blocks(title="VLIW Test") as demo:
-    gr.Markdown("# VLIW Optimizer - Test Mode")
+def start_training(model_name, num_steps):
+    """Start training in background thread."""
+    with training_lock:
+        if training_state["is_training"]:
+            return "Training already in progress"
 
-    with gr.Row():
-        status_box = gr.Textbox(label="Startup Status", value=get_status(), lines=15)
+    log_messages = []
+
+    def progress_callback(msg):
+        log_messages.append(msg)
+        with training_lock:
+            training_state["log"] = log_messages.copy()
+
+    thread = threading.Thread(
+        target=run_training,
+        args=(model_name, int(num_steps), progress_callback),
+        daemon=False,
+    )
+    thread.start()
+    return "Training started! Check progress below."
+
+
+def stop_training():
+    """Request training stop."""
+    with training_lock:
+        if not training_state["is_training"]:
+            return "No training in progress"
+        training_state["should_stop"] = True
+    return "Stop requested..."
+
+
+def get_progress():
+    """Get current training progress."""
+    with training_lock:
+        if not training_state["log"]:
+            return "No training started yet"
+        return "\n".join(training_state["log"][-20:])  # Last 20 messages
+
+
+# Gradio UI
+with gr.Blocks(title="VLIW Optimizer") as demo:
+    gr.Markdown("# VLIW Kernel Optimizer - RL Training")
+    gr.Markdown("Train a model to generate optimized VLIW/SIMD kernels using reinforcement learning.")
 
     with gr.Row():
-        model = gr.Dropdown(
-            ["Qwen/Qwen2.5-Coder-1.5B-Instruct", "Qwen/Qwen2.5-Coder-3B-Instruct"],
-            value="Qwen/Qwen2.5-Coder-1.5B-Instruct",
-            label="Model"
-        )
-        steps = gr.Slider(1, 5, value=1, label="Steps")
+        with gr.Column(scale=1):
+            status_box = gr.Textbox(
+                label="System Status",
+                value=get_status(),
+                lines=10,
+                interactive=False,
+            )
+
+        with gr.Column(scale=2):
+            model_dropdown = gr.Dropdown(
+                choices=[
+                    "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+                    "Qwen/Qwen2.5-Coder-3B-Instruct",
+                ],
+                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",
+            )
+
+            with gr.Row():
+                start_btn = gr.Button("Start Training", variant="primary")
+                stop_btn = gr.Button("Stop Training", variant="stop")
+
+    output_box = gr.Textbox(
+        label="Training Progress",
+        lines=15,
+        interactive=False,
+    )
+
+    # Auto-refresh progress
+    refresh_btn = gr.Button("Refresh Progress")
 
-    test_btn = gr.Button("Test Training Setup", variant="primary")
-    output = gr.Textbox(label="Output", lines=15)
+    start_btn.click(start_training, [model_dropdown, steps_slider], [output_box])
+    stop_btn.click(stop_training, [], [output_box])
+    refresh_btn.click(get_progress, [], [output_box])
 
-    test_btn.click(test_train, [model, steps], [output])
+    # Auto-refresh every 5 seconds when training
+    demo.load(get_progress, [], [output_box], every=5)
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)