nanoAhaMoment：单文件“LLM强化学习”库

阿米尔侯赛因·卡泽姆内贾德*，米拉德·阿加乔哈里*，亚历山德罗·索尔多尼，阿伦·库维尔，希瓦·雷迪

采用以下方式实现 DeepSeek R1-zero 风格的训练：

• 单张 80G 显卡（也支持多 GPU）
• 无需任何强化学习库
• 3B 基础模型（多 GPU 下也可支持 7B 模型）
• 全参数微调
• 高效（性能可与 VerL 竞争，但实现更简单）
• 多 GPU 下 3B 模型最高支持 32K 上下文长度（7B 模型为 16K）

新闻

• 2025 年 6 月：新增多 GPU 支持，以加快训练速度，并支持 7B 模型
• 2025 年 6 月：新增 VinePPO episode 生成功能（实验性）

受 TinyZero 和 Mini-R1 启发，但设计得更加 简单、整洁 和 高效，所有代码行都清晰可见且易于理解。

卡帕西风格的 YouTube 详细讲解

• nanoAhaMoment：从零开始，仅用一张 GPU 实现 LLM 强化学习 - 第 1 部分
• nanoAhaMoment：从零开始，仅用一张 GPU 实现 LLM 强化学习 - 第 2 部分

文件说明

• nano_r1.ipynb 是带有教程的交互式单文件 Jupyter 笔记本。
• nano_r1_script.py 实际上就是 nano_r1.ipynb，但为了方便使用 Python 运行以及支持多 GPU 而单独提供。
• notebooks/checkpoint_playground.ipynb 是一个用于比较不同模型检查点（包括我们训练的模型）并进行交互的笔记本。
• 🤗 McGill-NLP/nano-aha-moment-3b：使用上述脚本训练的 Hugging Face 模型（在 CountDown 任务上的准确率约为 60%）。

设置说明

1. 克隆仓库

   git clone https://github.com/McGill-NLP/nano-aha-moment.git
   

2. 安装依赖
   首先确保已安装 CUDA 12.4。

   安装 PyTorch：

   pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
   

   安装其余依赖：

   pip install -r requirements.txt
   

   使用 uv 的替代安装方法（可选）

   uv sync
   uv sync --extra compile  # 安装 flash-attention
   

3. 运行训练脚本
   打开 nano_r1.ipynb 或 nano_r1_script.py，即可开始训练。

   如果使用 uv，可以运行 uv run nano_r1_script.py，或者激活虚拟环境后执行 python nano_r1_script.py。

多 GPU 训练

以下是使用 4 张 GPU 运行训练脚本的命令：

python nano_r1_script.py --nproc 4  # 使用 4 张 GPU

不同上下文长度下的批量大小

上下文长度	3B 模型 (per_device_batch_size)	7B 模型 (per_device_batch_size)
1024	32	16
2048	16	8
4K	8	4
8K	4	2
16K	2	1
32K	1	无

注意：这些批量大小是针对 4 张 A100 80GB 显卡优化的。对于其他类型的显卡，可能需要相应调整批量大小。

待办事项

• 完整的评估套件
• 多 GPU 支持（2025 年 6 月添加）

致谢

我们衷心感谢 Lambda 公司通过其研究计算资助计划提供的算力支持。

引用

如果您在研究中使用了本代码库，请使用以下格式引用我们：

@misc{Kazemnejad2025:NanoAhaMoment,
  author       = {Amirhossein Kazemnejad and Milad Aghajohari and Alessandro Sordoni and Aaron Courville and Siva Reddy},
  title        = {Nano Aha! Moment: 单文件“LLM 强化学习”库},
  year         = {2025},
  howpublished = {\url{https://github.com/McGill-NLP/nano-aha-moment}},
  note         = {GitHub 仓库}
}

nano-aha-moment 快速上手指南

nanoAhaMoment 是一个单文件实现的"LLM 强化学习（RL）”库，旨在复现 DeepSeek R1-zero 风格的训练。它无需依赖复杂的 RL 库，支持单卡（80G）或多卡训练，代码简洁透明，适合学习和实验。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐)
• GPU:
  • 单卡模式：至少 1 张 80G 显存 GPU (如 A100/H100)。
  • 多卡模式：支持多卡并行，可训练 7B 模型或处理更长上下文（最高 32K）。
• CUDA 版本: 必须安装 CUDA 12.4。
• Python 环境: 建议创建独立的虚拟环境。

安装步骤

1. 克隆项目

git clone https://github.com/McGill-NLP/nano-aha-moment.git
cd nano-aha-moment

2. 安装依赖

方式 A：使用 pip (标准安装)

首先安装指定版本的 PyTorch (对应 CUDA 12.4)：

pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124

国内加速建议：如果下载速度慢，可使用清华或阿里镜像源替代官方源：

pip install torch==2.5.1 --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple
# 注意：需确保下载的 wheel 包兼容 cu124，若镜像源无特定版本，建议使用官方源或手动下载 whl 文件。

接着安装其他依赖：

pip install -r requirements.txt

方式 B：使用 uv (可选，推荐用于快速构建)

如果您安装了 uv 工具，可以使用以下命令同步环境并编译 Flash Attention：

uv sync
uv sync --extra compile  # 安装 flash-attention

基本使用

本项目提供了 Jupyter Notebook (nano_r1.ipynb) 和 Python 脚本 (nano_r1_script.py) 两种运行方式，功能一致。

1. 单卡训练 (最简单示例)

直接运行脚本即可开始训练（默认使用单卡）：

python nano_r1_script.py

或者，如果您使用 uv 管理环境：

uv run nano_r1_script.py

2. 多卡训练

如需使用多张 GPU 加速训练或运行更大模型（如 7B），通过 --nproc 参数指定 GPU 数量。例如，使用 4 张 GPU：

python nano_r1_script.py --nproc 4

3. 调整 Batch Size

根据您使用的模型大小（3B 或 7B）和上下文长度，可能需要调整 nano_r1_script.py 中的 per_device_batch_size 以避免显存溢出。参考配置如下（基于 4xA100 80G 优化）：

上下文长度	3B 模型建议 Batch Size	7B 模型建议 Batch Size
1024	32	16
2048	16	8
4K	8	4
8K	4	2
16K	2	1
32K	1	N/A

修改代码中的对应参数后，重新运行上述命令即可。