{ "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": [ "# CompilerOptimizationEnv, PASSES, reverse passes, terminal-weighted reward (toyir_rl_support)\n", "import os\n", "import sys\n", "\n", "if os.getcwd() not in sys.path:\n", " sys.path.insert(0, os.getcwd())\n", "\n", "from toyir_rl_support import (\n", " MCPEnvironment,\n", " StepResult,\n", " EpisodeStats,\n", " CompilerOptimizationEnv,\n", " PASSES,\n", " MOCK_PASSES,\n", " compute_shaped_reward,\n", " rollout_shaped_return,\n", " log_reverse_pass_stats_for_completion,\n", " ensure_import_path,\n", " MockEngine,\n", " MOCK_ENGINE,\n", " SAMPLE_PROGRAM,\n", ")\n", "\n", "ensure_import_path()\n", "print(\"โœ… Loaded CompilerOptimizationEnv, PASSES, and reward helpers from toyir_rl_support\")\n" ] }, { "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": [ "# Smoke test (MockEngine; same PASSES as training)\n", "engine = MockEngine()\n", "env = CompilerOptimizationEnv(\n", " engine, MOCK_PASSES, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS\n", ")\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", " \"peephole_optimization\", # no-op\n", " \"hallucinated_pass\", # invalid โ€” but won't kill episode yet\n", " \"constant_folding\",\n", " \"dead_code_elimination\",\n", " \"expand_constant\",\n", " \"STOP\",\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} is_terminal={result.info.get('is_terminal')}\")\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\", \"is_terminal\")}\n", " if relevant:\n", " print(f\" info: {relevant}\")\n", " print()\n", " if result.done:\n", " break\n", "\n", "print(\"โœ… Smoke test passed\")\n" ] }, { "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=CompilerOptimizationEnv.MAX_EPISODE_STEPS)\n", "env.reset(SAMPLE_PROGRAM)\n", "\n", "sequence = [\n", " \"constant_folding\", \"dead_code_elimination\", \"peephole_optimization\",\n", " \"expand_constant\", \"peephole_optimization\",\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\")\n" ] }, { "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", "import re\n", "from datasets import Dataset\n", "\n", "\n", "class Deliverable2_Formatter:\n", " 'State translation + robust action-array extraction for Role 2.'\n", "\n", " @staticmethod\n", " def translate_state(raw_json: list) -> str:\n", " pseudo_assembly = []\n", " for i, instruction in enumerate(raw_json):\n", " if not isinstance(instruction, dict):\n", " pseudo_assembly.append(f\"{i}. NOP\")\n", " continue\n", " op = str(instruction.get(\"op\", \"UNKNOWN\")).upper()\n", " args = \", \".join(str(arg) for arg in instruction.get(\"args\", []))\n", " dest = instruction.get(\"dest\", \"\")\n", " if dest:\n", " pseudo_assembly.append(f\"{i}. {dest} = {op} {args}\".rstrip())\n", " else:\n", " pseudo_assembly.append(f\"{i}. {op} {args}\".rstrip())\n", " return \"\\n\".join(pseudo_assembly) if pseudo_assembly else \"; (empty program)\"\n", "\n", " @staticmethod\n", " def extract_action_array(llm_output: str) -> list:\n", " text = (llm_output or \"\").strip()\n", " if not text:\n", " raise ValueError(\"Invalid JSON format\")\n", "\n", " try:\n", " parsed = json.loads(text)\n", " if isinstance(parsed, list):\n", " return parsed\n", " except json.JSONDecodeError:\n", " pass\n", "\n", " match = re.search(r\"\\[.*?\\]\", text, re.DOTALL)\n", " if match:\n", " try:\n", " parsed = json.loads(match.group(0))\n", " if isinstance(parsed, list):\n", " return parsed\n", " except json.JSONDecodeError:\n", " pass\n", "\n", " raise ValueError(\"Invalid JSON format\")\n", "\n", "\n", "def build_system_prompt(passes: dict) -> str:\n", " lines: list = []\n", " for name in sorted(passes.keys()):\n", " if name == \"expand_constant\":\n", " lines.append(\n", " \" - expand_constant: Splits a constant assignment into runtime arithmetic (e.g., CONST 8 becomes CONST 3 + ADD 5). May temporarily increase cycles but can enable forward passes to find better optimization chains. Use sparingly when standard passes seem stuck.\"\n", " )\n", " elif name == \"duplicate_computation\":\n", " lines.append(\n", " \" - duplicate_computation: Creates a redundant copy of a binary operation with a fresh variable. May temporarily increase cycles but can enable alternative dead code elimination paths. Use when redundancy might unlock further simplification.\"\n", " )\n", " else:\n", " lines.append(f\" - {name}\")\n", " action_block = \"\\n\".join(lines)\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:\\n{action_block}\\n\"\n", " \" - STOP (emit this in the array to stop; optional alias: \\\"done\\\". Stop early if no further improvement is possible.)\\n\\n\"\n", " \"Rules:\\n\"\n", " \" โ€ข Output only a JSON array of pass names (example: [\\\"constant_folding\\\"]).\\n\"\n", " \" โ€ข No explanation, no markdown, no extra text.\\n\"\n", " \" โ€ข Do not invent actions not listed above.\\n\"\n", " \" โ€ข Maximum 5 passes per response (hard episode cap: 5 optimization steps or STOP).\\n\"\n", " )\n", "\n", "\n", "def build_dataset(\n", " programs: list,\n", " engine,\n", " passes: dict,\n", ") -> Dataset:\n", " # Each row = one episode; GRPO samples K completions per row\n", " env = CompilerOptimizationEnv(\n", " engine, passes, max_steps=CompilerOptimizationEnv.MAX_EPISODE_STEPS\n", " )\n", " system_prompt = build_system_prompt(passes)\n", "\n", " rows = []\n", " for prog in programs:\n", " translated_state = Deliverable2_Formatter.translate_state(prog)\n", " prompt = [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": f\"Current program:\\n{translated_state}\\n\\nChoose optimization passes:\"},\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\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", " {\n", " \"op\": random.choice(ops),\n", " \"dest\": f\"v{i}\",\n", " \"args\": [f\"v{max(0, i-1)}\"],\n", " \"type\": \"int\",\n", " }\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())}\")\n" ] }, { "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 json\n", "from typing import Any\n", "\n", "# Rollout + terminal weighting; reverse-pass W&B window logging\n", "from toyir_rl_support import (\n", " CompilerOptimizationEnv,\n", " MOCK_PASSES,\n", " rollout_shaped_return,\n", " log_reverse_pass_stats_for_completion,\n", " SAMPLE_PROGRAM,\n", ")\n", "\n", "EP_CAP = CompilerOptimizationEnv.MAX_EPISODE_STEPS\n", "\n", "\n", "def _normalize_action_seq(actions: list, cap: int, passes: dict) -> list:\n", " out: list = []\n", " for x in actions[:cap]:\n", " raw = str(x).strip()\n", " s_low = raw.lower()\n", " if s_low in (\"stop\", \"done\"):\n", " out.append(\"done\" if s_low == \"done\" else \"STOP\")\n", " continue\n", " key = None\n", " for k in passes:\n", " if k.lower() == s_low:\n", " key = k\n", " break\n", " if key is None:\n", " out.append(raw)\n", " else:\n", " out.append(key)\n", " return out\n", "\n", "\n", "def make_reward_fn(engine, passes, max_steps: int = EP_CAP):\n", " # TRL GRPO reward: terminal-weighted `rollout_shaped_return` + reverse-pass logging\n", " _cap = min(int(max_steps), EP_CAP)\n", "\n", " def reward_fn(prompts, completions, **kwargs):\n", " programs = kwargs.get(\"program_json\", [None] * len(completions))\n", " rewards: list = []\n", "\n", " for completion, prog_json in zip(completions, programs):\n", " raw = completion if isinstance(completion, str) else completion[0][\"content\"]\n", " if prog_json is None:\n", " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", " continue\n", " program = (\n", " json.loads(prog_json) if isinstance(prog_json, str) else list(prog_json)\n", " )\n", " if not program:\n", " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", " continue\n", " try:\n", " actions = Deliverable2_Formatter.extract_action_array(raw)\n", " except ValueError:\n", " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", " continue\n", " if not actions:\n", " rewards.append(CompilerOptimizationEnv.INVALID_ACTION_PENALTY)\n", " continue\n", " norm = _normalize_action_seq(actions, _cap, passes)\n", " tr, n_e, n_d, _hs, _hc, ok = rollout_shaped_return(\n", " program, norm, engine, passes\n", " )\n", " if not ok:\n", " rewards.append(-1000.0)\n", " else:\n", " log_reverse_pass_stats_for_completion(n_e, n_d)\n", " rewards.append(tr)\n", " return rewards\n", "\n", " return reward_fn\n", "\n", "\n", "reward_fn = make_reward_fn(MockEngine(), MOCK_PASSES, max_steps=EP_CAP)\n", "print(\"โœ… Reward function factory ready\")\n", "\n", "# Quick sanity check\n", "test_completions = [\n", " '[\"constant_folding\"]',\n", " 'Here is my plan: [\"dead_code_elimination\", \"peephole_optimization\", \"done\"]',\n", " \"hallucinated_pass\",\n", "]\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!r} โ†’ {rew:+.2f}\")\n" ] }, { "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=5):\n", " \"\"\"Run inference and execute parsed pass arrays until episode termination.\"\"\"\n", " FastLanguageModel.for_inference(model)\n", "\n", " _ms = min(int(max_steps), CompilerOptimizationEnv.MAX_EPISODE_STEPS)\n", " env = CompilerOptimizationEnv(engine, passes, max_steps=_ms)\n", " obs = env.reset(program)\n", " system_prompt = build_system_prompt(passes)\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 optimization passes:\"},\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=32,\n", " temperature=0.1,\n", " do_sample=True,\n", " )\n", "\n", " raw_output = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True).strip()\n", "\n", " try:\n", " parsed_actions = Deliverable2_Formatter.extract_action_array(raw_output)\n", " except ValueError:\n", " parsed_actions = []\n", "\n", " if not parsed_actions:\n", " result = env.step(\"__invalid__\")\n", " actions_chosen.append(\"__invalid__\")\n", " rewards_earned.append(result.reward)\n", " obs = result.observation\n", " done = result.done\n", " print(f\"Step {len(actions_chosen)}: invalid output '{raw_output}' โ†’ reward={result.reward:+.2f}\")\n", " continue\n", "\n", " for action in parsed_actions[:CompilerOptimizationEnv.MAX_EPISODE_STEPS]:\n", " a = str(action).strip()\n", " al = a.lower()\n", " if al in (\"done\", \"stop\"):\n", " result = env.step(\"done\" if al == \"done\" else \"STOP\")\n", " actions_chosen.append(al)\n", " rewards_earned.append(result.reward)\n", " obs = result.observation\n", " done = result.done\n", " print(f\"Step {len(actions_chosen)}: stop ('{a}') โ†’ reward={result.reward:+.2f}\")\n", " break\n", " action = al\n", " if action not in passes:\n", " for k in passes:\n", " if k.lower() == al:\n", " action = k\n", " break\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", " print(f\"Step {len(actions_chosen)}: '{action}' โ†’ reward={result.reward:+.2f}\")\n", " if done:\n", " break\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", "| Equivalence / semantic failure | The shaped reward function returns `โˆ’1000.0` when a candidate fails verification (RLVR hard gate). |\n", "| Terminal reward | Final cycle savings (scaled) are given at `STOP`, step cap (5), or at end of a pass list; non-terminal steps use a small constant cost. |\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`" ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 5 }