[NPU] Add Search-R1 example by Fulin-Gao · Pull Request #311 · vllm-project/vime

Fulin-Gao · 2026-07-02T03:44:59Z

Search-R1 lite

This example is ported from slime (Megatron + SGLang) and adapted for vime (Megatron + vLLM). It provides minimal reproduction of Search-R1, demonstrating multi-turn conversation and tool-calling workflows in vime.

Overview

The Search-R1 example provides:

Multi-turn conversation with tool-calling (search/answer actions)
Dual search backend support: local dense retriever (FAISS + E5) or Google Search (serper.dev)
GRPO-based RL training with exact-match (EM) reward for QA tasks
TIS (Trajectory Importance Sampling) support for handling train/inference mismatch
Format-aware reward that evaluates both answer correctness and output structure

Files

File	Description
`generate_with_search.py`	Main generation function with multi-turn search + answer tool-calling, and reward function
`google_search_server.py`	Google Search backend via serper.dev API
`local_search_server.py`	Local search backend that wraps the retrieval server
`qa_em_format.py`	QA exact-match scoring with format validation and retrieval correctness check
`run_qwen3_4b_npu.sh`	Training launch script (NPU 8-card, Qwen2.5-3B, GRPO)
`local_dense_retriever/retrieval_server.py`	Dense retriever server (FAISS + E5 model, FastAPI)
`local_dense_retriever/download.py`	Download wiki-18 index and corpus from HuggingFace

Results

Closes #302

read-the-docs-community · 2026-07-02T03:46:17Z

Documentation build overview

📚 vime | 🛠️ Build #33489336 | 📁 Comparing f5ba4e4 against latest (5532763)

🔍 Preview build

43 files changed · + 2 added · ± 27 modified · - 14 deleted

+ Added

± Modified

- Deleted

gemini-code-assist

Code Review

This pull request introduces 'Search-R1 lite', a minimal reproduction of Search-R1 supporting both local dense retrieval (FAISS + E5) and Google Search backends, complete with GRPO-based RL training scripts, evaluation metrics, and documentation. The code review identified several critical issues, including a logic error in extract_solution that breaks zero-shot RL training, an indexing/alignment bug in the Google search server when items lack links, and an unpacking bug in the retrieval server when batch size is exactly two. Additionally, the reviewer pointed out a list length mismatch in the BM25 retriever, an ignored --repo_id argument in the download script, and type annotation mismatches in execute_predictions and is_retrieval_correct.

gemini-code-assist · 2026-07-02T03:47:20Z

+    # If there are 0 or exactly 1 matches, return None
+    if len(matches) <= 1:
+        return None


In zero-shot RL training, there is exactly 1 <answer> block in the generated response and 0 in the prompt. If len(matches) <= 1 returns None, then extract_solution will always return None for correct zero-shot responses, completely breaking the reward signal and preventing the model from receiving a positive reward. We should check if len(matches) == 0 instead.

Suggested change

# If there are 0 or exactly 1 matches, return None

if len(matches) <= 1:

return None

# If there are no matches, return None

if len(matches) == 0:

return None

Declined, not our scenario

gemini-code-assist · 2026-07-02T03:47:20Z

+        links = [item.get("link", "") for item in items if "link" in item]
+        web_contents = await fetch_all(links)


Filtering links using if "link" in item can result in len(links) being smaller than len(items). When iterating over items and accessing web_contents[i], this will raise an IndexError or cause incorrect alignment between the items and their fetched web contents. Since fetch already handles empty URLs gracefully, we should keep the list aligned by including empty strings for items without links.

Suggested change

links = [item.get("link", "") for item in items if "link" in item]

web_contents = await fetch_all(links)

links = [item.get("link", "") for item in items]

web_contents = await fetch_all(links)

gemini-code-assist · 2026-07-02T03:47:20Z

+    scores = []
+    try:
+        results, scores = tmp
+    except ValueError:
+        results = tmp


There is an unpacking bug when return_scores is False and the batch size is exactly 2. In this case, tmp is a list of length 2 (containing results for both queries). The try-except block will successfully unpack tmp into results and scores, which is incorrect because scores should be empty and results should contain the full list of results. We should explicitly check request.return_scores instead of relying on a try-except block.

Suggested change

scores = []

try:

results, scores = tmp

except ValueError:

results = tmp

if request.return_scores:

results, scores = tmp

else:

results = tmp

scores = []

gemini-code-assist · 2026-07-02T03:47:20Z

+    return action, content
+
+
+async def execute_predictions(prediction: str) -> str:


The type annotation for execute_predictions indicates that it returns a str, but the function actually returns a tuple of (str, bool) (i.e., next_obs, done). This type mismatch can cause static analysis or type-checking tools to report errors.

Suggested change

async def execute_predictions(prediction: str) -> str:

async def execute_predictions(prediction: str) -> tuple[str, bool]:

gemini-code-assist · 2026-07-02T03:47:21Z

+        scores = [hit.score for hit in hits]
+        if len(hits) < num:
+            warnings.warn("Not enough documents retrieved!", stacklevel=2)
+        else:
+            hits = hits[:num]


The scores list is populated from hits before hits is sliced to num. This results in a length mismatch where scores has more elements than results. We should compute scores after slicing hits to ensure their lengths are aligned.

Suggested change

scores = [hit.score for hit in hits]

if len(hits) < num:

warnings.warn("Not enough documents retrieved!", stacklevel=2)

