Skip to content

README.md

Latest commit

 

History

History
133 lines (103 loc) · 4.46 KB

README.md

File metadata and controls

133 lines (103 loc) · 4.46 KB
Raw

Search-R1: Train your LLMs to reason and call a search engine with reinforcement learning

Search-R1 is a reinforcement learning framework designed for training reasoning-and-searching interleaved LLMs—language models that learn to reason and make tool calls (e.g., to search engines) in a coordinated manner.

Built upon veRL, Search-R1 extends the ideas of DeepSeek-R1(-Zero) by incorporating interleaved search engine access and provides a fully open-source RL training pipeline. It serves as an alternative and open solution to OpenAI DeepResearch, enabling research and development in tool-augmented LLM reasoning.

EMA-PG: Search Engine Experiments

Run train_grpo_ema_pg.sh to use EMA Policy Gradient on Qwen-3B base.

To use EMA Anchor:

    +algorithm.ref_policy_ema_tau=0.9 \
    +algorithm.ref_policy_ema_update_period=10 \

To use Top-k KL:

    actor_rollout_ref.actor.kl_loss_type=full_reverse \
    actor_rollout_ref.actor.kl_topk_tokens=32 \
    actor_rollout_ref.actor.use_kl_iw=true \

Links

Installation

Search-r1 environment

conda create -n searchr1 python=3.9
conda activate searchr1
# install torch [or you can skip this step and let vllm to install the correct version for you]
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
# install vllm
pip3 install vllm==0.6.3 # or you can install 0.5.4, 0.4.2 and 0.3.1

# verl
pip install -e .

# flash attention 2
pip3 install flash-attn --no-build-isolation
pip install wandb

Retriever environment (optional)

If you would like to call a local retriever as the search engine, you can install the environment as follows. (We recommend using a seperate environment.)

conda create -n retriever python=3.10
conda activate retriever

# we recommend installing torch with conda for faiss-gpu
conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia
pip install transformers datasets pyserini

## install the gpu version faiss to guarantee efficient RL rollout
conda install -c pytorch -c nvidia faiss-gpu=1.8.0

## API function
pip install uvicorn fastapi

Quick start

Train a reasoning + search LLM on NQ dataset with e5 as the retriever and wikipedia as the corpus.

(1) Download the indexing and corpus.

save_path=/the/path/to/save
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) Process the NQ dataset.

python scripts/data_process/nq_search.py

(3) Launch a local retrieval server.

conda activate retriever
bash retrieval_launch.sh

(4) Run RL training (PPO) with Llama-3.2-3b-base.

conda activate searchr1
bash train_ppo.sh

Inference

You can play with the trained Search-R1 model with your own question.

(1) Launch a local retrieval server.

conda activate retriever
bash retrieval_launch.sh

(2) Run inference.

conda activate searchr1
python infer.py

You can modify the question on line 7 to something you're interested in.

Acknowledge

The concept of Search-R1 is inspired by Deepseek-R1 and TinyZero. Its implementation is built upon veRL and RAGEN. We sincerely appreciate the efforts of these teams for their contributions to open-source research and development.

Citations

@article{jin2025search,
  title={Search-r1: Training llms to reason and leverage search engines with reinforcement learning},
  author={Jin, Bowen and Zeng, Hansi and Yue, Zhenrui and Yoon, Jinsung and Arik, Sercan and Wang, Dong and Zamani, Hamed and Han, Jiawei},
  journal={arXiv preprint arXiv:2503.09516},
  year={2025}
}