feat(notebook): storytelling layer — origin story, research lineage, honest speedup framing, future vision

eval/space-src/arm_gym_grpo_colab.ipynb

@@ -2,13 +2,13 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "6e5601b0",
+   "id": "0f57be22",
    "metadata": {},
    "source": [
-    "# ARM-Gym — GRPO Training Notebook\n",
+    "# ARM-Gym — Teaching a 7B Model to Write Assembly\n",
+    "\n",
+    "> *Can a language model learn to write ARM assembly that a compiler's static rules would never produce?*\n",
     "\n",
-    "> **Teaching Qwen2.5-Coder-7B to write ARM AArch64 assembly that beats `clang-21 -O3`**\n",
-    ">\n",
     "> Meta / HuggingFace OpenEnv Hackathon India 2026 · Wild Card (Theme 5) · Team (dot)mkv\n",
     "\n",
     "| Resource | Link |\n",
@@ -25,7 +25,7 @@
     "1. **Setup** — installs ARM toolchain + Python packages\n",
     "2. **Environment overview** — shows ARM-Gym's kernel space (15 templates, 649 variants)\n",
     "3. **Reward & pipeline** — explains the 3-gate verifier and reward formula\n",
-    "4. **Training results** — V11 logs loaded from HF Hub, 5 plots generated ✅ (pre-run)\n",
+    "4. **Training results** — V11 logs loaded from HF Hub, 5 plots generated (pre-run) ✅\n",
     "5. **Run comparison** — V10 vs V11 head-to-head with all metrics\n",
     "6. **Evidence of learning** — explicit before/after proof the agent improved\n",
     "7. **Optional** — run GRPO training yourself (requires L40S GPU)\n",
@@ -35,7 +35,36 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fc889d63",
+   "id": "323776bb",
+   "metadata": {},
+   "source": [
+    "## The Problem — And Why a Language Model?\n",
+    "\n",
+    "Modern compilers like `clang -O3` are extraordinary engineering — decades of expert rules for loop unrolling, register allocation, instruction scheduling. But they are **heuristic-bound**. Every rule is a human decision about what is *usually* faster. For a specific kernel on a specific chip, the optimal instruction sequence may exist but violate those rules — and the compiler will never find it.\n",
+    "\n",
+    "This is the gap ARM-Gym is exploring: not replacing compilers, but using an LLM as a **probabilistic scout** that explores instruction orderings a human engineer would never hard-code. When the scout finds something, a deterministic verifier (assembler + QEMU + LLVM-MCA) confirms it. The AI proposes; the tools decide.\n",
+    "\n",
+    "### Standing on the shoulders of prior work\n",
+    "\n",
+    "This approach has a research lineage we built directly on:\n",
+    "\n",
+    "- **AlphaDev (DeepMind, Nature 2023)** — proved that RL over assembly search could find a sort3 algorithm 70% faster than LLVM's libc++, now merged into the production standard library. But AlphaDev used MCTS on x86 and required expensive per-problem search.\n",
+    "- **Meta LLM Compiler (Cummins et al., 2023)** — first language model trained for LLVM pass ordering. Achieved 3% code-size reduction over `-Oz` with *zero* extra compilations at inference. Showed LLMs can learn compiler semantics, not just pattern-match.\n",
+    "- **Compiler-R1 (NeurIPS 2025)** — applied R1-style GRPO to LLVM pass ordering, same algorithm we use. 8.46% instruction count reduction. Confirmed that RL recovers gains SFT leaves on the table.\n",
+    "- **SuperCoder (arXiv 2505.11480)** — the closest prior art. Trained Qwen2.5-Coder via GRPO to emit x86-64 assembly that exceeded `gcc -O3`, achieving 1.46x speedup at 95% correctness. The ARM ISA — despite powering AWS Graviton5, Azure Cobalt 100, and Apple Silicon — was untouched.\n",
+    "\n",
+    "**ARM-Gym ports SuperCoder's approach to AArch64**, replacing x86-64 with the Neoverse V2 microarchitecture, native execution with QEMU, and wall-clock timing with LLVM-MCA cycle estimates — a deterministic, Docker-safe reward signal that runs in milliseconds.\n",
+    "\n",
+    "### An honest framing\n",
+    "\n",
+    "ARM-Gym does not universally beat `clang-21 -O3`. That is not the claim. The compiler is very good. What our best run (V11) shows is that a 7B model, trained for ~110 minutes with GRPO, learned to write semantically correct AArch64 assembly — starting from 19% correctness and reaching 70% in the final 20 steps. On certain kernel shapes, at certain parameter combinations, it finds instruction sequences that LLVM-MCA estimates as faster. The V10 run found a +47.5% speedup on one kernel. V11 reached +14.5% while being far more consistent. **The trajectory is what matters** — correctness rising monotonically is the precondition for speedup to follow.\n",
+    "\n",
+    "The real bet is longer: with more training, larger models, and silicon validation, patterns discovered this way can be back-ported into LLVM as new deterministic rules — or eventually embedded as a `-O-ai` compiler pass that searches for the specific chip being targeted."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2bbf5824",
    "metadata": {},
    "source": [
     "## Training Run History\n",
@@ -91,12 +120,12 @@
     "- **1024-token window** removed the ceiling that was truncating longer kernel solutions in V10.\n",
     "- **LoRA r=32** (vs r=24) gave the adapter more capacity to represent AArch64 instruction patterns.\n",
     "\n",
-    "V11's lower best-speedup (14.5% vs 47.5%) reflects a trade-off: the model stopped taking syntactically lucky shortcuts and began generating semantically correct assembly consistently. With more training steps, speedup events would follow the correctness gains."
+    "V11's lower single-run best speedup (14.5% vs 47.5%) reflects a real trade-off: the model stopped taking syntactically lucky shortcuts and began generating semantically correct assembly consistently across a wider range of kernels. Correctness at 70% is the precondition for speedup to become reliable. With more training steps, the speedup events would follow. Both figures are LLVM-MCA Neoverse V2 estimates — directionally meaningful, not silicon-validated."
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "19a8c35a",
+   "id": "26f50165",
    "metadata": {},
    "source": [
     "## 1. System Setup (Linux / Colab)"
@@ -105,7 +134,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a9eda18a",
+   "id": "518ebd55",
    "metadata": {},
    "outputs": [
     {
@@ -148,7 +177,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ec507e68",
+   "id": "2e4b6918",
    "metadata": {},
    "outputs": [
     {
@@ -174,7 +203,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "953fb0f5",
+   "id": "7f59efda",
    "metadata": {},
    "source": [
     "## 2. ARM-Gym Environment Overview ✅"
@@ -183,7 +212,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "832b2dfc",
+   "id": "239b5a08",
    "metadata": {},
    "outputs": [
     {
@@ -215,7 +244,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "864a6964",
+   "id": "8b55fea8",
    "metadata": {},
    "outputs": [
     {
@@ -243,11 +272,15 @@
   },
   {
    "cell_type": "markdown",
-   "id": "b53a22c9",
+   "id": "baa3eda8",
    "metadata": {},
    "source": [
     "## 3. Reward & Training Pipeline ✅\n",
     "\n",
+    "The core design principle, borrowed from SuperCoder and Compiler-R1, is **verifiable reward without a judge**: every reward signal comes from a real tool, not another LLM. This is what makes the learning signal trustworthy — the model cannot sweet-talk an assembler or fake a QEMU execution.\n",
+    "\n",
+    "GRPO (Group Relative Policy Optimization) — the same algorithm behind DeepSeek-R1 — generates G=8 candidate assemblies per kernel, runs all of them through the verifier, and trains the model to favor the ones that scored higher relative to the group. No labelled dataset required. No reference solutions. Just the reward.\n",
+    "\n",
     "### The 3-Gate Verifier\n",
     "\n",
     "Every model attempt passes through **3 gates**. Fail any gate → reward = 0. Pass all → reward is the sum.\n",
@@ -281,7 +314,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "4810444b",
+   "id": "5654a348",
    "metadata": {},
    "outputs": [
     {
@@ -328,7 +361,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "f6e7a3bb",
+   "id": "13629c79",
    "metadata": {},
    "source": [
     "## 4. Training Results — V11 (Best Run) ✅\n",
@@ -340,7 +373,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a137d287",
+   "id": "ac75a677",
    "metadata": {},
    "outputs": [
     {
@@ -389,7 +422,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "33fd2fc1",
+   "id": "b0b32749",
    "metadata": {},
    "outputs": [
     {
@@ -439,7 +472,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d90db96d",
+   "id": "50e16611",
    "metadata": {},
    "outputs": [
     {
@@ -459,7 +492,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "571b43a4",
+   "id": "9304a206",
    "metadata": {},
    "source": [
     "### Plot 1 — V11 Training Loss ✅"
@@ -468,7 +501,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7c181ec7",
+   "id": "f29c3678",
    "metadata": {},
    "outputs": [
     {
@@ -498,7 +531,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "87a9dba1",
+   "id": "0bc57f41",
    "metadata": {},
    "source": [
     "### Plot 2 — All 4 Reward Components ✅"
@@ -507,7 +540,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "9c2b86c1",
+   "id": "7a0c2cec",
    "metadata": {},
    "outputs": [
     {
@@ -551,7 +584,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "fb6869a7",
+   "id": "20cb4f0d",
    "metadata": {},
    "source": [
     "### Plot 3 — QEMU Correctness Rate ✅"
@@ -560,7 +593,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a7e60b84",
+   "id": "0b5d9f9c",
    "metadata": {},
    "outputs": [
     {
@@ -601,7 +634,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "65315830",
+   "id": "d05f0589",
    "metadata": {},
    "source": [
     "### Plot 4 — Speedup Events & V10 vs V11 Comparison ✅"
@@ -610,7 +643,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "0f3fc5ae",
+   "id": "bdf294f1",
    "metadata": {},
    "outputs": [
     {
@@ -669,7 +702,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "10862fd0",
+   "id": "5b5b2681",
    "metadata": {},
    "source": [
     "### Plot 5 — Total Reward Trajectory ✅"
@@ -678,7 +711,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b36df2ce",
+   "id": "670ddd69",
    "metadata": {},
    "outputs": [
     {
@@ -716,11 +749,13 @@
   },
   {
    "cell_type": "markdown",
-   "id": "75c5c738",
+   "id": "69915eef",
    "metadata": {},
    "source": [
     "## 5. V10 vs V11 — Run Comparison ✅\n",
     "\n",
+    "Across 11 experiments we tuned one thing at a time. V10 → V11 was the most significant jump: more completions per step (G 6→8), longer generation window (768→1024 tokens), stronger LoRA adapter (r 24→32). The table below shows where each version wins.\n",
+    "\n",
     "| Metric | [V10](https://huggingface.co/ZDC-M01/arm-gym-v10-train-200) | [V11](https://huggingface.co/ZDC-M01/arm-gym-v11-train-250) | Winner |\n",
     "|--------|------|------|--------|\n",
     "| Steps | 200 | 250 | — |\n",
@@ -730,17 +765,19 @@
     "| Best speedup | **+47.5%** | +14.5% | V10 |\n",
     "| Mean total reward | 3.79 | **4.72** | V11 |\n",
     "| Correctness mean | 26.0% | **38.6%** | V11 |\n",
-    "| Correctness final 20 steps | 25.8% | **70.0%** | V11 🔥 |\n",
+    "| Correctness final 20 steps | 25.8% | **70.0%** | V11 |\n",
     "| Q4 reward | 3.96 | **6.50** | V11 |\n",
     "| Reward trajectory | flat | **monotonic ↑** | V11 |\n",
     "\n",
-    "> V10's +47.5% best speedup was a single lucky rollout — the overall reward trajectory was flat. V11's reward rose monotonically every quarter, ending at 70% correctness. V11 is the better model and the primary submission."
+    "> **Reading the table honestly:** V10's +47.5% best speedup was a single lucky rollout on one kernel — the overall reward trajectory was flat, meaning the model wasn't consistently improving. V11's lower single-run best (+14.5%) reflects a real trade-off: the model stopped relying on syntactic shortcuts and started generating semantically correct assembly more consistently. Correctness is the prerequisite. Speedup follows. V11 is the better model.\n",
+    "\n",
+    "Both results are **MCA-model speedups** (LLVM-MCA Neoverse V2 scheduling model) — not wall-clock measurements on silicon. They are estimates of cycle throughput, which correlate with real performance but are not the same thing."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "ec152508",
+   "id": "9a2eaf51",
    "metadata": {},
    "outputs": [
     {
@@ -784,7 +821,7 @@
   },
   {
    "cell_type": "markdown",
-   "id": "702e19ca",
+   "id": "ceb99547",
    "metadata": {},
    "source": [
     "## 6. Evidence of Learning — Did the Agent Actually Improve? ✅"
@@ -793,7 +830,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "191b9d29",
+   "id": "c0f0b54f",
    "metadata": {},
    "outputs": [
     {
@@ -852,29 +889,43 @@
   },
   {
    "cell_type": "markdown",
-   "id": "cc13ffdb",
+   "id": "5f8fac26",
    "metadata": {},
    "source": [
-    "## 7. Key Findings\n",
+    "## 7. Key Findings & What Comes Next\n",
+    "\n",
+    "### Results\n",
     "\n",
     "| Metric | V11 (Best Run) |\n",
     "|--------|----------------|\n",
     "| Training | 250 steps, NVIDIA L40S 48GB, ~110 min |\n",
-    "| **Best speedup** | **+14.5% over clang-21 -O3** (LLVM-MCA Neoverse V2) |\n",
+    "| Best speedup (MCA estimate) | +14.5% over clang-21 -O3 on one kernel |\n",
     "| Speedup win rate | 14/125 steps (11.2%) |\n",
-    "| Peak total reward | **9.03** (step 98) |\n",
-    "| Correctness: start → end | 19% → **70%** (3.7× improvement) |\n",
-    "| Reward trajectory | Q1=3.30 → Q4=6.50 (monotonic) |\n",
-    "| Fully-correct batches | 14 (all G=8 pass QEMU) |\n",
+    "| Peak total reward | 9.03 (step 98) |\n",
+    "| Correctness: start → end | 19% → 70% (monotonic, 3.7× improvement) |\n",
+    "| Reward trajectory | Q1=3.30 → Q4=6.50 |\n",
+    "| Fully-correct batches | 14 (all G=8 pass QEMU adversarial tests) |\n",
+    "\n",
+    "All speedup figures are **LLVM-MCA Neoverse V2 cycle estimates**, not wall-clock measurements on silicon. They are meaningful engineering signals but should be labeled as model speedups until validated on physical ARM hardware.\n",
     "\n",
-    "All speedup claims are **MCA-model speedup** (LLVM-MCA Neoverse V2) until silicon-validated on Graviton hardware.\n",
+    "### Where this is going\n",
+    "\n",
+    "The current results are a proof of concept, not a finished product. The interesting question is what happens next:\n",
+    "\n",
+    "**Near term — Pattern extraction.** When the model finds a faster instruction sequence on a specific kernel, that sequence can be studied and back-ported into LLVM as a new deterministic optimization rule. The AI acts as a research tool that discovers the exception; the compiler authors write the rule.\n",
+    "\n",
+    "**Medium term — A neural compiler pass.** Instead of `-O3`, a developer could run `-O-ai` and get a search-based optimization for the specific chip they are targeting. Not a replacement for the compiler — a specialist pass that runs on top of it, exploring the residual headroom that static heuristics miss.\n",
+    "\n",
+    "**Long term — Software-defined silicon.** Closing the loop between chip design and software. When ARM or AWS designs a new instruction (say, an SME2 matrix operation), an RL agent can probe: *can code actually take advantage of this?* That feedback loop, today done manually, could become automated.\n",
+    "\n",
+    "The bet is not that a 7B model beats compilers today. The bet is that the combination of LLM reasoning over assembly + deterministic hardware simulation as a reward signal is a viable path toward compiler optimization that no fixed rule set can match.\n",
     "\n",
     "> **Models:** [V11 — Best](https://huggingface.co/ZDC-M01/arm-gym-v11-train-250) · [V10 — Second Best](https://huggingface.co/ZDC-M01/arm-gym-v10-train-200)"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "2f743799",
+   "id": "a295411d",
    "metadata": {},
    "source": [
     "## 8. (Optional) Run GRPO Training\n",
@@ -887,7 +938,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "029cf632",
+   "id": "b02cc636",
    "metadata": {},
    "outputs": [
     {
@@ -917,13 +968,15 @@
   },
   {
    "cell_type": "markdown",
-   "id": "c5b28b24",
+   "id": "81852bbd",
    "metadata": {},
    "source": [
     "---\n",
     "\n",
     "## Resources\n",
     "\n",
+    "### This project\n",
+    "\n",
     "| Resource | Link |\n",
     "|----------|------|\n",
     "| HuggingFace Space | [kaori02/arm-gym](https://huggingface.co/spaces/kaori02/arm-gym) |\n",
@@ -931,7 +984,15 @@
     "| Second-best — V10 | [ZDC-M01/arm-gym-v10-train-200](https://huggingface.co/ZDC-M01/arm-gym-v10-train-200) |\n",
     "| Training scripts | [ZDC-M01/arm-gym-pkg](https://huggingface.co/datasets/ZDC-M01/arm-gym-pkg) |\n",
     "| Blog post | [arm-gym/blob/main/blog.md](https://huggingface.co/spaces/kaori02/arm-gym/blob/main/blog.md) |\n",
-    "| SuperCoder (prior art) | [arxiv 2505.11480](https://arxiv.org/abs/2505.11480) |\n",
+    "\n",
+    "### Prior work & research lineage\n",
+    "\n",
+    "| Paper | What it showed |\n",
+    "|-------|----------------|\n",
+    "| [SuperCoder (arXiv 2505.11480)](https://arxiv.org/abs/2505.11480) | GRPO-trained LLM beats gcc -O3 on x86-64, 1.46× speedup, 95% correctness — the direct blueprint ARM-Gym ports to AArch64 |\n",
+    "| [AlphaDev — DeepMind (Nature 2023)](https://www.nature.com/articles/s41586-023-06004-9) | RL over x86 assembly found a sort3 algorithm 70% faster than libc++ — now in production LLVM |\n",
+    "| [Compiler-R1 (NeurIPS 2025)](https://arxiv.org/abs/2410.10053) | GRPO applied to LLVM pass ordering, 8.46% instruction count reduction — validates our GRPO + verifiable reward approach |\n",
+    "| [Meta LLM Compiler (Cummins et al., 2023)](https://arxiv.org/abs/2309.07062) | First LLM trained for compiler pass ordering, 3% code-size reduction over -Oz with zero extra compilations at inference |\n",
     "\n",
     "**Team:** (dot)mkv · Meta / HuggingFace OpenEnv Hackathon India 2026 · Wild Card Theme"
    ]