else:

hits = hits[:num]

if len(hits) < num:

warnings.warn("Not enough documents retrieved!", stacklevel=2)

else:

hits = hits[:num]

scores = [hit.score for hit in hits]

gemini-code-assist · 2026-07-02T03:47:21Z

+    return [match.strip() for match in matches]
+
+
+def is_retrieval_correct(text: str, golden_answers: list[str]) -> list[str]:


The type annotation for is_retrieval_correct indicates that it returns list[str], but the function actually returns a boolean (True or False). This type mismatch should be corrected.

Suggested change

def is_retrieval_correct(text: str, golden_answers: list[str]) -> list[str]:

def is_retrieval_correct(text: str, golden_answers: list[str]) -> bool:

gemini-code-assist · 2026-07-02T03:47:21Z

+
+args = parser.parse_args()
+
+repo_id = "PeterJinGo/wiki-18-e5-index"


The command-line argument --repo_id is defined on line 24 but is completely ignored because repo_id is hardcoded to "PeterJinGo/wiki-18-e5-index" on line 29. We should use args.repo_id instead.

Suggested change

repo_id = "PeterJinGo/wiki-18-e5-index"

repo_id = args.repo_id

CalvinXKY · 2026-07-02T08:08:53Z

+SEMAPHORE = asyncio.Semaphore(SEARCH_R1_CONFIGS["search_concurrency"])
+
+
+def _build_inference_sampling_params(sampling_params: dict) -> dict:


Could we import the framework helper？ to avoid drift when rollout sampling mapping changes. vime.rollout.vllm_rollout._build_inference_sampling_params

CalvinXKY · 2026-07-02T08:10:48Z

Why do we see sharp peaks on the diff curve?

CalvinXKY · 2026-07-02T08:14:43Z

Results show training reward and logprob diff only. To claim Search-R1 effectiveness (not just pipeline health), can you add periodic eval EM on NQ/HotpotQA test (uncomment --eval-prompt-data) and a base vs trained EM table? Slime example has the hooks but no numbers either — we'd like vime ascend to be the first with a minimal reproducible EM curve.

Fulin-Gao · 2026-07-08T02:27:09Z

Why do we see sharp peaks on the diff curve?

We suspect the issue arises from partial timeouts in the retrieval service rather than excessive overall load. We have reduced the concurrency, rerun the experiment, and updated the reward and diff curves.

Fulin-Gao · 2026-07-08T02:32:17Z

Results show training reward and logprob diff only. To claim Search-R1 effectiveness (not just pipeline health), can you add periodic eval EM on NQ/HotpotQA test (uncomment --eval-prompt-data) and a base vs trained EM table? Slime example has the hooks but no numbers either — we'd like vime ascend to be the first with a minimal reproducible EM curve.

We provide evaluation results sampled at 50-step intervals within the first 200 training steps, which demonstrate a steady improvement in evaluation scores.

Signed-off-by: gaofulin <gaofulin1@huawei.com>

gemini-code-assist Bot reviewed Jul 2, 2026

View reviewed changes

Fulin-Gao marked this pull request as ready for review July 2, 2026 07:31

CalvinXKY reviewed Jul 2, 2026

View reviewed changes

Fulin-Gao mentioned this pull request Jul 8, 2026

[Ascend][RFC] vime-ascend Build and Roadmap #243

Open

20 tasks

gaofulin added 3 commits July 8, 2026 10:51

[NPU] Add Search-R1 example

b7e366f

Signed-off-by: gaofulin <gaofulin1@huawei.com>

[NPU] Add Search-R1 example

e556693

Signed-off-by: gaofulin <gaofulin1@huawei.com>

[NPU] Add Search-R1 example

3126f08

Signed-off-by: gaofulin <gaofulin1@huawei.com>

Fulin-Gao force-pushed the search-r1 branch from 24af89a to 3126f08 Compare July 8, 2026 02:53

[NPU] Add Search-R1 example

f5ba4e4

Signed-off-by: gaofulin <gaofulin1@huawei.com>

		links = [item.get("link", "") for item in items if "link" in item]
		web_contents = await fetch_all(links)

		return action, content


		async def execute_predictions(prediction: str) -> str:

		return [match.strip() for match in matches]


		def is_retrieval_correct(text: str, golden_answers: list[str]) -> list[str]:


		args = parser.parse_args()

		repo_id = "PeterJinGo/wiki-18-e5-index"

	repo_id = "PeterJinGo/wiki-18-e5-index"
	repo_id = args.repo_id

		SEMAPHORE = asyncio.Semaphore(SEARCH_R1_CONFIGS["search_concurrency"])


		def _build_inference_sampling_params(sampling_params: dict) -> dict:

Uh oh!

Conversation

Fulin-Gao commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Search-R1 lite

Overview

Files

Results

Uh oh!

read-the-docs-community Bot commented Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Fulin-Gao Jul 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 2, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CalvinXKY commented Jul 2, 2026

Uh oh!

CalvinXKY commented Jul 2, 2026

Uh oh!

Fulin-Gao commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fulin-Gao commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Fulin-Gao commented Jul 2, 2026 •

edited

Loading

read-the-docs-community Bot commented Jul 2, 2026 •

edited

Loading

Fulin-Gao Jul 2, 2026 •

edited

Loading

Fulin-Gao commented Jul 8, 2026 •

edited

Loading

Fulin-Gao commented Jul 8, 2026 •

edited

Loading