README.md

---
license: mit
base_model: openai/gpt-oss-120b
tags:
  - lora
  - peft
  - reinforcement-learning
  - gpu-kernels
  - triton
  - test-time-training
---

<p align="center">
  <h1 align="center">gpu-mode-trimul — LoRA checkpoint (step 30)</h1>
  <p align="center">
    <a href="https://arxiv.org/abs/2601.16175"><img src="https://img.shields.io/badge/arXiv-2601.16175-b31b1b.svg" alt="arXiv"></a>
    <a href="https://test-time-training.github.io/discover/"><img src="https://img.shields.io/badge/Project-Page-blue" alt="Project Page"></a>
    <a href="https://github.com/test-time-training/discover"><img src="https://img.shields.io/badge/GitHub-discover-black?logo=github" alt="GitHub"></a>
    <a href="https://colab.research.google.com/github/Pran-Ker/gpu-mode-trimul/blob/main/gpu_mode_trimul.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
  </p>
</p>

LoRA adapter for **gpt-oss-120b**, trained with reinforcement learning on the
[GPU Mode TriMul](https://www.gpumode.com/v2/leaderboard/496?tab=rankings) competition —
triangular matrix multiplication on H100.

Produced via **TTT-Discover** ([Yuksekgonul et al., 2026](https://arxiv.org/abs/2601.16175)):

> *"We perform reinforcement learning at test time, allowing the LLM to continue training
> with experience specific to the problem at hand … Our test-time training runs are performed
> using Tinker, an API by Thinking Machines, with a cost of only a few hundred dollars per problem."*

---

## What is TTT-Discover?

Instead of prompting a frozen model (like AlphaEvolve), TTT-Discover **keeps training** on your
specific problem at test time. The model earns a reward signal from real execution feedback
(`reward = 1500 / runtime_μs`) and learns to write faster Triton kernels through trial and error —
no human-written examples needed.

**Published results (H100, TriMul):**

| | A100 ↓ | H100 ↓ | B200 ↓ | MI300x ↓ |
|---|---|---|---|---|
| Best Human | 4531 μs | 1371 μs | 1005 μs | 2462 μs |
| **TTT-Discover** | **2198 μs** | **1161 μs** | **905 μs** | **1596 μs** |

[Verify on GPU Mode leaderboard →](https://www.gpumode.com/v2/leaderboard/496?tab=rankings)

This checkpoint reached **~3638 μs** at step 14 (best reward 0.412).
Starting from it instead of scratch saves ~14 steps of cold-start exploration.

---

## Reward trajectory (this run)

| Step | reward_max | reward_mean | runtime (best) |
|------|-----------|------------|----------------|
| 7  | 0.213 | 0.054 | ~7040 μs |
| 8  | 0.273 | 0.126 | ~5490 μs |
| 10 | 0.281 | 0.173 | ~5340 μs |
| 13 | 0.376 | 0.124 | ~3990 μs |
| **14** | **0.412** | **0.118** | **~3638 μs** ← best |
| 18 | 0.281 | 0.152 | ~5340 μs |
| 25 | 0.281 | 0.156 | ~5340 μs |
| 30 | 0.280 | 0.210 | ~5360 μs |

> Steps 15–17 had zero reward due to an eval cluster billing limit — not a training failure.

---

## Files

```
sampler_weights/
  adapter_model.safetensors   # LoRA weights (~5 GB) — use for inference
  adapter_config.json         # PEFT config (rank=32, target=all-linear)
  checkpoint_complete         # Completion marker
gpu_mode_trimul.ipynb         # Google Colab notebook
```

Training state (for resuming RL with fresh optimizer):
`tinker://681a070d-2ef4-5b8c-a216-d4f22dca1efb:train:0/weights/000030`

---

## Quick start

[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Pran-Ker/gpu-mode-trimul/blob/main/gpu_mode_trimul.ipynb)

Open the notebook above to run TTT-Discover on TriMul, fork from this checkpoint,
or plug in your own GPU kernel problem.

### What you need

| Service | Purpose | Get it |
|---------|---------|--------|
| [Tinker](https://tinker.thinkingmachines.dev) | Hosts `gpt-oss-120b` + LoRA training | Request access |
| [Modal](https://modal.com) | H100 GPU sandbox for kernel eval | Free tier |
| [Weights & Biases](https://wandb.ai) | Run tracking | Free account |

> No local GPU required — training runs on Tinker's cluster; kernel evals run on Modal H100s.

### Warm-start RL from this checkpoint

```python
import asyncio, os
from ttt_discover.rl.train import Config, main as rl_main
from ttt_discover.tinker_utils import misc_utils
from ttt_discover.tinker_utils.dataset_builder import DatasetConfig, get_single_problem_dataset_builder
from examples.gpu_mode.env import GpuModeEnv  # from github.com/test-time-training/discover

CHECKPOINT = "tinker://681a070d-2ef4-5b8c-a216-d4f22dca1efb:train:0/weights/000030"
EXPERIMENT  = "my-trimul-run"
log_path    = f"./tinker_log/{EXPERIMENT}"
os.makedirs(log_path, exist_ok=True)

dataset_builder = get_single_problem_dataset_builder(DatasetConfig(
    env_type=GpuModeEnv, problem_type="trimul",
    batch_size=4, group_size=16,
    model_name_for_tokenizer="openai/gpt-oss-120b",
    renderer_name="gpt_oss_high_reasoning",
    num_cpus_per_task=0, eval_timeout=530, log_path=log_path,
))

config = Config(
    env_type=GpuModeEnv, problem_type="trimul",
    learning_rate=4e-5, dataset_builder=dataset_builder,
    model_name="openai/gpt-oss-120b", lora_rank=32,
    wandb_project="gpu-mode", wandb_name=EXPERIMENT,
    log_path=log_path,
    load_checkpoint_path=CHECKPOINT,   # warm start ← key line
    num_epochs=20, save_every=1,
    kl_penalty_coef=0.1, phase1_max_tokens=26000,
    loss_fn="importance_sampling",
    adv_estimator="entropic_adaptive_beta", adv_estimator_beta=2.0,
    remove_constant_reward_groups=True, num_substeps=1, local_model_path=None,
)

misc_utils.check_log_dir(log_path, behavior_if_exists="resume")
asyncio.run(rl_main(config))
```

### Single-shot inference

```python
import tinker, asyncio

SAMPLER = "tinker://681a070d-2ef4-5b8c-a216-d4f22dca1efb:train:0/sampler_weights/000030"

async def ask(prompt):
    svc = tinker.ServiceClient(base_url=None)
    client = await svc.create_sampling_client_async(SAMPLER)
    resp = await client.sample_async(tinker.SampleRequest(
        model_input=tinker.ModelInput.from_text(prompt),
        sampling_params=tinker.SamplingParams(temperature=0.8, max_new_tokens=4096),
    ))
    return resp.completion_text

print(asyncio.run(ask("Write a fast Triton kernel for triangular matmul on H100.")))
```

---

## Paper

**Learning to Discover at Test Time**  
Mert Yuksekgonul\*, Daniel Koceja\*, Xinhao Li\*, Federico Bianchi\*, Jed McCaleb,
Xiaolong Wang, Jan Kautz, Yejin Choi, James Zou†, Carlos Guestrin†, Yu Sun  
*Stanford · NVIDIA · Astera Institute · UC San Diego · Together AI*  
[arXiv:2601.16175](https://arxiv.org/abs/2601.16175) · [Project page](https://test-time-training.github.io/discover/) · [PDF](https://test-time-training.github.io/discover.pdf)

```bibtex
@article{ttt-discover2026,
  title   = {Learning to Discover at Test Time},
  author  = {Yuksekgonul, Mert and Koceja, Daniel and Li, Xinhao
             and Bianchi, Federico and McCaleb, Jed and Wang, Xiaolong
             and Kautz, Jan and Choi, Yejin and Zou, James
             and Guestrin, Carlos and Sun, Yu},
  journal = {arXiv preprint arXiv:2601.16175},
  year    = {2026}
}
```

---

## Acknowledgments

- **[GPU Mode](https://github.com/gpu-mode)** — community for GPU kernel optimization and the TriMul competition
- **[Tinker](https://tinker.thinkingmachines.dev)** — LLM training and RL infrastructure by Thinking Machines