AI-MO/NuminaMath-72B-TIR

Model Card for NuminaMath 72B TIR

NuminaMath is a series of language models that are trained with two stages of supervised fine-tuning to solve math problems using chain of thought (CoT) and tool-integrated reasoning (TIR):

• Stage 1: fine-tune the base model on a large, diverse dataset of natural language math problems and solutions, where each solution is templated with Chain of Thought (CoT) to facilitate reasoning.
• Stage 2: fine-tune the model from Stage 1 on a synthetic dataset of tool-integrated reasoning, where each math problem is decomposed into a sequence of rationales, Python programs, and their outputs.

Model description

• Model type: A 72B parameter math LLM fine-tuned on a dataset with 860k+ math problem-solution pairs.
• Language(s) (NLP): Primarily English
• License: Tongyi Qianwen
• Finetuned from model: Qwen/Qwen2-72B

Model Sources

• Repository: https://github.com/project-numina/aimo-progress-prize

Intended uses & limitations

Here's how you can run the model using the pipeline() function from 🤗 Transformers:

import re
import torch
from transformers import pipeline

pipe = pipeline("text-generation", model="AI-MO/NuminaMath-72B-TIR", torch_dtype=torch.bfloat16, device_map="auto")

messages = [
    {"role": "user", "content": "For how many values of the constant $k$ will the polynomial $x^{2}+kx+36$ have two distinct integer roots?"},
]
prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

gen_config = {
    "max_new_tokens": 1024,
    "do_sample": False,
    "stop_strings": ["```output"], # Generate until Python code block is complete
    "tokenizer": pipe.tokenizer,
}

outputs = pipe(prompt, **gen_config)
text = outputs[0]["generated_text"]
print(text)

# WARNING: This code will execute the python code in the string. We show this for eductional purposes only.
# Please refer to our full pipeline for a safer way to execute code.
python_code = re.findall(r"```python(.*?)```", text, re.DOTALL)[0]
exec(python_code)

The above executes a single step of Python code - for more complex problems, you will want to run the logic for several steps to obtain the final solution.

Bias, Risks, and Limitations

NuminaMath 72B TIR was created to solve problems in the narrow domain of competition-level mathematics. As a result, the model should not be used for general chat applications. With greedy decoding, we find the model is capable of solving problems at the level of AMC 12, but often struggles generate a valid solution on harder problems at the AIME and Math Olympiad level. The model also struggles to solve geometry problems, likely due to it's limited capacity and lack of other modalities like vision.

Training procedure

Training hyperparameters

The following hyperparameters were used during training:

• learning_rate: 2e-05
• train_batch_size: 1
• eval_batch_size: 1
• seed: 42
• distributed_type: multi-GPU
• num_devices: 32
• total_train_batch_size: 32
• total_eval_batch_size: 32
• optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
• lr_scheduler_type: cosine
• lr_scheduler_warmup_ratio: 0.1
• num_epochs: 4

Framework versions

• Transformers 4.42.3
• Pytorch 2.3.0+cu121
• Datasets 2.18.0
• Tokenizers 0.19.1

Citation

If you find NuminaMath 7B TIR is useful in your work, please cite it with:

@misc{numina_math_7b,
  author = {Edward Beeching and Shengyi Costa Huang and Albert Jiang and Jia Li and Benjamin Lipkin and Zihan Qina and Kashif Rasul and Ziju Shen and Roman Soletskyi and Lewis Tunstall},
  title = {NuminaMath 7B TIR},
  year = {2024},
  publisher = {Numina & Hugging Face},
  journal = {Hugging Face repository},
  howpublished = {\url{https://huggingface.co/AI-MO/NuminaMath-7B-TIR}}
}