搜索工具集成

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

简介

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

快速重现

创建新的 Docker 容器

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

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

docker start -i sglang_{your-name}

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

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 上游版本

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 文档。 简要来说：

  • GPU 版本提供更高的准确性和速度；每个 GPU 使用约 5–7 GB 内存。

  • CPU 版本可用于简单测试，但检索精度较低，会降低训练性能。请参阅 search-R1 中的 retriever 文档 以获取详细信息。

  • 推荐使用 Conda 安装 faiss-gpu=1.8.0；venv 可能导致错误。

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

# 下载 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：

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 训练。

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

export WANDB_API_KEY={YOUR_WANDB_API_KEY}

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

预处理数据集

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

python3 examples/data_preprocess/preprocess_search_r1_dataset.py

在 8 x H20 上测试

# 确保 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 推理，请在您的配置中设置以下字段：

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 示例：

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 中自定义。

输入格式：
{
  "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）