TinyZero

⚠️ 弃用通知： 此仓库不再积极维护。如需运行强化学习实验，请直接使用最新的 veRL 库。
存档的原始文档请参见 OLD_README.md。

TinyZero 是对 DeepSeek R1 Zero 在倒计时和乘法任务上的复现。我们基于 veRL 构建。

通过强化学习，30亿参数的基础语言模型能够自主发展出自我验证和搜索能力。

你可以用不到 30 美元的成本亲自体验这个“顿悟”时刻。

Twitter 讨论串：https://x.com/jiayi_pirate/status/1882839370505621655

完整实验日志：https://wandb.ai/jiayipan/TinyZero

📢：我们发布了 自适应并行推理，探索了扩展推理模型的新维度。

安装

conda create -n zero python=3.9
# 安装 torch [你也可以跳过此步，让 vllm 为你安装正确的版本]
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
# 安装 vllm
pip3 install vllm==0.6.3 # 你也可以安装 0.5.4、0.4.2 或 0.3.1
pip3 install ray

# verl
pip install -e .

# flash attention 2
pip3 install flash-attn --no-build-isolation
# 提升体验的工具
pip install wandb IPython matplotlib

倒计时任务

数据准备

conda activate zero
python ./examples/data_preprocess/countdown.py --local_dir {你的数据集路径}

运行训练

conda activate zero

对于以下代码，如果遇到显存不足（VRAM）的问题，可以尝试在脚本中添加 critic.model.enable_gradient_checkpointing=True，并查看此处的讨论。

单 GPU

适用于模型 <= 1.5B。对于 Qwen2.5-0.5B 基础模型，已知其无法学会推理。

export N_GPUS=1
export BASE_MODEL={你的模型路径}
export DATA_DIR={你的数据集路径}
export ROLLOUT_TP_SIZE=1
export EXPERIMENT_NAME=countdown-qwen2.5-0.5b
export VLLM_ATTENTION_BACKEND=XFORMERS

bash ./scripts/train_tiny_zero.sh

3B+ 模型
在这种情况下，基础模型能够发展出复杂的推理技能。

export N_GPUS=2
export BASE_MODEL={你的模型路径}
export DATA_DIR={你的数据集路径}
export ROLLOUT_TP_SIZE=2
export EXPERIMENT_NAME=countdown-qwen2.5-3b
export VLLM_ATTENTION_BACKEND=XFORMERS

bash ./scripts/train_tiny_zero.sh

指令模型消融实验

我们也对 Qwen-2.5-3B Instruct 模型进行了实验。
数据准备
为了遵循聊天模板，我们需要重新处理数据：

conda activate zero
python examples/data_preprocess/countdown.py --template_type=qwen-instruct --local_dir={你的数据集路径}

训练

export N_GPUS=2
export BASE_MODEL={你的模型路径}
export DATA_DIR={你的数据集路径}
export ROLLOUT_TP_SIZE=2
export EXPERIMENT_NAME=countdown-qwen2.5-3b-instruct
export VLLM_ATTENTION_BACKEND=XFORMERS

bash ./scripts/train_tiny_zero.sh

致谢

• 我们的实验基于 veRL 运行。
• 我们使用了 Qwen2.5 系列基础模型 Qwen2.5。

引用

@misc{tinyzero,
author       = {Jiayi Pan and Junjie Zhang and Xingyao Wang and Lifan Yuan and Hao Peng and Alane Suhr},
title        = {TinyZero},
howpublished = {https://github.com/Jiayi-Pan/TinyZero},
note         = {Accessed: 2025-01-24},
year         = {2025}
}