Search-R1：强化学习训练大模型搜索与推理的工具

综合介绍

Search-R1 是一个开源项目，由 PeterGriffinJin 在 GitHub 上开发，基于 veRL 框架构建。它通过强化学习（RL）技术训练大语言模型（LLM），让模型自主学会推理和调用搜索引擎解决问题。项目支持 Qwen2.5-3B 和 Llama3.2-3B 等基础模型，扩展了 DeepSeek-R1 和 TinyZero 的方法。用户可以用它训练模型处理单轮或多轮任务，提供代码、数据集和实验日志。官方论文发布于 2025 年 3 月，项目模型和数据可在 Hugging Face 下载，适合研究人员和开发者使用。

功能列表

• 通过强化学习训练大模型，提升推理和搜索能力。
• 支持调用谷歌、必应、Brave 等搜索引擎 API。
• 提供 LoRA 调优和监督微调功能，优化模型表现。
• 内置现成重排序器，提高搜索结果准确性。
• 包含详细实验日志和论文，支持结果复现。
• 提供本地检索服务器功能，方便自定义搜索。
• 支持用户上传自定义数据集和语料库。

使用帮助

Search-R1 面向有编程和机器学习基础的用户。以下是详细安装和使用指南，让你快速上手。

安装流程

使用 Search-R1 需要先搭建环境。步骤如下：

1. 创建 Search-R1 环境
   在终端运行：

conda create -n searchr1 python=3.9
conda activate searchr1

这会创建一个 Python 3.9 的虚拟环境。

2. 安装 PyTorch
   输入以下命令安装 PyTorch 2.4.0（支持 CUDA 12.1）：

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

3. 安装 vLLM
   vLLM 是运行大模型的关键库，安装 0.6.3 版本：

pip3 install vllm==0.6.3

也可选择 0.5.4、0.4.2 或 0.3.1 版本。

4. 安装 veRL
   在项目根目录运行：

pip install -e .

这会安装 veRL 框架。

5. 安装可选依赖
   为提升性能，安装 Flash Attention 和 Wandb：

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

6. 安装检索器环境（可选）
   如果需要本地检索服务器，另建环境：

conda create -n retriever python=3.10
conda activate retriever
conda install pytorch==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install transformers datasets
conda install -c pytorch -c nvidia faiss-gpu=1.8.0
pip install uvicorn fastapi

快速开始

以下是基于 NQ 数据集训练模型的步骤：

1. 下载索引和语料
   设置保存路径并运行：

save_path=/你的保存路径
python scripts/download.py --save_path $save_path
cat $save_path/part_* > $save_path/e5_Flat.index
gzip -d $save_path/wiki-18.jsonl.gz

2. 处理 NQ 数据
   运行脚本生成训练数据：

python scripts/data_process/nq_search.py

3. 启动检索服务器
   在检索器环境运行：

conda activate retriever
bash retrieval_launch.sh

4. 运行 RL 训练
   在 Search-R1 环境运行：

conda activate searchr1
bash train_ppo.sh

这会用 Llama-3.2-3B 基础模型进行 PPO 训练。

使用自定义数据集

1. 准备 QA 数据
   数据需为 JSONL 格式，每行包含以下字段：

{
"data_source": "web",
"prompt": [{"role": "user", "content": "问题"}],
"ability": "fact-reasoning",
"reward_model": {"style": "rule", "ground_truth": "答案"},
"extra_info": {"split": "train", "index": 1}
}

参考 <scripts/data_process/nq_search.py>。

2. 准备语料库
   语料需为 JSONL 格式，每行包含 id 和 contents，如：

   {"id": "0", "contents": "文本内容"}
   

   可参考 <example/corpus.jsonl>。

3. 索引语料（可选）
   如果使用本地检索，运行：

   bash search_r1/search/build_index.sh
   

调用自定义搜索引擎

1. 修改 <search_r1/search/retriever_server.py>，配置 API。
2. 启动服务器：

   python search_r1/search/retriever_server.py
   

3. 模型会通过 http://127.0.0.1:8000/retrieve 调用搜索。

推理操作

1. 启动检索服务器：

   bash retrieval_launch.sh
   

2. 运行推理：

   python infer.py
   

3. 修改 <infer.py> 第 7 行的 question，输入你想问的问题。

注意事项

• 训练需要 GPU，至少 24GB 显存（如 NVIDIA A100）。
• 确保 API 密钥有效，网络连接稳定。
• 官方论文和实验日志（<Full experiment log 1> 和 <Full experiment log 2>）提供更多细节。

通过这些步骤，你可以用 Search-R1 训练一个能推理和搜索的模型，处理各种任务。

应用场景

1. 科研实验
   研究人员可以用 Search-R1 复现论文结果，探索强化学习在模型训练中的应用。
2. 智能助手开发
   开发者可以训练模型，集成到聊天工具中，提供搜索和推理功能。
3. 知识查询
   用户可以用它快速回答复杂问题，通过搜索获取最新信息。

QA

1. Search-R1 支持哪些模型？
   目前支持 Qwen2.5-3B 和 Llama3.2-3B 基础模型，其他模型需自行适配。
2. 训练需要多长时间？
   视数据集和硬件而定，在 24GB 显存 GPU 上，NQ 数据集训练约需数小时。
3. 如何验证训练效果？
   查看 <Preliminary results> 中的图表，或检查 Wandb 日志。