You can deploy this model and start earning money today!

Pran-Ker/
gpu-mode-trimul

peft

safetensors

lora

reinforcement-learning

Views

No views yet

gpu-mode-trimul — LoRA checkpoint (step 30)

arXiv

GitHub

Open In Colab

LoRA adapter for gpt-oss-120b, trained with reinforcement learning on the GPU Mode TriMul competition — triangular matrix multiplication on H100.

Produced via TTT-Discover (Yuksekgonul et al., 2026):

"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 →

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

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	Hosts `gpt-oss-120b` + LoRA training	Request access
Modal	H100 GPU sandbox for kernel eval	Free tier
Weights & Biases	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