Full training app with verified imports

app.py

@@ -1,101 +1,371 @@
 """
 HF Spaces app for VLIW kernel optimization via RL.
-Minimal version for debugging.
+Deploy to HF Spaces Pro (A10G GPU).
 """
 import os
 import sys
+import re
+import threading
+import time
+import random
+from datetime import datetime
+
 import gradio as gr
 
-# Check imports on startup
-startup_log = []
+# Thread lock for safe state access
+training_state_lock = threading.Lock()
+
+# Add simulator path
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+PERF_TAKEHOME_PATH = os.path.join(SCRIPT_DIR, "original_performance_takehome")
+if os.path.exists(PERF_TAKEHOME_PATH):
+    sys.path.insert(0, PERF_TAKEHOME_PATH)
+
+# Constants
+BASELINE_CYCLES = 147734
+TARGET_CYCLES = 1363
+SCORE_SCALE = 3000.0
+
+# Training state
+training_state = {
+    "running": False,
+    "step": 0,
+    "total_steps": 0,
+    "best_cycles": BASELINE_CYCLES,
+    "best_code": None,
+    "log": [],
+    "start_time": None,
+    "results": [],
+}
+
+SYSTEM_PROMPT = '''Write optimized VLIW/SIMD kernel. OUTPUT ONLY ONE ```python CODE BLOCK.
+
+ARCHITECTURE: 12 ALU + 6 VALU (VLEN=8) + 2 load + 2 store + 1 flow slots per cycle.
+
+API:
+- alloc_scratch(name, length) -> addr
+- add(engine, slot): engine in {alu, valu, load, store, flow}
+- valu ops work on 8 elements at once
+- build(slots, vliw=True): pack into VLIW bundle
+
+ALGORITHM: 16 rounds x 256 items, hash tree traversal.
+
+OPTIMIZATION: Use vload/vstore (8 elements), pack 6 VALU ops/cycle, unroll loops.
+'''
+
 
-def check_import(name, import_fn):
+def extract_code_block(text: str) -> str:
+    """Extract python code from markdown."""
+    pattern = r"```python\s*(.*?)```"
+    matches = re.findall(pattern, text, re.DOTALL)
+    if matches:
+        return matches[-1].strip()
+    pattern = r"```\s*(.*?)```"
+    matches = re.findall(pattern, text, re.DOTALL)
+    if matches:
+        return matches[-1].strip()
+    return text.strip()
+
+
+def verify_code(code: str) -> dict:
+    """Verify kernel code and return metrics."""
     try:
-        result = import_fn()
-        startup_log.append(f"✓ {name}: {result}")
-        return True
+        if not code or "def run" not in code:
+            return {"score": 0.0, "correctness": 0.0, "cycles": None, "msg": "Invalid code"}
+
+        if "OptimizedKernelBuilder" not in code:
+            return {"score": 0.0, "correctness": 0.0, "cycles": None, "msg": "No OptimizedKernelBuilder"}
+
+        exec_globals = {"FOREST_HEIGHT": 10, "ROUNDS": 16, "BATCH_SIZE": 256}
+
+        setup_code = f'''
+import sys
+sys.path.insert(0, "{PERF_TAKEHOME_PATH}")
+from problem import Machine, Tree, Input, build_mem_image, N_CORES, reference_kernel2
+from perf_takehome import KernelBuilder, HASH_STAGES, BASELINE
+import random
+'''
+        full_code = setup_code + "\n" + code
+        exec(full_code, exec_globals)
+
+        if "OptimizedKernelBuilder" not in exec_globals:
+            return {"score": 0.0, "correctness": 0.0, "cycles": None, "msg": "Class not defined"}
+
+        random.seed(123)
+        from problem import Tree, Input, Machine, build_mem_image, N_CORES, reference_kernel2
+
+        forest = Tree.generate(10)
+        inp = Input.generate(forest, 256, 16)
+        mem = build_mem_image(forest, inp)
+
+        ref_mem = None
+        for ref_mem in reference_kernel2(list(mem)):
+            pass
+
+        if ref_mem is None:
+            return {"score": 0.0, "correctness": 0.0, "cycles": None, "msg": "Reference failed"}
+
+        kb = exec_globals["OptimizedKernelBuilder"]()
+        kb.build_kernel(10, len(forest.values), 256, 16)
+        machine = Machine(list(mem), kb.instrs, kb.debug_info(), n_cores=N_CORES)
+        machine.enable_pause = False
+        machine.enable_debug = False
+        machine.run()
+
+        cycles = machine.cycle
+
+        if cycles <= 100 or cycles > 200000:
+            return {"score": 0.0, "correctness": 0.0, "cycles": cycles, "msg": f"Bad cycles: {cycles}"}
+
+        inp_values_p = ref_mem[6]
+        expected = ref_mem[inp_values_p : inp_values_p + len(inp.values)]
+        actual = machine.mem[inp_values_p : inp_values_p + len(inp.values)]
+
+        if expected != actual:
+            return {"score": 0.0, "correctness": 0.0, "cycles": cycles, "msg": "Wrong output"}
+
+        score = SCORE_SCALE / cycles
+        return {"score": score, "correctness": 1.0, "cycles": cycles, "msg": f"OK: {cycles} cycles"}
+
     except Exception as e:
-        startup_log.append(f"✗ {name}: {str(e)[:100]}")
-        return False
-
-# Test imports
-check_import("torch", lambda: __import__("torch").__version__)
-check_import("transformers", lambda: __import__("transformers").__version__)
-check_import("datasets", lambda: __import__("datasets").__version__)
-check_import("peft", lambda: __import__("peft").__version__)
-check_import("trl", lambda: __import__("trl").__version__)
-check_import("accelerate", lambda: __import__("accelerate").__version__)
-
-# Try GRPO import
-try:
-    from trl import GRPOConfig, GRPOTrainer
-    startup_log.append("✓ GRPOTrainer: imported from trl")
-except ImportError as e:
-    startup_log.append(f"✗ GRPOTrainer from trl: {e}")
+        return {"score": 0.0, "correctness": 0.0, "cycles": None, "msg": f"Error: {str(e)[:100]}"}
+
+
+def log(msg: str):
+    """Thread-safe logging."""
+    timestamp = datetime.now().strftime("%H:%M:%S")
+    formatted = f"[{timestamp}] {msg}"
+    with training_state_lock:
+        training_state["log"].append(formatted)
+    print(formatted)
+
+
+def reward_function(completions: list[str], **kwargs) -> list[float]:
+    """Compute rewards."""
+    rewards = []
+    for completion in completions:
+        try:
+            code = extract_code_block(completion)
+            result = verify_code(code)
+            reward = result["score"]
+
+            if result["correctness"] > 0:
+                reward += 1.0
+                cycles = result.get("cycles")
+                if cycles:
+                    with training_state_lock:
+                        training_state["results"].append({"cycles": cycles, "time": time.time()})
+                        if cycles < training_state["best_cycles"]:
+                            training_state["best_cycles"] = cycles
+                            training_state["best_code"] = code
+                            log(f"NEW BEST: {cycles:,} cycles ({BASELINE_CYCLES/cycles:.2f}x)")
+
+            rewards.append(reward)
+        except Exception as e:
+            log(f"Reward error: {str(e)[:50]}")
+            rewards.append(0.0)
+
+    return rewards
+
+
+def run_training(model_name: str, num_steps: int, batch_size: int, lr: float, lora_rank: int):
+    """Main training loop."""
+    with training_state_lock:
+        training_state["running"] = True
+        training_state["step"] = 0
+        training_state["total_steps"] = num_steps
+        training_state["best_cycles"] = BASELINE_CYCLES
+        training_state["best_code"] = None
+        training_state["log"] = []
+        training_state["results"] = []
+        training_state["start_time"] = time.time()
+
+    log(f"Starting: {model_name}")
+    log(f"Steps: {num_steps}, Batch: {batch_size}, LR: {lr}")
+
     try:
