|
|
--- |
|
|
base_model: |
|
|
- deepseek-ai/DeepSeek-R1-Distill-Qwen-14B |
|
|
license: apache-2.0 |
|
|
pipeline_tag: text-generation |
|
|
library_name: transformers |
|
|
tags: |
|
|
- math |
|
|
- reasoning |
|
|
- llm |
|
|
- mathematical-reasoning |
|
|
- aimo |
|
|
datasets: |
|
|
- RabotniKuma/Fast-Math-R1-SFT |
|
|
- RabotniKuma/Fast-Math-R1-GRPO |
|
|
- open-r1/OpenR1-Math-220k |
|
|
- hoanganhpham/openr1_hard |
|
|
- qihoo360/Light-R1-SFTData |
|
|
language: |
|
|
- en |
|
|
metrics: |
|
|
- pass@1 |
|
|
--- |
|
|
|
|
|
# Kaggle AI Mathematical Olympiad - Progress Prize 2 - 9th Place Solution (Fast-Math-R1-14B) |
|
|
|
|
|
This model was presented in the paper [A Practical Two-Stage Recipe for Mathematical LLMs: Maximizing Accuracy with SFT and Efficiency with Reinforcement Learning](https://huggingface.co/papers/2507.08267). |
|
|
|
|
|
## Team |
|
|
- Hiroshi Yoshihara @ [Aillis Inc.](https://aillis.jp/en), [The Univ. of Tokyo](https://publichealth.f.u-tokyo.ac.jp/#page_home) |
|
|
- Yuichi Inoue @ [Sakana AI](https://sakana.ai) |
|
|
- Taiki Yamaguchi @ [Rist Inc.](https://www.rist.co.jp/en/) |
|
|
|
|
|
## Summary |
|
|
By applying SFT and GRPO on difficult math problems, we enhanced the performance of `DeepSeek-R1-Distill-Qwen-14B` and developed `Fast-Math-R1-14B`, |
|
|
which achieves up to 60% (on average approx. 30%) faster inference while maintaining accuracy. |
|
|
|
|
|
In addition, we trained and open-sourced `Fast-OpenMath-Nemotron-14B`, an efficiency-optimized version of NVIDIA’s `OpenMath-Nemotron-14B`, following the same approach. |
|
|
|
|
|
Technical details can be found in [Kaggle Discussion](https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/571252) and [Github](https://github.com/analokmaus/kaggle-aimo2-fast-math-r1). |
|
|
|
|
|
## Evaluation |
|
|
<img src="https://github.com/analokmaus/kaggle-aimo2-fast-math-r1/blob/master/assets/pass1_aime_all.png?raw=true" max-height="400px"> |
|
|
|
|
|
### DS-R1-Qwen-14B vs Fast-Math-R1-14B (Ours) |
|
|
| | | AIME 2024 | | AIME 2025 | | |
|
|
| ---------------------------- | ------------ | ---------------- | ------------------ | ---------------- | ------------------ | |
|
|
| Model | Token budget | Pass@1 (avg. 64) | Mean output tokens | Pass@1 (avg. 64) | Mean output tokens | |
|
|
| DeepSeek-R1-Distill-Qwen-14B | 32000 | 66.9 | 11026 | 49.9 | 12310 | |
|
|
| | 24000 | 65.7 | 10784 | 49.7 | 11978 | |
|
|
| | 16000 | 61 | 9708 | 46.2 | 10567 | |
|
|
| | 12000 | 53.7 | 8472 | 39.9 | 9008 | |
|
|
| | 8000 | 41.8 | 6587 | 31.1 | 6788 | |
|
|
| Fast-Math-R1-14B | 32000 | 68 | 8217 | 49.6 | 9663 | |
|
|
| | 24000 | 67.9 | 8209 | 49.6 | 9627 | |
|
|
| | 16000 | 66.7 | 8017 | 48.4 | 9083 | |
|
|
| | 12000 | 61.9 | 7362 | 45.2 | 8048 | |
|
|
| | 8000 | 51.4 | 5939 | 36.3 | 6174 | |
|
|
|
|
|
### OpenMath-Nemotron-14B vs Fast-OpenMath-Nemotron-14B (Ours) |
|
|
| | | AIME 2024 | | AIME 2025 | | |
|
|
| -------------------------- | ------------ | ---------------- | ------------------ | ---------------- | ------------------ | |
|
|
| Model | Token budget | Pass@1 (avg. 64) | Mean output tokens | Pass@1 (avg. 64) | Mean output tokens | |
|
|
| OpenMath-Nemotron-14B | 32000 | 76.2 | 11493 | 64.5 | 13414 | |
|
|
| | 24000 | 75.4 | 11417 | 63.4 | 13046 | |
|
|
| | 16000 | 66 | 10399 | 54.2 | 11422 | |
|
|
| | 12000 | 55 | 9053 | 40 | 9609 | |
|
|
| | 8000 | 36 | 6978 | 27.2 | 7083 | |
|
|
| [Fast-OpenMath-Nemotron-14B](https://huggingface.co/RabotniKuma/Fast-OpenMath-Nemotron-14B) | 32000 | 70.7 | 9603 | 61.4 | 11424 | |
|
|
| | 24000 | 70.6 | 9567 | 60.9 | 11271 | |
|
|
| | 16000 | 66.6 | 8954 | 55.3 | 10190 | |
|
|
| | 12000 | 59.4 | 7927 | 45.6 | 8752 | |
|
|
| | 8000 | 47.6 | 6282 | 33.8 | 6589 | |
|
|
|
|
|
### Qwen3-14B vs Fast-Math-Qwen3-14B |
|
|
| | | AIME 2024 | | AIME 2025 | | |
|
|
| ------------------- | ------------ | ---------------- | ------------------ | ---------------- | ------------------ | |
|
|
| Model | Token budget | Pass@1 (avg. 64) | Mean output tokens | Pass@1 (avg. 64) | Mean output tokens | |
|
|
| Qwen3-14B | 32000 | 79.3 | 13669 | 69.5 | 16481 | |
|
|
| | 24000 | 75.9 | 13168 | 65.6 | 15235 | |
|
|
| | 16000 | 64.5 | 11351 | 50.4 | 12522 | |
|
|
| | 12000 | 49.7 | 9746 | 36.3 | 10353 | |
|
|
| | 8000 | 28.4 | 7374 | 19.5 | 7485 | |
|
|
| [Fast-Math-Qwen3-14B](https://huggingface.co/RabotniKuma/Fast-Math-Qwen3-14B) | 32000 | 77.6 | 9740 | 66.6 | 12281 | |
|
|
| | 24000 | 76.5 | 9634 | 65.3 | 11847 | |
|
|
| | 16000 | 72.6 | 8793 | 60.1 | 10195 | |
|
|
| | 12000 | 65.1 | 7775 | 49.4 | 8733 | |
|
|
| | 8000 | 50.7 | 6260 | 36 | 6618 | |
|
|
|
|
|
## Dataset |
|
|
- [Our first stage SFT dataset](https://huggingface.co/datasets/RabotniKuma/Fast-Math-R1-SFT) |
|
|
- [Our second stage GRPO dataset](https://huggingface.co/datasets/RabotniKuma/Fast-Math-R1-GRPO) |
|
|
|
|
|
## Inference |
|
|
### vLLM |
|
|
```python |
|
|
from vllm import LLM, SamplingParams |
|
|
from transformers import AutoTokenizer |
|
|
|
|
|
|
|
|
model_path = 'RabotniKuma/Fast-Math-R1-14B' |
|
|
vllm_engine = LLM( |
|
|
model=model_path, |
|
|
max_model_len=8192, |
|
|
gpu_memory_utilization=0.9, |
|
|
trust_remote_code=True, |
|
|
) |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path) |
|
|
|
|
|
|
|
|
sampling_params = SamplingParams( |
|
|
temperature=1.0, |
|
|
top_p=0.90, |
|
|
min_p=0.05, |
|
|
max_tokens=8192, |
|
|
stop='</think>', # For even faster inference, applying early stopping at the </think> tag and extracting the final boxed content is recommended. |
|
|
) |
|
|
messages = [ |
|
|
{ |
|
|
'role': 'user', |
|
|
'content': ( |
|
|
'Solve the problem, and put the answer in \\\\boxed{{}}. ' |
|
|
'Sarah is twice as old as her youngest brother. If the difference between their ages is 15 years. How old is her youngest brother?' |
|
|
) |
|
|
} |
|
|
] |
|
|
messages = tokenizer.apply_chat_template( |
|
|
conversation=messages, |
|
|
tokenize=False, |
|
|
add_generation_prompt=True |
|
|
) |
|
|
response = vllm_engine.generate(messages, sampling_params=sampling_params) |
|
|
``` |
|
|
|
|
|
## Training models |
|
|
### 1. Installation |
|
|
```bash |
|
|
poetry lock |
|
|
poetry install --no-root |
|
|
``` |
|
|
|
|
|
### 2. First stage training |
|
|
Training time: approx. 10 hours (8× H200 GPUs) |
|
|
```bash |
|
|
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ |
|
|
accelerate launch --config_file accelerate_configs/deepspeed_zero3.yaml --num_processes 8 \ |
|
|
experiments/train_first_stage.py |
|
|
``` |
|
|
<img src="https://github.com/analokmaus/kaggle-aimo2-fast-math-r1/blob/master/assets/wandb_stage1.png?raw=true" max-height="300px"> |
|
|
|
|
|
### 3. Second stage training |
|
|
Training time: approx. 10 hours (8× H200 GPUs) |
|
|
```bash |
|
|
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ |
|
|
accelerate launch --config_file accelerate_configs/deepspeed_zero2.yaml --num_processes 8 \ |
|
|
experiments/train_second_stage.py |
|
|
``` |
|
|
<img src="https://github.com/analokmaus/kaggle-aimo2-fast-math-r1/blob/master/assets/wandb_stage2.png?raw=true" max-height="600px"> |
|
|
|
|
|
### (Optional) Token scheduler training |
|
|
Training time: approx. 1 hours (8× H200 GPUs) |
|
|
|
|
|
The token scheduler is a lightweight model that predicts the difficulty of a problem, measured by how many tokens the R1 model requires before reaching the final answer. See [Kaggle discussion](https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/571252) for details. |
|
|
```bash |
|
|
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ |
|
|
accelerate launch --config_file accelerate_configs/deepspeed_zero3.yaml --num_processes 8 \ |
|
|
experiments/train_token_scheduler.py |
|
|
``` |
|
|
<img src="https://github.com/analokmaus/kaggle-aimo2-fast-math-r1/blob/master/assets/wandb_token_scheduler.png?raw=true" max-height="300px"> |
|
|
|
|
|
### (Optional) Fast-OpenMath-Nemotron-14B |
|
|
Training time: approx. 12 hours (8× H200 GPUs) |
|
|
```bash |
|
|
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ |
|
|
accelerate launch --config_file accelerate_configs/deepspeed_zero3.yaml --num_processes 8 \ |
|
|
experiments/train_fast_nemotron_14b.py |
|
|
``` |
|
|
|
|
|
### (Optional) Fast-Math-Qwen3-14B |
|
|
Training time: approx. 12 hours (8× H200 GPUs) |
|
|
|
|
|
**Note:** You’ll need to update your dependencies to train any of the Qwen3 series models. |
|
|
```bash |
|
|
# Update environment |
|
|
cp dev/pyproject_qwen3.toml pyproject.toml |
|
|
poetry lock |
|
|
poetry install --no-root |
|
|
# Train |
|
|
CUDA_VISIBLE_DEVICES=0,1,2,3 \ |
|
|
accelerate launch --config_file accelerate_configs/deepspeed_zero3_cpu_offload.yaml --num_processes 4 \ |
|
|
experiments/train_fast_qwen3_14b.py & |
|
|
CUDA_VISIBLE_DEVICES=4,5,6,7 trl vllm-serve --model Qwen/Qwen3-14B --tensor_parallel_size 2 --data_parallel_size 2 & |
|
|
wait |
|
|
``` |
|
|
|
|
|
## Technical details |
|
|
Detailed report is available on [Kaggle Disucussion](https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/571252). |
|
|
|
|
|
### First stage: intensive SFT using a high-difficulty dataset |
|
|
#### Dataset |
|
|
- [OpenR1 Math](https://huggingface.co/datasets/open-r1/OpenR1-Math-220k): We randomly sampled 3000 examples where the R1’s trace had more than 12800 tokens and an accuracy of over 50%, along with another 3000 examples where the accuracy ranged between 50% and 75%. |
|
|
- [openr1_hard](https://huggingface.co/datasets/hoanganhpham/openr1_hard): "~2.5k hard samples from open-r1-math-220k. Samples deemed as hard were unsolvable by r1-distill-32b after 4 tries." |
|
|
- [Light-R1-SFTData](https://huggingface.co/datasets/qihoo360/Light-R1-SFTData): We used the 2nd stage data from Light-R1-SFTData. |
|
|
|
|
|
We merged all the datasets mentioned above, removed duplicates, and selected the correct generation with the shortest token length. For samples in the Light-R1 dataset where ground truth answers were not provided, we extracted and substituted the answers from the R1 traces. As a result, we constructed a **high-difficulty dataset consisting of 7900 problem - R1 trace - answer sets**. |
|
|
|
|
|
[Our first stage SFT dataset](https://huggingface.co/datasets/RabotniKuma/Fast-Math-R1-SFT) |
|
|
|
|
|
#### Training |
|
|
A full-parameter supervised fine-tuning training was conducted on a machine with 8 H200 GPUs, using the SFTTrainer from the trl library. |
|
|
|
|
|
### Second stage: GRPO for more efficient reasoning |
|
|
#### Dataset |
|
|
- [Light-R1-SFTData](https://huggingface.co/datasets/qihoo360/Light-R1-SFTData): We extracted the answers from the 2nd stage SFT data of Light-R1. |
|
|
|
|
|
[Our second stage GRPO dataset](https://huggingface.co/datasets/RabotniKuma/Fast-Math-R1-GRPO) |
|
|
|
|
|
#### Training |
|
|
We used the [faster implementation of trl GRPOTrainer](https://github.com/nhannguyen2709/open-r1). |
|
|
|
|
|
Reward functions: |
|
|
1. Format reward |
|
|
|
|
|
In order to save output tokens, we forced the model to give an answer in the end of reasoning block before `</think>` by rewarding the pattern `r"^.*?oxed{(.*?)}.*?</think>.*?$"`. Generation is stopped at `</think>` during inference. |
|
|
|
|
|
2. Cosine reward |
|
|
|
|
|
Compared to a normal accuracy-based reward, cosine reward applies a continuous penalty to longer correct reasoning traces and shorter incorrect ones. |
|
|
|
|
|
3. Length reward |
|
|
|
|
|
Length-based rewards to discourage overthinking and promote token efficiency. |
|
|
Paper: https://arxiv.org/abs/2501.12599 |
