DOCSEARCHbyAlgolia
Submit Docs

ZeroSearch Documentation

repository·main·Indexed 23 days ago

https://github.com/alibaba-nlp/zerosearch

A reinforcement learning framework designed to incentivize LLMs to develop search capabilities using simulated searches during training. It supports RL training via REINFORCE, GRPO, and PPO algorithms and integrates with the verl (Volcano Engine Reinforcement Learning) library for model optimization, parallelism (TP, DP, PP), and standardized data exchange using DataProto.

Tokens
2.9K
Snippets
4
Records
15
Agent score
79%

What's inside ZeroSearch

  1. Understand the RLHF dataset format

    main

    The RLHF datasets are stored in Parquet files. Data is organized using a chat-based format within the prompt field to support multi-turn conversations. To facilitate answer extraction, instruction-following text is often included directly in the prompt (e.g., asking the model to output the final answer after a specific delimiter like ####).

    Key fields in the dataset schema include:

    • data_source: The origin of the data (e.g., openai/gsm8k).
    • prompt: A list of message objects containing role and content (standard chat format).
    • ability: The capability being tested (e.g., math).
    • reward_model: A configuration object defining how the model is evaluated. It includes:
      • style: The reward calculation method (e.g., rule).
      • ground_truth: A list of expected correct answers.
    {
    "data_source": "openai/gsm8k",
    "prompt": [{"role": "user", "content": "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Let's think step by step and output the final answer after \"####\""}],
    "ability": "math",
    "reward_model": {
        "style": "rule",
        "ground_truth": ["72"]
    }
}

  2. Conduct RL training with ZeroSearch

    main

    Reinforcement Learning (RL) training can be performed using REINFORCE, GRPO, or PPO algorithms.

    Before training, ensure you have activated your conda environment and exported your SER_API_KEY.

    There are two simulation modes:

    1. simulate_prompt: Uses a prompt-based simulation model.
    2. simulate_sft: Uses a fine-tuning-based simulation model.

    Key parameters for the training scripts:

    • NUM_GPUS_PER_NODE: Number of GPUs available.
    • MODEL_PATH: The path to the model being trained (e.g., Qwen2.5-3B-Instruct).
    • DATA_PATH: Path to the ZeroSearch_dataset.
    • SEARCH_MODE: Either simulate_prompt or simulate_sft.
    • SIMULATION_LLM: The model used for simulation.
    • START_THRESHOLD / END_THRESHOLD: Define the difficulty levels for the curriculum rollout.
    # Activate the conda environment
conda activate zerosearch

# Set your Google Search API key
export SER_API_KEY=your_api_key

## Example: Prompt-based simulation with GRPO
bash train_grpo.sh NUM_GPUS_PER_NODE 4 MODEL_PATH Qwen2.5-3B-Instruct DATA_PATH ZeroSearch_dataset TOTAL_STEPS 203 IP localhost SEARCH_MODE simulate_prompt SIMULATION_LLM Qwen2.5-14B-Instruct START_THRESHOLD 0 END_THRESHOLD 0.5 SEARCH_ENGINE google MAX_TURNS 5 TOPK 5

## Example: Fine-tuning-based simulation with PPO
bash train_ppo.sh NUM_GPUS_PER_NODE 4 MODEL_PATH Qwen2.5-3B-Instruct DATA_PATH ZeroSearch_dataset TOTAL_STEPS 203 IP localhost SEARCH_MODE simulate_sft SIMULATION_LLM Simulation_LLM_google_14B START_THRESHOLD 0 END_THRESHOLD 0.5 SEARCH_ENGINE google MAX_TURNS 5 TOPK 5

  3. Launch a local simulation server using sglang

    main

    ZeroSearch uses sglang to host simulation models. You can launch a server using either a prompt-based approach (using a standard Instruct model) or a fine-tuning-based approach (using a specialized Simulation LLM).

    # Prompt-based simulation
python -m sglang.launch_server --model-path Qwen2.5-14B-Instruct --host 0.0.0.0 --tp 2 --dp 2 --port 6001

# Fine-tuning-based simulation
python -m sglang.launch_server --model-path Simulation_LLM_google_14B --host 0.0.0.0 --tp 2 --dp 2 --port 6001

  4. How to add a new Huggingface model to verl

    main

    To integrate a new Huggingface model into verl, you must port the model file to the verl/models/hf directory and optimize it for the verl design principles (parallelizable, highly-optimized, and using packed inputs).

    Follow these steps:

    1. Copy the model file: Create a new file under verl/models/hf and copy only the specific model file from the huggingface/transformers/models directory.
    2. Optimize for packed inputs:
      • Remove all inference-related code (such as KV cache management).
      • Update the model's input signature to accept only:
        • input_ids (shape: (total_nnz,))
        • cu_seqlens (shape: (total_nnz + 1,))
        • max_seqlen_in_batch (type: int)
      • Note: This optimization requires using Flash Attention with a causal mask.
    3. Verify with tests: Add a test in tests/models/hf to compare the output of your new verl model version against the original Huggingface version to ensure correctness.
    4. Implement Parallelism: Implement functions to support Tensor Parallelism (TP), Data Parallelism (DP), and Pipeline Parallelism (PP) as detailed below.

  6. Download ZeroSearch datasets and simulation LLMs

    main

    ZeroSearch requires specific datasets and simulation LLMs to function. Use huggingface-cli to download the main training dataset and (optionally) the Simulation Tuning dataset if you intend to train your own simulation LLMs.

    For simulation LLMs, you can choose different parameter sizes; the 14B version is recommended for stable and reliable performance.

  7. Implement Tensor Parallelism for verl models

    main

    When adding Tensor Parallelism (TP) to a model in verl, note that native PyTorch TP is not automatic. You must specify how model parameters and inputs/outputs are resharded using configurations. These configurations are then registered as hooks to perform the resharding before and after the model's forward pass.

    Refer to the following PyTorch documentation for implementation details:

  8. Install ZeroSearch dependencies

    main

    To set up the ZeroSearch environment, create a new Conda environment with Python 3.9 and install the required packages including torch, vllm, wandb, and serpapi. You must also install the local project in editable mode, flash-attn with no build isolation, and sglang.

    Note: If you encounter package conflicts when installing sglang, it is recommended to create a separate new environment specifically for sglang.

    conda create -n zerosearch python=3.9
conda activate zerosearch
pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cu121
pip install vllm==0.6.3
pip install wandb
pip install serpapi

# verl
pip install -e .

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

# sglang
pip install sglang[all]

  10. Use DataProto for standardized data exchange

    main

    The DataProto class is the core data structure for exchanging data between functions and modules in verl. It provides a unified protocol containing:

    • batch: A TensorDict containing tensors that share the same batch dimension.
    • non_tensor_batch: A dictionary of numpy.ndarray (with dtype=object) for non-tensor data.
    • meta_info: A dictionary for metadata.

    Common operations include moving to a device, selecting subsets of keys, and concatenating multiple DataProto objects.

  11. Handle asynchronous data with DataProtoFuture

    main

    For asynchronous execution (e.g., using Ray), DataProtoFuture allows the driver to avoid waiting for actual data fetching. It holds ray.ObjectRef futures and defines how to collect and dispatch them once they are ready.

    • DataProtoFuture.concat(data: List[ray.ObjectRef]): Creates a future that will eventually collect and concatenate a list of Ray object references into a single DataProto using DataProto.concat.
    • get(): Blocks until the futures are resolved, then applies the collect_fn and dispatch_fn to return the final DataProto.

  12. Manipulate DataProto: select, pop, rename, and union

    main

    The DataProto class provides several methods to transform or subset the data:

    • select(batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False): Returns a new DataProto containing only the specified keys. Use deepcopy=True to ensure the new object doesn't share memory with the original non-tensor/meta data.
    • pop(batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None): Removes and returns the specified keys as a new DataProto object.
    • rename(old_keys, new_keys): Renames keys within the batch (TensorDict) in-place.
    • union(other): Merges another DataProto into the current one. It performs a union on batch, non_tensor_batch, and meta_info. It raises an error if there are conflicting keys with different values or if batch sizes do not match.