-        from trl.trainer.grpo_trainer import GRPOConfig, GRPOTrainer
-        startup_log.append("✓ GRPOTrainer: imported from trl.trainer.grpo_trainer")
-    except ImportError as e2:
-        startup_log.append(f"✗ GRPOTrainer alt: {e2}")
-
-# Check CUDA
-try:
-    import torch
-    if torch.cuda.is_available():
-        startup_log.append(f"✓ CUDA: {torch.cuda.get_device_name(0)}")
-    else:
-        startup_log.append("✗ CUDA: Not available")
-except:
-    startup_log.append("✗ CUDA: Could not check")
-
-# Check simulator
-try:
-    SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
-    PERF_PATH = os.path.join(SCRIPT_DIR, "original_performance_takehome")
-    if os.path.exists(PERF_PATH):
-        sys.path.insert(0, PERF_PATH)
-        from problem import Machine, Tree
-        startup_log.append("✓ Simulator: loaded")
-    else:
-        startup_log.append(f"✗ Simulator: path not found ({PERF_PATH})")
-except Exception as e:
-    startup_log.append(f"✗ Simulator: {e}")
-
-
-def get_startup_log():
-    return "\n".join(startup_log)
-
-
-def dummy_train(model, steps):
-    return f"Would train {model} for {steps} steps\n\nImport status:\n" + get_startup_log()
-
-
-# Simple UI
-with gr.Blocks(title="VLIW Optimizer") as demo:
-    gr.Markdown("# VLIW Kernel Optimizer - Debug Mode")
-    gr.Markdown("Checking if all imports work...")
+        import torch
+        from datasets import Dataset
+        from transformers import AutoTokenizer, BitsAndBytesConfig, TrainerCallback
+        from peft import LoraConfig
+        from trl import GRPOConfig, GRPOTrainer
 
