aimo-progress-prize

<div align="center"> <img src="https://yellow-cdn.veclightyear.com/835a84d5/0e98e0e9-bafe-4d78-b689-b38c12a1201b.png" width="60%" alt="Numina" /> </div> <p align="center"> 🌐 <a href="https://www.projectnumina.ai/">网站</a>&nbsp&nbsp | 🗂️ <a href="https://huggingface.co/collections/AI-MO/numinamath-6697df380293bcfdbc1d978c">数据集</a>&nbsp&nbsp | &nbsp&nbsp 📑 <a href="https://github.com/project-numina/aimo-progress-prize/blob/main/report/numina_dataset.pdf">数据集技术报告</a> &nbsp&nbsp <br> 🤗 <a href="https://huggingface.co/AI-MO">Hugging Face</a>&nbsp&nbsp | &nbsp&nbsp🤗 <a href="https://huggingface.co/blog/winning-aimo-progress-prize">博客文章</a>&nbsp&nbsp | &nbsp&nbsp🖥️ <a href="https://huggingface.co/spaces/AI-MO/math-olympiad-solver">演示</a> </p> <br><br>

Numina 对第一届 AIMO 进展奖的解决方案

本仓库包含了复现我们在AI 数学奥林匹克 - 进展奖 1中获胜解决方案所需的训练和推理代码。

我们的解决方案包括四个主要组成部分：

一个用于微调 DeepSeekMath-Base 7B 的配方，使其能够使用工具集成推理（TIR）解决数学问题
两个高质量的训练数据集，包含约 100 万个数学问题及其解答
一个自洽解码算法，用于生成带有代码执行反馈的解答候选（SC-TIR）
四个精心选择的验证集，来自 AMC、AIME 和 MATH，用于指导模型选择并避免过拟合公开排行榜

我们使用了一系列开源库来训练我们的模型，主要包括 TRL、PyTorch、vLLM 和 DeepSpeed。在一个包含 8 个 H100 GPU 的节点上，我们的模型训练时间为 10 小时。您可以在 Hugging Face Hub 的 NuminaMath 集合中找到这些数据集和模型。

如何浏览这个项目

这个项目设计简洁，主要包括：

training：用于微调和量化模型的脚本。
kaggle-solution：一个包含我们在 Kaggle 提交过程中使用的 SC-TIR 推理代码的笔记本。

安装说明

要运行本项目中的代码，首先使用例如 Conda 创建一个 Python 虚拟环境：

conda create -n aimo python=3.10 && conda activate aimo

接下来，安装 PyTorch v2.1.2 - 精确的版本对于可重现性很重要！由于这依赖于硬件，我们建议您参考 PyTorch 安装页面。

然后，您可以按如下方式安装剩余的包依赖：

pip install -r requirements.txt

您还需要安装 Flash Attention 2，可以通过运行以下命令完成：

python -m pip install flash-attn --no-build-isolation

注意：如果您的机器 RAM 少于 96GB 且 CPU 核心较多，请减少 MAX_JOBS 参数，例如 MAX_JOBS=4 pip install flash-attn --no-build-isolation

接下来，按如下方式登录您的 Hugging Face 账户：

huggingface-cli login

最后，安装 Git LFS，以便您可以将模型推送到 Hugging Face Hub：

sudo apt-get install git-lfs

训练方法

训练分为两个阶段，遵循 MuMath-Code 配方：

第一阶段 在数学问题和文本解答上进行思维链（CoT）训练。
第二阶段 在数学问题和代码解答上进行工具集成推理（TIR）训练。

数据集

我们使用了两个数据集来微调我们的模型：

