Improve model card: Add pipeline tag, library name, update paper link and enhance details (#1)

- Improve model card: Add pipeline tag, library name, update paper link and enhance details (6e8c7ddd6983c984051d14a3100a36e2e0e6e08d)


Co-authored-by: Niels Rogge <[email protected]>

README.md

@@ -1,25 +1,48 @@
 ---
-license: apache-2.0
 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
+## 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
+## 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)
+### 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 | 
@@ -34,7 +57,7 @@ Technical details can be found in [Kaggle Discussion](https://www.kaggle.com/com
 |                              | 12000        | 61.9             | 7362               | 45.2             | 8048               | 
 |                              | 8000         | 51.4             | 5939               | 36.3             | 6174               | 
 
-## OpenMath-Nemotron-14B vs Fast-OpenMath-Nemotron-14B (Ours)
+### 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 | 
@@ -49,7 +72,7 @@ Technical details can be found in [Kaggle Discussion](https://www.kaggle.com/com
 |                            | 12000        | 59.4             | 7927               | 45.6             | 8752               | 
 |                            | 8000         | 47.6             | 6282               | 33.8             | 6589               | 
 
-## Qwen3-14B vs Fast-Math-Qwen-14B
+### 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 | 
@@ -64,13 +87,12 @@ Technical details can be found in [Kaggle Discussion](https://www.kaggle.com/com
 |                     | 12000        | 65.1             | 7775               | 49.4             | 8733               | 
 |                     | 8000         | 50.7             | 6260               | 36               | 6618               | 
 
-
-# Dataset
+## 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
+## Inference
+### vLLM
 ```python
 from vllm import LLM, SamplingParams
 from transformers import AutoTokenizer
@@ -97,7 +119,7 @@ messages = [
     {
         'role': 'user', 
         'content': (
-            'Solve the problem, and put the answer in \boxed{{}}. '
+            '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?'
         )
     }
@@ -108,4 +130,104 @@ messages = tokenizer.apply_chat_template(
     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