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

上次更新于：05/30/2025。

简介
----
- 我们已在 Multi-Turn RL 中添加了搜索工具调用功能，使模型能够在 Actor rollout 阶段发起检索请求，并直接使用检索结果进行训练。**我们支持使用本地密集检索器作为检索工具，也支持与您自己的本地检索引擎集成。**



快速重现
----------

创建新的 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}

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

.. 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

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

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

-  GPU 版本提供更高的准确性和速度；每个 GPU 使用约 5–7 GB 内存。
-  CPU 版本可用于简单测试，但检索精度较低，会降低训练性能。请参阅 search-R1 中的 `retriever 文档 <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

   # 创建并激活retriever环境，使用 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 内存，剩余部分用于 multi-turn RL 训练。

.. 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}

   # 定义时间戳函数
   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() 函数已定义
   # 创建日志目录
   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 &

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

要启用 multi-turn 推理，请在您的配置中设置以下字段：

.. 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

   输入格式：
   {
     "queries": ["What is Python?", "Tell me about neural networks."],
     "topk": 3,
     "return_scores": true
   }

   输出格式（当 return_scores=True 时，返回相似度得分）：
   {
       "result": [
           [   # 每个查询的结果
               {
                   "document": doc, "score": score
               },
               # ... 更多文档
           ],
           # ... 其他查询的结果
       ]
   }

注意事项
-------

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