NuminaMath-CoT： 约86万道数学题，每个解答都以思维链（CoT）的方式格式化。数据集来源范围从中国高中数学练习到美国和国际数学奥林匹克竞赛题。数据主要从在线试卷PDF和数学讨论论坛收集。处理步骤包括：(a)从原始PDF进行OCR，(b)分割成问题-解答对，(c)翻译成英语，(d)重新调整以产生CoT推理格式，(e)最终答案格式化。
NuminaMath-TIR： 工具集成推理（TIR）在本次竞赛中起着关键作用。然而，收集和注释此类数据既昂贵又耗时。为解决这个问题，我们从NuminaMath-CoT数据集中选择了约7万道问题，重点关注那些有数值输出的问题，其中大部分是整数。然后我们利用GPT-4构建了一个管道来生成类似TORA的推理路径，执行代码并产生结果，直到解答完成。我们过滤掉最终答案与参考不匹配的解答，并重复这个过程三次以确保准确性和一致性。这种迭代方法使我们能够高效地生成高质量的TORA数据。

模型训练

如上所述，模型训练分两个步骤进行：

对DeepSeekMath-Base 7B应用SFT，在NuminaMath-CoT数据集上进行微调。结果是一个SFT模型，如NuminaMath-7B-CoT。
在NuminaMath-TIR数据集上微调第1阶段的SFT模型，以学习工具集成推理。结果是一个"推理代理"，如NuminaMath-7B-TIR。

以下是使用DeepSpeed ZeRO-3训练这些模型的命令。请注意，你需要8个GPU，每个有80GB的显存，才能使用我们的设置训练完整模型。

第1阶段

accelerate launch --config_file=training/configs/deepspeed_zero3.yaml training/sft.py training/configs/stage-1-cot.yaml

第2阶段

accelerate launch --config_file=training/configs/deepspeed_zero3.yaml training/sft.py training/configs/stage-2-tir.yaml

训练后量化

模型训练完成后，我们使用AutoGPTQ将其量化为8位精度，以提高在Kaggle的T4 GPU上使用vLLM的性能（因为它们不支持bfloat16类型）。这一步是可选的，当以16位精度加载时，模型应该有更好的性能：

python training/quantization.py --model_id AI-MO/NuminaMath-7B-TIR --calibration_dataset data/NuminaMath-TIR

项目结构

aimo-progress-prize/
├── LICENSE
├── README.md               <- 开发者使用本项目的顶级README
├── images
├── kaggle-solution.ipynb   <- 我们Kaggle提交中使用的推理代码笔记本
├── requirements.txt        <- 项目依赖
└── training
    ├── configs             <- 训练的超参数配置
    ├── numina              <- 本项目使用的源代码
    ├── quantization.py     <- 使用AutoGPTQ量化模型的脚本
    └── sft.py              <- 微调我们模型的脚本

硬件（以下规格用于创建原始解决方案）

操作系统

NAME="Ubuntu"
VERSION="20.04.6 LTS (Focal Fossa)"
ID=ubuntu
ID_LIKE=debian
PRETTY_NAME="Ubuntu 20.04.6 LTS"
VERSION_ID="20.04"
HOME_URL="https://www.ubuntu.com/"
SUPPORT_URL="https://help.ubuntu.com/"
BUG_REPORT_URL="https://bugs.launchpad.net/ubuntu/"
PRIVACY_POLICY_URL="https://www.ubuntu.com/legal/terms-and-policies/privacy-policy"
VERSION_CODENAME=focal
UBUNTU_CODENAME=focal

GPU

一个节点，包含8个H100，每个有80GB显存，96个CPU，1TB内存。

软件（Python包详见`requirements.txt`）：

Python 3.10.14
CUDA 12.2
Nvidia驱动 v.535.104.12

引用

@misc{numina_math_datasets,
  author = {Jia LI, Edward Beeching, Lewis Tunstall, Ben Lipkin, Roman Soletskyi, Shengyi Costa Huang, Kashif Rasul, Longhui Yu, Albert Jiang, Ziju Shen, Zihan Qin, Bin Dong, Li Zhou, Yann Fleureau, Guillaume Lample, and Stanislas Polu},
  title = {NuminaMath},
  year = {2024},
  publisher = {Numina},
  journal = {GitHub repository},
  howpublished = {\url{[https://github.com/project-numina/aimo-progress-prize](https://github.com/project-numina/aimo-progress-prize/blob/main/report/numina_dataset.pdf)}}
}