{ "nbformat": 4, "nbformat_minor": 0, "metadata": { "colab": { "provenance": [], "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# TTT-Discover: GPU Kernel Optimization\n", "\n", "[![arXiv](https://img.shields.io/badge/arXiv-2601.16175-b31b1b.svg)](https://arxiv.org/abs/2601.16175)\n", "[![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://test-time-training.github.io/discover/)\n", "[![HF Checkpoint](https://img.shields.io/badge/HF-Pran--Ker/gpu--mode--trimul-yellow)](https://huggingface.co/Pran-Ker/gpu-mode-trimul)\n", "[![GitHub](https://img.shields.io/badge/GitHub-discover-black?logo=github)](https://github.com/test-time-training/discover)\n", "\n", "Run [TTT-Discover](https://arxiv.org/abs/2601.16175) on GPU kernel optimization tasks, starting from a pre-trained checkpoint that already knows how to write fast Triton kernels.\n", "\n", "**What TTT-Discover does:** performs reinforcement learning *at test time* — the model keeps training on your specific problem, earning rewards from real GPU execution, until it finds a great solution.\n", "\n", "## Available checkpoints\n", "\n", "| Checkpoint | Step | Best reward | Best runtime (H100) | Tinker path |\n", "|---|---|---|---|---|\n", "| `step_061` | 61 | 0.412 | ~3638 μs | `tinker://55bc74de-...:train:0/weights/000061` |\n", "| `step_030` | 30 | 0.412 | ~3638 μs | `tinker://681a070d-...:train:0/weights/000030` |\n", "\n", "**step_061 is recommended** — it is 31 more steps of RL training beyond step_030.\n", "\n", "---\n", "\n", "### Prerequisites\n", "\n", "| Service | Purpose | Get it |\n", "|---------|---------|--------|\n", "| [Tinker](https://tinker.thinkingmachines.dev) | LLM training + inference (hosts `gpt-oss-120b`) | Request access |\n", "| [Modal](https://modal.com) | GPU sandbox for kernel evaluation | Free tier available |\n", "| [Weights & Biases](https://wandb.ai) | Run tracking | Free account |\n", "\n", "> **No local GPU required** — training runs on Tinker's cluster; kernel evaluation runs on Modal H100s." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Install dependencies" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!pip install -q ttt-discover modal wandb\n", "!git clone --depth 1 https://github.com/test-time-training/discover.git 2>/dev/null || echo 'already cloned'\n", "!pip install -q -e discover/" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Configure API keys\n", "\n", "Use the Colab **Secrets panel** (🔑 left sidebar) to store these, then enable each for this notebook:\n", "- `TINKER_API_KEY`\n", "- `MODAL_TOKEN_ID` and `MODAL_TOKEN_SECRET`\n", "- `WANDB_API_KEY`\n", "- `WANDB_ENTITY` (your W&B username)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "from google.colab import userdata\n", "\n", "for key in [\"TINKER_API_KEY\", \"MODAL_TOKEN_ID\", \"MODAL_TOKEN_SECRET\", \"WANDB_API_KEY\", \"WANDB_ENTITY\"]:\n", " try:\n", " os.environ[key] = userdata.get(key)\n", " except Exception:\n", " print(f\" ⚠ {key} not found in secrets — set it before running training cells\")\n", "\n", "print(\"Keys loaded.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Choose a checkpoint\n", "\n", "Run this cell once to select which pre-trained checkpoint to start from." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# --- Pick one ---\n", "CHECKPOINT = \"step_061\" # recommended: 61 steps of RL training\n", "# CHECKPOINT = \"step_030\" # earlier checkpoint, 30 steps\n", "# CHECKPOINT = None # start from scratch (base gpt-oss-120b)\n", "\n", "CHECKPOINT_PATHS = {\n", " \"step_061\": \"tinker://55bc74de-c858-54ce-9756-e6f54d7a5a8d:train:0/weights/000061\",\n", " \"step_030\": \"tinker://681a070d-2ef4-5b8c-a216-d4f22dca1efb:train:0/weights/000030\",\n", "}\n", "SAMPLER_PATHS = {\n", " \"step_061\": \"tinker://55bc74de-c858-54ce-9756-e6f54d7a5a8d:train:0/sampler_weights/000061\",\n", " \"step_030\": \"tinker://681a070d-2ef4-5b8c-a216-d4f22dca1efb:train:0/sampler_weights/000030\",\n", "}\n", "\n", "load_checkpoint_path = CHECKPOINT_PATHS.get(CHECKPOINT) if CHECKPOINT else None\n", "sampler_path = SAMPLER_PATHS.get(CHECKPOINT) if CHECKPOINT else None\n", "print(f\"Checkpoint: {CHECKPOINT or 'scratch'}\")\n", "print(f\"State path: {load_checkpoint_path or '(none — fresh LoRA)'}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4a. Run TTT-Discover on TriMul (H100)\n", "\n", "Continues RL training from the selected checkpoint on the GPU Mode TriMul triangular matrix multiplication competition." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import asyncio, os, sys\n", "sys.path.insert(0, \"discover\")\n", "\n", "from ttt_discover.rl.train import Config, main as rl_main\n", "from ttt_discover.tinker_utils import misc_utils\n", "from ttt_discover.tinker_utils.dataset_builder import DatasetConfig, get_single_problem_dataset_builder\n", "from examples.gpu_mode.env import GpuModeEnv\n", "\n", "EXPERIMENT_NAME = f\"gpu-mode-trimul-from-{CHECKPOINT or 'scratch'}\" # change to taste\n", "log_path = f\"./tinker_log/{EXPERIMENT_NAME}\"\n", "os.makedirs(log_path, exist_ok=True)\n", "\n", "dataset_builder = get_single_problem_dataset_builder(DatasetConfig(\n", " env_type=GpuModeEnv,\n", " problem_type=\"trimul\",\n", " batch_size=4, # groups per batch — reduce if you hit resource limits\n", " group_size=16, # rollouts per group\n", " model_name_for_tokenizer=\"openai/gpt-oss-120b\",\n", " renderer_name=\"gpt_oss_high_reasoning\",\n", " num_cpus_per_task=0,\n", " eval_timeout=530,\n", " log_path=log_path,\n", "))\n", "\n", "config = Config(\n", " env_type=GpuModeEnv,\n", " problem_type=\"trimul\",\n", " learning_rate=4e-5,\n", " dataset_builder=dataset_builder,\n", " model_name=\"openai/gpt-oss-120b\",\n", " lora_rank=32,\n", " temperature=1.0,\n", " wandb_project=\"gpu-mode\",\n", " wandb_name=EXPERIMENT_NAME,\n", " log_path=log_path,\n", " load_checkpoint_path=load_checkpoint_path, # None = fresh start\n", " kl_penalty_coef=0.1,\n", " num_substeps=1,\n", " save_every=1,\n", " num_epochs=20,\n", " loss_fn=\"importance_sampling\",\n", " adv_estimator=\"entropic_adaptive_beta\",\n", " adv_estimator_beta=2.0,\n", " remove_constant_reward_groups=True,\n", " phase1_max_tokens=26000,\n", " local_model_path=None,\n", ")\n", "\n", "misc_utils.check_log_dir(log_path, behavior_if_exists=\"resume\")\n", "asyncio.run(rl_main(config))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 4b. Bring your own GPU kernel problem\n", "\n", "Plug in any GPU kernel problem by defining your own reward function and prompt." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import subprocess\n", "from ttt_discover import Environment, BaseRewardEvaluator, State, DiscoverConfig, discover\n", "\n", "\n", "class MyKernelReward(BaseRewardEvaluator):\n", " \"\"\"Reward = 1 / runtime_ms. Replace the evaluation block with your own benchmark.\"\"\"\n", "\n", " def get_reward(self, code: str, state: State) -> dict:\n", " if \"@triton.jit\" not in code:\n", " return {\"reward\": 0.0, \"correctness\": 0.0, \"raw_score\": -1e6,\n", " \"msg\": \"No @triton.jit found.\", \"result_construction\": [], \"stdout\": \"\"}\n", "\n", " # ── replace this block with your benchmark ──────────────────────────\n", " with open(\"/tmp/kernel.py\", \"w\") as f:\n", " f.write(code)\n", " result = subprocess.run(\n", " [\"python\", \"/tmp/kernel.py\"], capture_output=True, text=True, timeout=30\n", " )\n", " runtime_ms = 1.0 # TODO: parse from result.stdout\n", " # ────────────────────────────────────────────────────────────────────\n", "\n", " return {\n", " \"reward\": 1.0 / runtime_ms,\n", " \"correctness\": 1.0,\n", " \"raw_score\": runtime_ms,\n", " \"msg\": f\"runtime: {runtime_ms:.3f} ms\",\n", " \"result_construction\": [],\n", " \"stdout\": result.stdout,\n", " }\n", "\n", "\n", "class MyKernelEnv(Environment):\n", " reward_function = MyKernelReward\n", " state_type = State\n", "\n", " def get_question(self) -> str:\n", " return \"\"\"\n", "You are an expert GPU kernel engineer. Your task is to implement a fast Triton kernel for:\n", "\n", " \n", "\n", "Requirements:\n", "- Use Triton 3.3.1 on H100\n", "- Implement in a single ```python``` code block\n", "- Include a short docstring summarizing your algorithm\n", "\"\"\"\n", "\n", "\n", "discover(DiscoverConfig(\n", " env_type=MyKernelEnv,\n", " problem_type=\"my_kernel\",\n", " experiment_name=\"my-kernel-discover\",\n", " wandb_project=\"my-kernels\",\n", " group_size=8,\n", " groups_per_batch=2,\n", " num_epochs=30,\n", "))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Single-shot inference (no RL loop)\n", "\n", "Just want to generate one kernel from the checkpoint without running the full training loop?" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import tinker, asyncio\n", "\n", "async def ask(prompt: str, sampler: str) -> str:\n", " svc = tinker.ServiceClient(base_url=None)\n", " client = await svc.create_sampling_client_async(sampler)\n", " resp = await client.sample_async(tinker.SampleRequest(\n", " model_input=tinker.ModelInput.from_text(prompt),\n", " sampling_params=tinker.SamplingParams(temperature=0.8, max_new_tokens=4096),\n", " ))\n", " return resp.completion_text\n", "\n", "\n", "PROMPT = \"\"\"Write a fast Triton kernel for triangular matrix multiplication on an H100.\n", "The operation computes: out[i,j] = sum_k A[i,k] * B[k,j] where A is lower triangular.\n", "Implement in a single ```python``` block with @triton.jit.\"\"\"\n", "\n", "result = asyncio.run(ask(PROMPT, sampler_path))\n", "print(result)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "## Citation\n", "\n", "```bibtex\n", "@article{ttt-discover2026,\n", " title = {Learning to Discover at Test Time},\n", " author = {Yuksekgonul, Mert and Koceja, Daniel and Li, Xinhao\n", " and Bianchi, Federico and McCaleb, Jed and Wang, Xiaolong\n", " and Kautz, Jan and Choi, Yejin and Zou, James\n", " and Guestrin, Carlos and Sun, Yu},\n", " journal = {arXiv preprint arXiv:2601.16175},\n", " year = {2026}\n", "}\n", "```\n", "\n", "[Paper](https://arxiv.org/abs/2601.16175) · [Project Page](https://test-time-training.github.io/discover/) · [GitHub](https://github.com/test-time-training/discover) · [Checkpoint](https://huggingface.co/Pran-Ker/gpu-mode-trimul)" ] } ] }