-    with gr.Row():
-        with gr.Column():
-            status = gr.Textbox(
-                label="Startup Log",
-                value=get_startup_log(),
-                lines=20
+        if torch.cuda.is_available():
+            log(f"GPU: {torch.cuda.get_device_name(0)}")
+        else:
+            log("WARNING: No GPU!")
+
+        log("Loading tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        prompt = f"{SYSTEM_PROMPT}\n\nCURRENT: {BASELINE_CYCLES} cycles. TARGET: <{TARGET_CYCLES}."
+        dataset = Dataset.from_dict({"prompt": [prompt] * 32})
+
+        peft_config = LoraConfig(
+            r=lora_rank,
+            lora_alpha=lora_rank * 2,
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        )
+
+        output_dir = f"./output/{datetime.now().strftime('%Y%m%d-%H%M%S')}"
+        os.makedirs(output_dir, exist_ok=True)
+
+        training_args = GRPOConfig(
+            output_dir=output_dir,
+            num_train_epochs=num_steps,
+            per_device_train_batch_size=batch_size,
+            learning_rate=lr,
+            logging_steps=1,
+            save_steps=max(1, num_steps // 5),
+            max_completion_length=2048,
+            temperature=0.7,
+            num_generations=4,
+            beta=0.1,
+            bf16=True,
+            report_to="none",
+        )
+
+        quant_config = None
+        if "7B" in model_name or "7b" in model_name:
+            log("Using 4-bit quantization")
+            quant_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
             )
-            refresh_btn = gr.Button("Refresh Status")
-            refresh_btn.click(get_startup_log, outputs=[status])
 
+        log("Loading model...")
+        model_kwargs = {}
+        if quant_config:
+            model_kwargs["quantization_config"] = quant_config
+
+        class StopCallback(TrainerCallback):
+            def on_step_end(self, args, state, control, **kwargs):
+                if not training_state["running"]:
+                    log("Stopping...")
+                    control.should_training_stop = True
+                return control
+
+        trainer = GRPOTrainer(
+            model=model_name,
+            reward_funcs=[reward_function],
+            args=training_args,
+            train_dataset=dataset,
+            peft_config=peft_config,
+            processing_class=tokenizer,
+            model_init_kwargs=model_kwargs,
+            callbacks=[StopCallback()],
+        )
+
+        log("Model loaded! Training...")
+        trainer.train()
+        log("Training complete!")
+
+        trainer.save_model(os.path.join(output_dir, "final"))
+        log(f"Saved to {output_dir}")
+
+        if training_state["best_code"]:
+            with open(os.path.join(output_dir, "best_code.py"), "w") as f:
+                f.write(training_state["best_code"])
+
+    except Exception as e:
+        import traceback
+        log(f"ERROR: {e}")
+        log(traceback.format_exc()[:500])
+
+    finally:
+        with training_state_lock:
+            training_state["running"] = False
+            elapsed = time.time() - (training_state["start_time"] or time.time())
+            best = training_state["best_cycles"]
+        log(f"Time: {elapsed/60:.1f} min, Best: {best:,} cycles")
+
+
+def start_training(model_name, num_steps, batch_size, lr, lora_rank):
+    if training_state["running"]:
+        return "Already running!"
+    thread = threading.Thread(
+        target=run_training,
+        args=(model_name, int(num_steps), int(batch_size), float(lr), int(lora_rank)),
+        daemon=False
+    )
+    thread.start()
+    return "Training started!"
+
+
+def stop_training():
+    with training_state_lock:
+        training_state["running"] = False
+    return "Stop signal sent."
+
+
+def get_status():
+    with training_state_lock:
+        if not training_state["start_time"]:
+            return "### Not started"
+        elapsed = time.time() - training_state["start_time"]
+        best = max(training_state["best_cycles"], 1)
+        is_running = training_state["running"]
+        logs = training_state["log"][-20:]
+
+    speedup = BASELINE_CYCLES / best
+    return f"""### {'Running' if is_running else 'Stopped'}
+| Metric | Value |
+|--------|-------|
+| Time | {elapsed/60:.1f} min |
+| Best | **{best:,}** cycles |
+| Speedup | **{speedup:.2f}x** |
+| Target | {TARGET_CYCLES:,} |
+
+```
+{chr(10).join(logs)}
+```"""
+
+
+def get_best_code():
+    with training_state_lock:
+        return training_state["best_code"] or "# No valid code yet"
+
+
+# UI
+with gr.Blocks(title="VLIW Optimizer") as demo:
+    gr.Markdown("# VLIW Kernel Optimizer via RL")
+    gr.Markdown(f"**Baseline:** {BASELINE_CYCLES:,} | **Target:** {TARGET_CYCLES:,} (108x speedup)")
+
+    with gr.Row():
         with gr.Column():
             model = gr.Dropdown(
-                choices=["Qwen/Qwen2.5-Coder-1.5B-Instruct", "Qwen/Qwen2.5-Coder-3B-Instruct"],
-                value="Qwen/Qwen2.5-Coder-1.5B-Instruct",
+                ["Qwen/Qwen2.5-Coder-7B-Instruct", "Qwen/Qwen2.5-Coder-3B-Instruct"],
+                value="Qwen/Qwen2.5-Coder-3B-Instruct",
                 label="Model"
             )
-            steps = gr.Slider(1, 10, value=3, label="Steps")
-            train_btn = gr.Button("Test Train", variant="primary")
-            output = gr.Textbox(label="Output", lines=10)
-            train_btn.click(dummy_train, inputs=[model, steps], outputs=[output])
+            steps = gr.Slider(1, 100, value=50, step=1, label="Steps")
+            batch = gr.Slider(1, 8, value=4, step=1, label="Batch")
+            lr = gr.Number(value=2e-4, label="LR")
+            lora = gr.Slider(8, 64, value=32, step=8, label="LoRA Rank")
+            with gr.Row():
+                start_btn = gr.Button("Start", variant="primary")
+                stop_btn = gr.Button("Stop", variant="stop")
+
+        with gr.Column():
+            status = gr.Markdown("### Not started")
+            refresh = gr.Button("Refresh")
+
+    with gr.Row():
+        code_out = gr.Code(label="Best Code", language="python", lines=20)
+        code_btn = gr.Button("Show Best Code")
 
+    start_btn.click(start_training, [model, steps, batch, lr, lora], [status])
+    stop_btn.click(stop_training, outputs=[status])
+    refresh.click(get_status, outputs=[status])
+    code_btn.click(get_best_code, outputs=[code_out])
+    demo.load(get_status, outputs=[status], every=5)
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)