Qwen3-1.7B Triton kernel writer, GRPO run 9 (step 1698, branch main)
Writes Triton GPU kernels for PyTorch modules. Part of an independent reimplementation of TritonRL (Woo et al., arXiv 2510.17891): supervised fine-tuning on verified teacher kernels, then GRPO with a reward that compiles and runs each candidate kernel against the PyTorch reference.
- Start: cgpadwick2020/qwen3-1.7b-triton-sft-v2 (Qwen3-1.7B, SFT on about 20K verified DeepSeek kernels over KernelBook tasks).
- GRPO run 9: 2 epochs over 6,799 tasks, 1,698 steps: the KernelBook training tasks plus size-augmented variants of them (the same modules at KernelBench-like input sizes, up to 2^25 elements and rows of 65,537). A completion on an augmented task is rewarded only if it is correct at the large size and at the original. 8 tasks x 16 samples per step, lr 3e-6, 4,096-token completions, full fine-tune; checkpoint saved in fp32, uploaded in bf16.
- Branches:
main= final (step 1698);step-1600= best held-out validation;step-1500.
Held-out validation during training (152 tasks x 8 samples, temperature 0.6, pass@1 / pass@8): step 1500 52.5 / 69.7, step 1600 52.5 / 70.4, step 1698 51.7 / 67.8 (SFT start 29.0 / 55.3).
Final checkpoint: test split (152 tasks x 16 samples) pass@1 51.3, pass@16 70.4; KernelBench at the problem sizes of early 2025 (the version the paper used) with KernelBench's own evaluator, pass@10: Level 1 80%, Level 2 41%.
Prompt: the reference module in a Python code block after "Here is the reference PyTorch module. get_inputs()
and get_init_inputs() show how it will be constructed and called.", ending "Write ModelNew as specified.";
chat template with an empty think block. Sample at temperature 0.6, top-k 20, top-p 0.95.
- Downloads last month
- 38
Model tree for cgpadwick2020/qwen3-1.7b-triton-grpo-run9
Base model
Qwen/Qwen3-1.7B-Base