{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" }, "colab": { "provenance": [], "gpuType": "T4" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "id": "title-cell", "metadata": {}, "source": [ "# ๐Ÿ”ง Compiler Optimization RL Environment\n", "### OpenEnv Hackathon 2026 โ€” Theme #2: Long-Horizon Planning\n", "\n", "**What this notebook does:**\n", "1. Defines a fully OpenEnv-compliant `CompilerOptimizationEnv`\n", "2. Loads `Qwen2.5-3B-Instruct` via Unsloth (4-bit QLoRA)\n", "3. Trains with GRPO (TRL) โ€” LLM learns to pick compiler passes that reduce CPU cycles\n", "4. Runs a smoke-test with a mock engine so you can verify reward logic without real hardware\n", "5. Plots reward curves\n", "\n", "---\n", "**Stack:** `unsloth` ยท `trl` ยท `openenv` ยท `wandb` ยท `matplotlib`\n", "\n", "> **Runtime:** Google Colab T4 GPU recommended. For the smoke-test only, CPU is fine." ] }, { "cell_type": "markdown", "id": "install-header", "metadata": {}, "source": [ "## ๐Ÿ“ฆ Cell 1 โ€” Install Dependencies" ] }, { "cell_type": "code", "execution_count": null, "id": "install-cell", "metadata": {}, "outputs": [], "source": [ "# Install all required packages\n", "# Unsloth must be installed before trl to get the right CUDA kernels\n", "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" --quiet\n", "!pip install trl datasets transformers accelerate peft bitsandbytes --quiet\n", "!pip install wandb matplotlib --quiet\n", "\n", "# Optional: install openenv for production use\n", "# !pip install openenv --quiet\n", "\n", "print(\"โœ… All packages installed\")" ] }, { "cell_type": "markdown", "id": "env-header", "metadata": {}, "source": [ "## ๐ŸŒ Cell 2 โ€” CompilerOptimizationEnv (OpenEnv-Compliant)" ] }, { "cell_type": "code", "execution_count": null, "id": "env-cell", "metadata": {}, "outputs": [], "source": [ "\"\"\"\n", "CompilerOptimizationEnv โ€” OpenEnv-compatible RL Environment\n", "OpenEnv Hackathon 2026 | Theme #2: Long-Horizon Planning\n", "\n", "Key improvements over naive baseline:\n", " โœ“ Inherits from MCPEnvironment (OpenEnv API compliant)\n", " โœ“ Dynamic crash_penalty scaled to reward range (not hardcoded -1000)\n", " โœ“ Soft termination: invalid actions give 3 chances before episode ends\n", " โœ“ No-op detection: penalises passes that change nothing\n", " โœ“ Terminal bonus: rewards cumulative improvement, not just greedy steps\n", " โœ“ StepResult / EpisodeStats dataclasses for clean interfaces\n", " โœ“ Curriculum level support\n", "\"\"\"\n", "\n", "import copy\n", "import math\n", "from dataclasses import dataclass, field\n", "from typing import Any, Dict, List, Optional, Tuple\n", "\n", "\n", "# โ”€โ”€ Stub base class โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "# In production replace with: from openenv import MCPEnvironment\n", "class MCPEnvironment:\n", " \"\"\"Minimal stub. In production: `from openenv import MCPEnvironment`\"\"\"\n", " def reset(self, *args, **kwargs): raise NotImplementedError\n", " def step(self, *args, **kwargs): raise NotImplementedError\n", " def state(self): raise NotImplementedError\n", "\n", "\n", "# โ”€โ”€ Data classes โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "@dataclass\n", "class StepResult:\n", " observation: str\n", " reward: float\n", " done: bool\n", " info: Dict[str, Any] = field(default_factory=dict)\n", "\n", "\n", "@dataclass\n", "class EpisodeStats:\n", " steps_taken: int = 0\n", " total_reward: float = 0.0\n", " passes_applied: List[str] = field(default_factory=list)\n", " invalid_actions: int = 0\n", " no_ops: int = 0\n", " baseline_cycles: int = 0\n", " final_cycles: int = 0\n", "\n", " @property\n", " def total_improvement_pct(self) -> float:\n", " if self.baseline_cycles == 0:\n", " return 0.0\n", " return ((self.baseline_cycles - self.final_cycles) / self.baseline_cycles) * 100.0\n", "\n", "\n", "# โ”€โ”€ Core Environment โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "class CompilerOptimizationEnv(MCPEnvironment):\n", " \"\"\"\n", " RL environment where an LLM agent sequentially applies compiler\n", " optimization passes to minimise CPU cycle count while preserving\n", " program semantics.\n", "\n", " State : Pseudo-assembly representation of the current IR (text)\n", " Actions: Named compiler passes (strings) from the `passes` dictionary\n", " Reward : % cycle improvement per step โˆ’ time_tax, with terminal bonus\n", " Done : Semantic violation | max steps reached | too many invalid actions\n", " \"\"\"\n", "\n", " # Hyperparameters\n", " TIME_TAX: float = 1.0\n", " NO_OP_PENALTY: float = -2.0\n", " INVALID_ACTION_PENALTY: float = -5.0\n", " MAX_INVALID_ACTIONS: int = 3\n", " TERMINAL_BONUS_SCALE: float = 0.5\n", "\n", " def __init__(\n", " self,\n", " role1_engine,\n", " role3_passes: Dict[str, Any],\n", " max_steps: int = 10,\n", " curriculum_level: int = 1,\n", " ):\n", " self.engine = role1_engine\n", " self.passes = role3_passes\n", " self.max_steps = max_steps\n", " self.curriculum_level = curriculum_level\n", " self._valid_actions = frozenset(self.passes.keys())\n", "\n", " # Episode state\n", " self._stats: Optional[EpisodeStats] = None\n", " self.original_program = None\n", " self.current_program = None\n", " self.previous_cycles = 0\n", " self._consecutive_invalid = 0\n", "\n", " # โ”€โ”€ OpenEnv API โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " def reset(self, new_program_json: List[Dict]) -> str:\n", " self.original_program = copy.deepcopy(new_program_json)\n", " self.current_program = copy.deepcopy(new_program_json)\n", " self.previous_cycles = self._safe_count_cycles(self.current_program)\n", " self._consecutive_invalid = 0\n", " self._stats = EpisodeStats(\n", " baseline_cycles=self.previous_cycles,\n", " final_cycles=self.previous_cycles,\n", " )\n", " return self.state()\n", "\n", " def state(self) -> str:\n", " return self._program_to_pseudoasm(self.current_program)\n", "\n", " def step(self, action_string: str) -> StepResult:\n", " assert self._stats is not None, \"Call reset() before step().\"\n", " self._stats.steps_taken += 1\n", "\n", " if action_string not in self._valid_actions:\n", " return self._handle_invalid_action(action_string)\n", "\n", " candidate_program = self.passes[action_string](\n", " copy.deepcopy(self.current_program)\n", " )\n", "\n", " is_valid = self.engine.verify_equivalence(self.original_program, candidate_program)\n", " if not is_valid:\n", " return self._handle_semantic_violation()\n", "\n", " new_cycles = self._safe_count_cycles(candidate_program)\n", " reward, info = self._compute_reward(action_string, new_cycles)\n", "\n", " self.current_program = candidate_program\n", " self.previous_cycles = new_cycles\n", " self._stats.final_cycles = new_cycles\n", " self._stats.total_reward += reward\n", " self._stats.passes_applied.append(action_string)\n", " self._consecutive_invalid = 0\n", "\n", " done = self._stats.steps_taken >= self.max_steps\n", " if done:\n", " terminal_bonus = self._terminal_bonus()\n", " reward += terminal_bonus\n", " info[\"terminal_bonus\"] = terminal_bonus\n", " info[\"reason\"] = \"max_steps_reached\"\n", " info[\"episode_stats\"] = self._episode_summary()\n", "\n", " return StepResult(self.state(), reward, done, info)\n", "\n", " # โ”€โ”€ Reward logic โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " def _compute_reward(self, action: str, new_cycles: int) -> Tuple[float, Dict]:\n", " info: Dict[str, Any] = {\"action\": action}\n", " if self.previous_cycles == 0:\n", " return -self.TIME_TAX, {**info, \"note\": \"zero_baseline\"}\n", "\n", " delta_pct = ((self.previous_cycles - new_cycles) / self.previous_cycles) * 100.0\n", "\n", " if new_cycles == self.previous_cycles:\n", " reward = self.NO_OP_PENALTY\n", " self._stats.no_ops += 1\n", " info[\"no_op\"] = True\n", " else:\n", " reward = delta_pct - self.TIME_TAX\n", " info[\"delta_pct\"] = round(delta_pct, 3)\n", "\n", " info[\"prev_cycles\"] = self.previous_cycles\n", " info[\"new_cycles\"] = new_cycles\n", " return reward, info\n", "\n", " def _terminal_bonus(self) -> float:\n", " return max(0.0, self._stats.total_improvement_pct * self.TERMINAL_BONUS_SCALE)\n", "\n", " def _compute_crash_penalty(self) -> float:\n", " # 2ร— best possible episode reward โ€” always catastrophic, never overwhelming\n", " return -2.0 * (100.0 * self.max_steps)\n", "\n", " # โ”€โ”€ Error handlers โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " def _handle_invalid_action(self, action: str) -> StepResult:\n", " self._consecutive_invalid += 1\n", " self._stats.invalid_actions += 1\n", " done = self._consecutive_invalid >= self.MAX_INVALID_ACTIONS\n", " info = {\n", " \"error\": f\"Unknown action: '{action}'\",\n", " \"valid_actions\": sorted(self._valid_actions),\n", " \"consecutive_invalid\": self._consecutive_invalid,\n", " }\n", " if done:\n", " info[\"reason\"] = \"too_many_invalid_actions\"\n", " info[\"episode_stats\"] = self._episode_summary()\n", " return StepResult(self.state(), self.INVALID_ACTION_PENALTY, done, info)\n", "\n", " def _handle_semantic_violation(self) -> StepResult:\n", " return StepResult(\n", " self.state(),\n", " self._compute_crash_penalty(),\n", " True,\n", " {\n", " \"error\": \"Semantic equivalence check FAILED.\",\n", " \"reason\": \"semantic_violation\",\n", " \"episode_stats\": self._episode_summary(),\n", " },\n", " )\n", "\n", " # โ”€โ”€ State representation โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " @staticmethod\n", " def _program_to_pseudoasm(program: List[Dict]) -> str:\n", " if not program:\n", " return \"; (empty program)\"\n", " lines = []\n", " for i, instr in enumerate(program):\n", " op = instr.get(\"op\", \"NOP\")\n", " args = instr.get(\"args\", [])\n", " dest = instr.get(\"dest\")\n", " typ = instr.get(\"type\", \"\")\n", " arg_str = \", \".join(str(a) for a in args)\n", " type_hint = f\":{typ}\" if typ else \"\"\n", " if dest:\n", " lines.append(f\" {i:>3}: {dest}{type_hint} = {op} {arg_str}\")\n", " else:\n", " lines.append(f\" {i:>3}: {op} {arg_str}\")\n", " return \"\\n\".join(lines)\n", "\n", " # โ”€โ”€ Utilities โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " def _safe_count_cycles(self, program: List[Dict]) -> int:\n", " return max(0, int(self.engine.execute_and_count_cycles(program)))\n", "\n", " def _episode_summary(self) -> Dict:\n", " s = self._stats\n", " return {\n", " \"steps\": s.steps_taken,\n", " \"total_reward\": round(s.total_reward, 3),\n", " \"passes_applied\": s.passes_applied,\n", " \"invalid_actions\": s.invalid_actions,\n", " \"no_ops\": s.no_ops,\n", " \"baseline_cycles\": s.baseline_cycles,\n", " \"final_cycles\": s.final_cycles,\n", " \"total_improvement_pct\": round(s.total_improvement_pct, 3),\n", " }\n", "\n", " def available_actions(self) -> List[str]:\n", " return sorted(self._valid_actions)\n", "\n", "\n", "print(\"โœ… CompilerOptimizationEnv defined\")" ] }, { "cell_type": "markdown", "id": "smoke-header", "metadata": {}, "source": [ "## ๐Ÿงช Cell 3 โ€” Smoke Test (No GPU / Real Engine Needed)\n", "Validates the entire reward pipeline with a mock engine. Run this before spending compute." ] }, { "cell_type": "code", "execution_count": null, "id": "smoke-test-cell", "metadata": {}, "outputs": [], "source": [ "class MockEngine:\n", " \"\"\"Stub engine: cycles = instruction count, all programs semantically valid.\"\"\"\n", " def execute_and_count_cycles(self, program):\n", " return len(program)\n", "\n", " def verify_equivalence(self, original, candidate):\n", " return True\n", "\n", "\n", "MOCK_PASSES = {\n", " \"constant_folding\": lambda p: p[:-1] if len(p) > 1 else p,\n", " \"dead_code_elimination\": lambda p: p[:-1] if len(p) > 2 else p,\n", " \"loop_unrolling\": lambda p: p, # intentional no-op for testing\n", "}\n", "\n", "SAMPLE_PROGRAM = [\n", " {\"op\": \"const\", \"dest\": \"x\", \"args\": [\"5\"], \"type\": \"int\"},\n", " {\"op\": \"const\", \"dest\": \"y\", \"args\": [\"3\"], \"type\": \"int\"},\n", " {\"op\": \"add\", \"dest\": \"z\", \"args\": [\"x\", \"y\"], \"type\": \"int\"},\n", " {\"op\": \"mul\", \"dest\": \"w\", \"args\": [\"z\", \"x\"], \"type\": \"int\"},\n", " {\"op\": \"ret\", \"args\": [\"w\"]},\n", "]\n", "\n", "engine = MockEngine()\n", "env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=6)\n", "obs = env.reset(SAMPLE_PROGRAM)\n", "\n", "print(\"โ•\" * 55)\n", "print(\"SMOKE TEST\")\n", "print(\"โ•\" * 55)\n", "print(f\"Initial state:\\n{obs}\")\n", "print(f\"\\nBaseline cycles : {env.previous_cycles}\")\n", "print(f\"Available actions: {env.available_actions()}\")\n", "print()\n", "\n", "actions_to_try = [\n", " \"constant_folding\",\n", " \"dead_code_elimination\",\n", " \"loop_unrolling\", # no-op\n", " \"hallucinated_pass\", # invalid โ€” but won't kill episode yet\n", " \"constant_folding\",\n", " \"dead_code_elimination\",\n", "]\n", "\n", "for action in actions_to_try:\n", " result = env.step(action)\n", " tag = \"โœ—\" if result.reward < 0 else \"โœ“\"\n", " print(f\"{tag} '{action}'\")\n", " print(f\" reward={result.reward:+.2f} done={result.done}\")\n", " relevant = {k: v for k, v in result.info.items()\n", " if k in (\"delta_pct\", \"error\", \"no_op\", \"reason\",\n", " \"terminal_bonus\", \"episode_stats\")}\n", " if relevant:\n", " print(f\" info: {relevant}\")\n", " print()\n", " if result.done:\n", " break\n", "\n", "print(\"โœ… Smoke test passed\")" ] }, { "cell_type": "markdown", "id": "reward-plot-header", "metadata": {}, "source": [ "## ๐Ÿ“Š Cell 4 โ€” Visualise Reward Across a Mock Episode" ] }, { "cell_type": "code", "execution_count": null, "id": "reward-plot-cell", "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "import matplotlib.ticker as ticker\n", "\n", "# Run a full episode and collect data\n", "env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10)\n", "env.reset(SAMPLE_PROGRAM)\n", "\n", "sequence = [\n", " \"constant_folding\", \"dead_code_elimination\", \"loop_unrolling\",\n", " \"constant_folding\", \"dead_code_elimination\", \"loop_unrolling\",\n", " \"dead_code_elimination\", \"loop_unrolling\", \"constant_folding\",\n", " \"dead_code_elimination\",\n", "]\n", "\n", "rewards, cycle_counts, actions_log = [], [], []\n", "for act in sequence:\n", " r = env.step(act)\n", " rewards.append(r.reward)\n", " cycle_counts.append(env.previous_cycles)\n", " actions_log.append(act)\n", " if r.done:\n", " break\n", "\n", "steps = list(range(1, len(rewards) + 1))\n", "\n", "fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)\n", "fig.suptitle(\"Compiler Optimization Episode โ€” Mock Engine\", fontsize=14, fontweight=\"bold\")\n", "\n", "# Reward per step\n", "colors = [\"#2ecc71\" if r >= 0 else \"#e74c3c\" for r in rewards]\n", "ax1.bar(steps, rewards, color=colors, edgecolor=\"white\", linewidth=0.5)\n", "ax1.axhline(0, color=\"grey\", linewidth=0.8, linestyle=\"--\")\n", "ax1.set_ylabel(\"Reward\")\n", "ax1.set_title(\"Reward per Step (green = positive, red = negative)\")\n", "ax1.yaxis.set_major_formatter(ticker.FormatStrFormatter(\"%.1f\"))\n", "\n", "# Cycle count over time\n", "ax2.plot(steps, cycle_counts, marker=\"o\", color=\"#3498db\", linewidth=2, markersize=6)\n", "ax2.set_xlabel(\"Step\")\n", "ax2.set_ylabel(\"CPU Cycles\")\n", "ax2.set_title(\"CPU Cycle Count Over Episode (lower = better)\")\n", "ax2.set_xticks(steps)\n", "ax2.set_xticklabels(\n", " [a.replace(\"_\", \"\\n\") for a in actions_log],\n", " fontsize=7,\n", ")\n", "\n", "plt.tight_layout()\n", "plt.savefig(\"episode_reward_curve.png\", dpi=150, bbox_inches=\"tight\")\n", "plt.show()\n", "print(\"๐Ÿ“ˆ Plot saved as episode_reward_curve.png\")" ] }, { "cell_type": "markdown", "id": "model-header", "metadata": {}, "source": [ "## ๐Ÿค– Cell 5 โ€” Load Model with Unsloth (QLoRA 4-bit)\n", "> **Requires T4 GPU.** Skip to Cell 9 if you only want to test the environment." ] }, { "cell_type": "code", "execution_count": null, "id": "model-cell", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from unsloth import FastLanguageModel\n", "\n", "MODEL_NAME = \"unsloth/Qwen2.5-3B-Instruct\" # swap to 7B if VRAM allows\n", "MAX_SEQ_LEN = 1024\n", "LORA_RANK = 16\n", "\n", "model, tokenizer = FastLanguageModel.from_pretrained(\n", " model_name = MODEL_NAME,\n", " max_seq_length = MAX_SEQ_LEN,\n", " dtype = None, # auto bf16/fp16\n", " load_in_4bit = True,\n", ")\n", "\n", "model = FastLanguageModel.get_peft_model(\n", " model,\n", " r = LORA_RANK,\n", " target_modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n", " \"gate_proj\", \"up_proj\", \"down_proj\"],\n", " lora_alpha = LORA_RANK * 2,\n", " lora_dropout = 0.0,\n", " bias = \"none\",\n", " use_gradient_checkpointing = \"unsloth\",\n", " random_state = 42,\n", ")\n", "\n", "print(f\"โœ… Loaded {MODEL_NAME} with QLoRA rank={LORA_RANK}\")\n", "print(f\" GPU memory allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")" ] }, { "cell_type": "markdown", "id": "prompt-header", "metadata": {}, "source": [ "## ๐Ÿ’ฌ Cell 6 โ€” System Prompt & Dataset Builder" ] }, { "cell_type": "code", "execution_count": null, "id": "prompt-cell", "metadata": {}, "outputs": [], "source": [ "import json\n", "from datasets import Dataset\n", "\n", "\n", "def build_system_prompt(env: CompilerOptimizationEnv) -> str:\n", " actions = \"\\n\".join(f\" - {a}\" for a in env.available_actions())\n", " return (\n", " \"You are a compiler optimization agent. Your goal is to reduce \"\n", " \"CPU cycle count by applying optimization passes to the program below.\\n\\n\"\n", " f\"Available actions (respond with EXACTLY one per turn):\\n{actions}\\n\"\n", " \" - done (stop early if no further improvement is possible)\\n\\n\"\n", " \"Rules:\\n\"\n", " \" โ€ข Output only the action name. No explanation, no markdown.\\n\"\n", " \" โ€ข Do not invent actions not listed above.\\n\"\n", " \" โ€ข Repeating a pass that does nothing wastes a step.\\n\"\n", " )\n", "\n", "\n", "def build_dataset(\n", " programs: list,\n", " engine,\n", " passes: dict,\n", ") -> Dataset:\n", " \"\"\"\n", " Each row = one episode's initial state.\n", " GRPO samples K completions (actions) per row to estimate group-relative advantage.\n", " \"\"\"\n", " env = CompilerOptimizationEnv(engine, passes, max_steps=10)\n", " system_prompt = build_system_prompt(env)\n", "\n", " rows = []\n", " for prog in programs:\n", " obs = env.reset(prog)\n", " prompt = [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": f\"Current program:\\n{obs}\\n\\nChoose an action:\"},\n", " ]\n", " rows.append({\"prompt\": prompt, \"program_json\": json.dumps(prog)})\n", "\n", " return Dataset.from_list(rows)\n", "\n", "\n", "# โ”€โ”€ Demo: build dataset from mock programs โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "def make_mock_programs(n: int = 20) -> list:\n", " \"\"\"Generate N random mock IR programs for demo purposes.\"\"\"\n", " import random\n", " ops = [\"add\", \"mul\", \"sub\", \"const\", \"load\"]\n", " progs = []\n", " for _ in range(n):\n", " length = random.randint(4, 12)\n", " prog = [\n", " {\"op\": random.choice(ops),\n", " \"dest\": f\"v{i}\",\n", " \"args\": [f\"v{max(0,i-1)}\"],\n", " \"type\": \"int\"}\n", " for i in range(length)\n", " ]\n", " prog.append({\"op\": \"ret\", \"args\": [f\"v{length-1}\"]})\n", " progs.append(prog)\n", " return progs\n", "\n", "\n", "mock_programs = make_mock_programs(n=30)\n", "train_dataset = build_dataset(mock_programs, MockEngine(), MOCK_PASSES)\n", "\n", "print(f\"โœ… Dataset built: {len(train_dataset)} episodes\")\n", "print(f\" Sample prompt keys: {list(train_dataset[0].keys())}\")" ] }, { "cell_type": "markdown", "id": "reward-fn-header", "metadata": {}, "source": [ "## ๐ŸŽฏ Cell 7 โ€” Reward Function Factory (for GRPO)" ] }, { "cell_type": "code", "execution_count": null, "id": "reward-fn-cell", "metadata": {}, "outputs": [], "source": [ "import re\n", "from typing import Any\n", "\n", "\n", "def make_reward_fn(engine, passes, max_steps=10):\n", " \"\"\"\n", " Returns a reward function compatible with TRL's GRPOTrainer.\n", " Each GRPO rollout gets its own fresh env instance โ€” no state leakage.\n", "\n", " Action parsing strips punctuation the model might add\n", " (e.g. 'constant_folding.' โ†’ 'constant_folding').\n", " \"\"\"\n", " def reward_fn(prompts, completions, **kwargs):\n", " programs = kwargs.get(\"program_json\", [None] * len(completions))\n", " rewards = []\n", "\n", " for completion, prog_json in zip(completions, programs):\n", " # Parse action from model output\n", " raw = completion if isinstance(completion, str) else completion[0][\"content\"]\n", " action = raw.strip().lower().split()[0] if raw.strip() else \"__invalid__\"\n", " action = re.sub(r\"[^a-z0-9_]\", \"\", action) # strip punctuation\n", "\n", " # Fresh env per rollout\n", " env = CompilerOptimizationEnv(\n", " role1_engine = engine,\n", " role3_passes = passes,\n", " max_steps = max_steps,\n", " )\n", " program = json.loads(prog_json) if prog_json else []\n", " env.reset(program)\n", "\n", " if action == \"done\":\n", " rewards.append(0.0) # neutral stop\n", " else:\n", " result = env.step(action)\n", " rewards.append(result.reward)\n", "\n", " return rewards\n", "\n", " return reward_fn\n", "\n", "\n", "reward_fn = make_reward_fn(MockEngine(), MOCK_PASSES, max_steps=10)\n", "print(\"โœ… Reward function factory ready\")\n", "\n", "# Quick sanity check\n", "test_completions = [\"constant_folding\", \"hallucinated_pass\", \"loop_unrolling\"]\n", "test_programs = [json.dumps(SAMPLE_PROGRAM)] * 3\n", "test_rewards = reward_fn(\n", " prompts = [\"\"] * 3,\n", " completions = test_completions,\n", " program_json = test_programs,\n", ")\n", "print(\"\\nReward sanity check:\")\n", "for act, rew in zip(test_completions, test_rewards):\n", " print(f\" '{act}' โ†’ {rew:+.2f}\")" ] }, { "cell_type": "markdown", "id": "trainer-header", "metadata": {}, "source": [ "## ๐Ÿš€ Cell 8 โ€” GRPO Trainer Config & Training" ] }, { "cell_type": "code", "execution_count": null, "id": "trainer-cell", "metadata": {}, "outputs": [], "source": [ "from trl import GRPOConfig, GRPOTrainer\n", "\n", "# Optional W&B โ€” comment out if not using\n", "try:\n", " import wandb\n", " wandb.init(project=\"openenv-compiler-opt\", name=\"grpo-qwen2.5-3b\")\n", " REPORT_TO = \"wandb\"\n", "except Exception:\n", " REPORT_TO = \"none\"\n", "\n", "\n", "grpo_config = GRPOConfig(\n", " # โ”€โ”€ Generation โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " num_generations = 4, # K rollouts per prompt for group-relative advantage\n", " max_new_tokens = 16, # Actions are 1 word; don't waste context\n", " temperature = 0.9,\n", " top_p = 0.95,\n", "\n", " # โ”€โ”€ Optimisation โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " learning_rate = 5e-6,\n", " per_device_train_batch_size = 2,\n", " gradient_accumulation_steps = 4, # effective batch = 8\n", " num_train_epochs = 3,\n", " max_grad_norm = 0.5,\n", "\n", " # โ”€โ”€ GRPO-specific โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " beta = 0.04, # KL penalty; keeps policy near reference\n", "\n", " # โ”€โ”€ Logging / checkpointing โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " output_dir = \"./grpo_compiler_checkpoints\",\n", " logging_steps = 10,\n", " save_steps = 100,\n", " report_to = REPORT_TO,\n", "\n", " # โ”€โ”€ Reproducibility โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", " seed = 42,\n", ")\n", "\n", "trainer = GRPOTrainer(\n", " model = model,\n", " tokenizer = tokenizer,\n", " config = grpo_config,\n", " train_dataset = train_dataset,\n", " reward_funcs = reward_fn,\n", ")\n", "\n", "print(\"โœ… Trainer configured\")\n", "print(f\" num_generations (K) = {grpo_config.num_generations}\")\n", "print(f\" effective batch size = \"\n", " f\"{grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}\")\n", "print(f\" KL beta = {grpo_config.beta}\")\n", "print()\n", "print(\"Starting training... (this will take a while on T4)\")\n", "trainer.train()" ] }, { "cell_type": "markdown", "id": "save-header", "metadata": {}, "source": [ "## ๐Ÿ’พ Cell 9 โ€” Save Model" ] }, { "cell_type": "code", "execution_count": null, "id": "save-cell", "metadata": {}, "outputs": [], "source": [ "SAVE_PATH = \"./grpo_compiler_final\"\n", "\n", "model.save_pretrained(SAVE_PATH)\n", "tokenizer.save_pretrained(SAVE_PATH)\n", "\n", "print(f\"โœ… Model saved to {SAVE_PATH}\")\n", "\n", "# Optional: push to HuggingFace Hub\n", "# model.push_to_hub(\"your-hf-username/compiler-opt-grpo\")\n", "# tokenizer.push_to_hub(\"your-hf-username/compiler-opt-grpo\")" ] }, { "cell_type": "markdown", "id": "curriculum-header", "metadata": {}, "source": [ "## ๐Ÿ“ˆ Cell 10 โ€” Curriculum Callback & Reward Tracking" ] }, { "cell_type": "code", "execution_count": null, "id": "curriculum-cell", "metadata": {}, "outputs": [], "source": [ "class CurriculumCallback:\n", " \"\"\"\n", " Tracks rolling mean reward and promotes curriculum level\n", " when the agent has mastered the current difficulty.\n", "\n", " Usage: call .record(reward) after every episode.\n", " Read .level to get current difficulty (1=easy, 2=medium, 3=hard).\n", " \"\"\"\n", " def __init__(self, reward_threshold: float = 5.0, window: int = 50):\n", " self.threshold = reward_threshold\n", " self.window = window\n", " self._history = []\n", " self.level = 1\n", " self._promotions = []\n", "\n", " def record(self, reward: float, step: int = None):\n", " self._history.append(reward)\n", " if len(self._history) >= self.window:\n", " mean = sum(self._history[-self.window:]) / self.window\n", " if mean >= self.threshold and self.level < 3:\n", " self.level += 1\n", " self._promotions.append((step or len(self._history), self.level))\n", " print(f\"[Curriculum] โ–ฒ Promoted to level {self.level} \"\n", " f\"(rolling mean={mean:.2f})\")\n", "\n", " def plot(self):\n", " import matplotlib.pyplot as plt\n", " import numpy as np\n", "\n", " history = self._history\n", " steps = list(range(len(history)))\n", " window = self.window\n", " rolling = [\n", " sum(history[max(0,i-window):i+1]) / min(i+1, window)\n", " for i in steps\n", " ]\n", "\n", " fig, ax = plt.subplots(figsize=(10, 4))\n", " ax.plot(steps, history, alpha=0.3, color=\"#3498db\", label=\"Episode reward\")\n", " ax.plot(steps, rolling, color=\"#e74c3c\", linewidth=2,\n", " label=f\"Rolling mean (w={window})\")\n", " ax.axhline(self.threshold, linestyle=\"--\", color=\"grey\",\n", " linewidth=1, label=f\"Promotion threshold ({self.threshold})\")\n", " for step, level in self._promotions:\n", " ax.axvline(step, color=\"green\", linewidth=1.5, linestyle=\":\")\n", " ax.text(step, ax.get_ylim()[1]*0.9, f\" L{level}\",\n", " color=\"green\", fontsize=9)\n", " ax.set_xlabel(\"Episode\")\n", " ax.set_ylabel(\"Reward\")\n", " ax.set_title(\"Training Reward + Curriculum Progression\")\n", " ax.legend()\n", " plt.tight_layout()\n", " plt.savefig(\"curriculum_reward_curve.png\", dpi=150)\n", " plt.show()\n", " print(\"๐Ÿ“ˆ Saved curriculum_reward_curve.png\")\n", "\n", "\n", "# โ”€โ”€ Demo: simulate 200 episodes of improving reward โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€\n", "import random\n", "cb = CurriculumCallback(reward_threshold=5.0, window=50)\n", "for ep in range(200):\n", " # Simulate reward slowly improving\n", " synthetic_reward = -5 + ep * 0.08 + random.gauss(0, 2)\n", " cb.record(synthetic_reward, step=ep)\n", "\n", "cb.plot()" ] }, { "cell_type": "markdown", "id": "inference-header", "metadata": {}, "source": [ "## ๐Ÿ” Cell 11 โ€” Inference: Before vs After Training" ] }, { "cell_type": "code", "execution_count": null, "id": "inference-cell", "metadata": {}, "outputs": [], "source": [ "def run_inference_episode(model, tokenizer, engine, passes, program, max_steps=10):\n", " \"\"\"Run a greedy inference episode and return the action sequence + total improvement.\"\"\"\n", " from transformers import TextStreamer\n", " FastLanguageModel.for_inference(model)\n", "\n", " env = CompilerOptimizationEnv(engine, passes, max_steps=max_steps)\n", " obs = env.reset(program)\n", " system_prompt = build_system_prompt(env)\n", "\n", " actions_chosen, rewards_earned = [], []\n", " done = False\n", "\n", " print(f\"\\nBaseline cycles: {env.previous_cycles}\")\n", " print(f\"Initial state:\\n{obs}\\n\")\n", "\n", " while not done:\n", " messages = [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": f\"Current program:\\n{obs}\\n\\nChoose an action:\"},\n", " ]\n", " inputs = tokenizer.apply_chat_template(\n", " messages,\n", " tokenize=True,\n", " add_generation_prompt=True,\n", " return_tensors=\"pt\",\n", " ).to(model.device)\n", "\n", " with torch.no_grad():\n", " outputs = model.generate(\n", " input_ids = inputs,\n", " max_new_tokens = 16,\n", " temperature = 0.1, # greedy-ish for inference\n", " do_sample = True,\n", " )\n", "\n", " raw_action = tokenizer.decode(\n", " outputs[0][inputs.shape[-1]:], skip_special_tokens=True\n", " ).strip()\n", " action = re.sub(r\"[^a-z0-9_]\", \"\", raw_action.lower().split()[0])\n", "\n", " result = env.step(action)\n", " actions_chosen.append(action)\n", " rewards_earned.append(result.reward)\n", " obs = result.observation\n", " done = result.done\n", "\n", " print(f\"Step {len(actions_chosen)}: '{action}' โ†’ reward={result.reward:+.2f}\")\n", "\n", " summary = env._episode_summary()\n", " print(f\"\\n{'โ”€'*40}\")\n", " print(f\"Total improvement: {summary['total_improvement_pct']:.1f}%\")\n", " print(f\"Final cycles: {summary['final_cycles']} (was {summary['baseline_cycles']})\")\n", " return summary\n", "\n", "\n", "# Uncomment after training:\n", "# summary = run_inference_episode(\n", "# model, tokenizer, MockEngine(), MOCK_PASSES, SAMPLE_PROGRAM\n", "# )\n", "\n", "print(\"โœ… Inference cell ready. Uncomment the last block after training to run.\")" ] }, { "cell_type": "markdown", "id": "tips-header", "metadata": {}, "source": [ "---\n", "## ๐Ÿ“ Notes & Tips\n", "\n", "| What | Why it matters |\n", "|------|----------------|\n", "| `num_generations=4` | GRPO needs Kโ‰ฅ2 rollouts per prompt to compute group-relative advantage. K=4 balances diversity vs. compute. |\n", "| `beta=0.04` | KL penalty keeping policy close to reference. Too low โ†’ mode collapse. Too high โ†’ no learning. |\n", "| `max_new_tokens=16` | Actions are one word. This prevents wasted computation and keeps the model from adding explanations. |\n", "| Dynamic crash penalty | Computed as `โˆ’2 ร— (100 ร— max_steps)` so it always dominates the best possible episode reward without washing out all other reward signal. |\n", "| Terminal bonus | Rewards cumulative improvement, not just greedy single-step gains. Critical for long-horizon tasks. |\n", "| Soft invalid-action termination | 3 consecutive invalid actions โ†’ end. Single mistakes don't kill the episode; the agent can recover. |\n", "| `deepcopy` on all pass inputs | Role 3's passes mutate dicts in-place. Without this, `original_program` gets corrupted and the verifier fails spuriously. |\n", "\n", "**Next steps:**\n", "- Replace `MockEngine` with Role 1's real engine\n", "- Replace `MOCK_PASSES` with Role 3's real passes\n", "- Push environment to HuggingFace Spaces: `openenv init && openenv deploy`\n", "- Add W&B sweep to tune `beta`, `learning_rate`, `num_generations`" ] } ] }"