=======================
搜索工具集成
=======================

最后更新：2025年5月30日。

简介
------------
- 我们为多轮强化学习（Multi-Turn RL）增加了搜索工具调用功能，使模型能够在 Actor roll out 过程中发起检索请求，并直接使用检索结果进行训练。**我们支持使用本地稠密检索器（local dense retriever）作为检索工具，同时也支持与您自己的本地检索引擎进行集成。**

快速复现
------------------

创建新的 Docker 容器
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code:: bash

   docker run \
       -it \
       --shm-size 32g \
       --gpus all \
       -v {Huggingface-Cache-Path}:/root/.cache \
       --ipc=host \
       --network=host \
       --privileged \
       --name sglang_{your-name} \
       lmsysorg/sglang:dev \
       /bin/zsh

如果退出容器后需要重启：

.. code:: bash

   docker start -i sglang_{your-name}

更新 Python 并使用 uv 配置虚拟环境
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. code:: bash

   apt update
   apt install -y python3.10 python3.10-venv

   # 创建虚拟环境
   python3 -m venv ~/.python/verl-multiturn-rollout

   # 激活虚拟环境
   source ~/.python/verl-multiturn-rollout/bin/activate

   # 安装 uv
   python3 -m pip install uv

安装 verl 上游代码
~~~~~~~~~~~~~~~~~~~~~

.. code:: bash

   cd ~
   git clone https://github.com/volcengine/verl.git
   cd verl

   # 安装 verl
   python3 -m uv pip install .
   python3 -m uv pip install -r ./requirements_sglang.txt

   # 手动安装 flash-attn
   python3 -m uv pip install wheel
   python3 -m uv pip install packaging
   python3 -m uv pip install flash-attn --no-build-isolation --no-deps

设置本地检索引擎
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

如果您使用自己的本地检索服务，可以跳过此步骤。我们选择了 search-R1 示例中提供的本地稠密检索器；详细说明请参见 `searchR1 文档 <https://raw.githubusercontent.com/PeterGriffinJin/Search-R1/refs/heads/main/docs/retriever.md>`__。简而言之：

-  GPU 版本检索器提供更高的准确性和速度；每个 GPU 使用约 5-7 GB 内存。
-  CPU 版本可用于简单测试，但检索精度较低，这将影响训练性能。有关详细信息，请参阅 search-R1 中的 `检索器文档 <https://github.com/PeterGriffinJin/Search-R1/blob/main/docs/retriever.md>`__。
-  建议使用 Conda 安装 faiss-gpu=1.8.0；venv 可能会导致错误。

**注意**：启动训练进程和本地检索服务时，我们会启动两个独立的 Python 环境。训练使用 verl-multiturn-rollout 环境中的 uv，而检索器使用 conda 安装 ``faiss-gpu``。

.. code:: bash

   # 下载 Miniconda 安装脚本
   wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh

   # 以批处理模式安装到 $HOME/miniconda3
   bash ~/miniconda.sh -b -p $HOME/miniconda3

   # 激活 conda（仅在当前shell中）
   eval "$($HOME/miniconda3/bin/conda shell.bash hook)"

   # （可选）将 conda 添加到您的默认 shell 启动配置
   conda init

   # 重新加载 shell 配置
   source ~/.bashrc

   # 创建并激活带 Python 3.10 的检索器环境
   conda create -n retriever python=3.10 -y
   conda activate retriever

   # 安装 PyTorch（支持 GPU）及相关库
   conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia -y

   # 安装其他 Python 包
   pip install transformers datasets pyserini huggingface_hub

   # 安装 GPU 版本的 faiss
   conda install faiss-gpu=1.8.0 -c pytorch -c nvidia -y

   # 安装 API 服务框架
   pip install uvicorn fastapi

下载索引和语料库
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

本地检索文件很大，请准备足够的磁盘空间。下载大约需要 60-70 GB，解压后大约需要 132 GB：

.. code:: bash

   conda activate retriever

   save_path=/the/path/to/save  # 请替换为您实际的保存路径
   python examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py --save_path $save_path
   cat $save_path/part_* > $save_path/e5_Flat.index
   gzip -d $save_path/wiki-18.jsonl.gz

启动本地 flat e5 检索服务器
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

1. 首次启动将下载模型并加载索引。
2. 除下载外，启动大约需要 1-2 分钟。
3. 启动后，每个 GPU 使用约 5-7 GB 内存，剩余内存可供多轮强化学习训练使用。

.. code:: bash

   conda activate retriever

   index_file=$save_path/e5_Flat.index
   corpus_file=$save_path/wiki-18.jsonl
   retriever_name=e5
   retriever_path=intfloat/e5-base-v2

   python examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py \
     --index_path $index_file \
     --corpus_path $corpus_file \
     --topk 3 \
     --retriever_name $retriever_name \
     --retriever_model $retriever_path \
     --faiss_gpu

设置 WANDB_API_KEY
~~~~~~~~~~~~~~~~~~~~

.. code:: bash

   export WANDB_API_KEY={YOUR_WANDB_API_KEY}  # 请替换为您的 WANDB API 密钥

   # 定义一个时间戳函数
   function now() {
       date '+%Y-%m-%d-%H-%M'
   }

**预处理数据集**
~~~~~~~~~~~~~~~~~~~~~~~~~~

   **注意：** 以下数据处理和训练命令必须在 verl-multiturn-rollout 环境中运行。

.. code:: bash

   python3 examples/data_preprocess/preprocess_search_r1_dataset.py

在 8 x H20 上进行测试
~~~~~~~~~~~~~~~~~~

.. code:: bash

   # 确保 now() 函数已定义
   # 创建 logs 目录
   mkdir -p logs

   # 设置 GPU 并使用合适的日志路径运行
   export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

   nohup bash examples/sglang_multiturn/search_r1_like/run_qwen2.5-3b_instruct_search_multiturn.sh \
     trainer.experiment_name=qwen2.5-3b-it_rm-searchR1-like-sgl-multiturn-$(now) \
     > logs/searchR1-like$(now).log 2>&1 &

自定义搜索配置
---------------------------

要启用多轮推理，请在您的配置文件中设置以下字段：

.. code:: yaml

   actor_rollout_ref:
     rollout:
       name: "sglang"
       multi_turn:
         enable: True

您必须在 ``examples/sglang_multiturn/config/tool_config/search_tool_config.yaml`` 中指定 ``retrieval_service_url``，并正确配置并发。有关并发的更多详细信息，请参阅 Sandbox Fusion 示例：

.. code:: yaml

   tools:
     - class_name: verl.tools.search_tool.SearchTool
       config:
         retrieval_service_url: http://127.0.0.1:8000/retrieve
         num_workers: 120
         rate_limit: 120
         timeout: 30

检索器的输入/输出格式如下。如果您的服务参数匹配，只需修改 ``retrieval_service_url``。您也可以在 ``search_r1_like_utils.py`` 中进行自定义。

.. code:: python

   Input format:
   {
     "queries": ["What is Python?", "Tell me about neural networks."],
     "topk": 3,
     "return_scores": true
   }

   Output format (when return_scores=True, similarity scores are returned):
   {
       "result": [
           [   # Results for each query
               {
                   "document": doc, "score": score
               },
               # ... more documents
           ],
           # ... other queries' results
       ]
   }

备注
-----

1. 总训练时间约 27 小时；同时，验证数据集非常大（51k），每次验证耗时约 6000 秒。（因此，默认情况下 ``val_before_train=False``）。