TinyZero：低成本复现 DeepSeeK-R1 Zero 的顿悟效果

综合介绍

TinyZero 是一个基于 veRL 的强化学习模型，旨在复现 DeepSeeK-R1 Zero 在倒计时和乘法任务中的表现。令人惊讶的是，该项目仅需 30 美元的运行成本（使用 2xH200，每小时 6.4 美元，不到 5 小时），就能实现与 DeepSeeK-R1 Zero 相同的顿悟效果。通过强化学习（RL），3B 基础语言模型（LM）能够自主开发自我验证和搜索能力。用户可以通过简单的安装步骤和训练过程，体验到 TinyZero 的强大功能和创新性。

功能列表

• 倒计时任务：支持数据准备和训练过程，帮助模型在倒计时任务中进行学习。
• 乘法任务：支持数据准备和训练过程，帮助模型在乘法任务中进行学习。
• 单 GPU 支持：适用于模型参数小于等于 1.5B 的情况。
• 多 GPU 支持：适用于更大参数的模型，能够开发复杂的推理能力。
• Instruct Ablation：支持 QWen-2.5-3B Instruct 模型的实验。
• 质量提升工具：包括 flash-attn、wandb、IPython 和 matplotlib 等工具，提升模型训练和使用体验。

使用帮助

安装流程

1. 创建虚拟环境：

   conda create -n zero python=3.9
   

2. 安装 PyTorch（可选）：

   pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
   

3. 安装 vllm：

   pip3 install vllm==0.6.3
   

4. 安装 ray：

   pip3 install ray
   

5. 安装 verl：

   pip install -e .
   

6. 安装 flash-attn：

   pip3 install flash-attn --no-build-isolation
   

7. 安装质量提升工具：

   pip install wandb IPython matplotlib
   

功能操作流程

倒计时任务

1. 数据准备：

   conda activate zero
   python ./examples/data_preprocess/countdown.py --local_dir {path_to_your_dataset}
   

2. 训练过程：

   conda activate zero
   export N_GPUS=1
   export BASE_MODEL={path_to_your_model}
   export DATA_DIR={path_to_your_dataset}
   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+ 模型训练

1. 数据准备：

   conda activate zero
   python examples/data_preprocess/countdown.py --template_type=qwen-instruct --local_dir={path_to_your_dataset}
   

2. 训练过程：

   conda activate zero
   export N_GPUS=2
   export BASE_MODEL={path_to_your_model}
   export DATA_DIR={path_to_your_dataset}
   export ROLLOUT_TP_SIZE=2
   export EXPERIMENT_NAME=countdown-qwen2.5-3b-instruct
   export VLLM_ATTENTION_BACKEND=XFORMERS
   bash ./scripts/train_tiny_zero.sh