一页看懂
问题设定与形式化
论文把 LLM router 的现状概括为 single-round, one-to-one mapping:每个 query 独立地交给一个模型。它指出两个困难:第一,每一轮"调用哪个 LLM"是离散选择,梯度方法(RouterDC、GraphRouter、FrugalGPT、HybridLLM 这类单次 router)扩展到多轮会迅速不可行,所以用 RL;第二,多跳问题需要"推理 ↔ 选模型"交错进行,而不是选一次、拿结果、结束。
于是 router 本身被实例化为一个 LLM(policy)。它的一条生成轨迹 y 里穿插了池子 𝒫 中模型的回复,优化目标是带 KL 正则的期望奖励(论文 Eq. 1):
这是通用形式(论文说可以套 PPO / GRPO);实际用 veRL 的 PPO + GAE,β 在代码里是 algorithm.kl_ctrl.kl_coef=0.001,KL 作为 token 级惩罚加在 reward 上。动作空间完全由四种 tag 定义:
| tag | 谁写 | 含义 | 代码里的处理 |
|---|---|---|---|
| <think>…</think> | policy | 内部推理。prompt 要求每次收到新信息都先 think,并说明"为什么需要外部信息、选哪个模型"。 | 不解析内容;格式奖励要求整段回复以 <think> 开头、至少一个且内容非空。 |
| <search>LLM-Name: query</search> | policy | 路由动作:选一个候选模型,并给它一个具体子问题。论文正文把它写成 <route>,代码沿用 Search-R1 的 <search>。 | postprocess_predictions 用正则取第一个 tag,按冒号切成模型名和子问题;名字经 check_llm_name 子串匹配映射到 API id。 |
| <information>…</information> | 环境 | 被调模型的回复(论文写作 <info>),拼回上下文。 | 截断到 600 token;这些 token 在 PPO loss 里被 mask 掉(state masking)。 |
| <answer>…</answer> | policy | 最终答案,结束 episode。 | extract_solution 取最后一个 answer 块,归一化后做 EM / F1。 |
_postprocess_responses 在第一个 </search>(否则第一个 </answer>)处截断,re.search 只取第一个 tag。"并行子问题 A、B"在这套协议里不存在,要做我们的 evidence graph 必须扩展解析器和环境步。动作协议与 prompt
整个"动作空间"就是 prompt 里的几条规则加 6 段模型描述(descriptor)。descriptor 是让 GPT-4o 读 NVIDIA 的 model card 写成的能力简介,论文称其只是 cold-start prior:训练中 policy 通过交互和 reward 自己学到各模型的真实强弱;新模型只要把描述追加进 prompt 就能零重训接入。
训练 prompt 模板要求:每次收到新信息先在 <think> 里推理;缺知识就写 <search>Candidate LLM: Query</search>;每次调用前必须说明"为什么需要外部信息"和"哪个 LLM 最合适";不要只盯着一个模型反复调,鼓励探索不同模型、也允许对同一问题多次调用。
注意:论文的 Intro 与 Limitations 说 descriptor 包含 pricing, latency, example performance,但附录给出的 descriptor 原文和代码里的字符串都只有规模、语言和特长,没有价格和延迟。成本信号完全来自 reward。
2Answer the given question. \ 3Every time you receive new information, you must first conduct reasoning inside <think> ... </think>. \ 4After reasoning, if you find you lack some knowledge, you can call a specialized LLM by writing a query inside <search> LLM-Name:Your-Query </search>. \ 5 6!!! STRICT FORMAT RULES for <search>: !!! 7 + You MUST replace LLM-Name with the EXACT name of a model selected from [Qwen2.5-7B-Instruct, LLaMA-3.1-8B-Instruct, LLaMA-3.1-70B-Instruct, Mistral-7B-Instruct, Mixtral-8x22B-Instruct, Gemma-2-27B-Instruct]. \ 8 + You MUST replace Your-Query with a CONCRETE QUESTION that helps answer the original question below. \ 9 + NEVER copy or paste model descriptions into <search>.
每个候选一段,共 6 段(Qwen2.5-7B、LLaMA-3.1-8B/70B、Mistral-7B、Mixtral-8x22B、Gemma-2-27B)。附录 Table 6 把它们压缩成一句话后平均 EM 从 0.416 变成 0.419:descriptor 的丰富度几乎不影响结果。
24Qwen2.5-7B-Instruct:\ 25Qwen2.5-7B-Instruct is a powerful Chinese-English instruction-tuned large language model designed for tasks in language, \ 26coding, mathematics, and reasoning. As part of the Qwen2.5 series, it features enhanced knowledge, stronger coding and \ 27math abilities, improved instruction following, better handling of long and structured texts, and supports up to 128K \ 28context tokens. It also offers multilingual capabilities across over 29 languages.\
{candidates_intro} 只是概念上的占位。池子里的模型只收到子问题,看不到原问题、历史或图像。prompt 要求它尽力回答或给背景,回答控制在 512 token 内,并且完全答不了时必须明说并让 base model 去问别的 LLM——附录案例里"先问 8B、失败再升级到 70B"的级联行为,升级信号正是从这句话来的。
91AGENT_PROMPT = """ 92You are a helpful assistant. \ 93You are participating in a multi-agent reasoning process, where a base model delegates sub-questions to specialized models like you. \ 94\nYour task is to do your **absolute best** to either: \n 95 + Answer the question directly, if possible, and provide a brief explanation; or \n 96 + Offer helpful and relevant context, background knowledge, or insights related to the question, even if you cannot fully answer it. \ 97 98If you are completely unable to answer the question or provide any relevant or helpful information, you must: \n 99 + Clearly state that you are unable to assist with this question, and \n 100 + Explicitly instruct the base model to consult other LLMs for further assistance. \ 101 102**Important Constraints**: \n 103 + Keep your response clear, concise, and informative (preferably under 512 tokens). Your response will help guide the base model’s reasoning and next steps. \n 104 + Stay strictly on-topic. Do not include irrelevant or generic content. \ 105 106\n\nHere is the sub-question for you to assist with: {query}\n 107"""
AGENT_PROMPT:被调 LLM 看到的全部内容。代码写的是 "multi-agent reasoning process",论文附录是 "multi-round"。多轮交互循环
<information> 拼回上下文),要么作答。路由最多 4 轮;第 4 轮之后若仍未作答,再生成一次但不允许路由。整条轨迹的 reward 只在最后一个 token 上给出。这套循环(LLMGenerationManager.run_llm_loop)是从 Search-R1 继承的,Router-R1 把"调检索器"换成"调模型池"。整批样本同步推进:每一轮 vLLM 对所有仍活跃的样本生成一次,然后解析动作、统一调 API、把结果拼回去。下面四段代码是循环里的四个关键动作。
vLLM 没有配置 stop token,每轮会一直生成到 1024 token 或 EOS,再在 Python 里截断到第一个 </search>(否则 </answer>),然后重新 tokenize。这意味着训练用的 token 序列并不严格等于采样出的序列(old log-prob 是在重新 tokenize 后的序列上算的,所以 PPO ratio 自洽)。
56 def _postprocess_responses(self, responses: torch.Tensor) -> torch.Tensor: 57 """Process responses to stop at search operation or answer operation.""" 58 responses_str = self.tokenizer.batch_decode( 59 responses, 60 skip_special_tokens=True 61 ) 62 63 responses_str = [resp.split('</search>')[0] + '</search>' 64 if '</search>' in resp 65 else resp.split('</answer>')[0] + '</answer>' 66 if '</answer>' in resp 67 else resp 68 for resp in responses_str]
只认第一个 <search> 或 <answer>。三种路由写法被标为无效:照抄占位符 LLM-Name / Your-Query、没有冒号、冒号后为空。无效动作在环境里得到空观测,模型下一轮接着自己的输出继续写,最终靠格式奖励 −1 来惩罚。
443 for prediction in predictions: 444 if isinstance(prediction, str): # for llm output 445 pattern = r'<(search|answer)>(.*?)</\1>' 446 match = re.search(pattern, prediction, re.DOTALL) 447 if match: 448 content = match.group(2).strip() # Return only the content inside the tags 449 action = match.group(1) 450 if action == "search" and ("llm-name" in content.strip().lower() or "your-query" in content.strip().lower()): 451 action = "route invalid-1" 452 elif action == "search" and ":" not in content: 453 action = "route invalid-2" 454 elif action == "search" and content.strip().lower().split(":")[-1].strip() == "": 455 action = "route invalid-3" 456 else: 457 content = '' 458 action = None 459 else: 460 raise ValueError(f"Invalid prediction type: {type(prediction)}") 461 462 actions.append(action) 463 contents.append(content) 464 465 return actions, contents
postprocess_predictions:动作 = 第一个 tag 的名字,内容 = tag 里的文本。所有 search 动作的子问题一次性交给 batch_route(内部是 10 线程的 API 调用)。模型名非法或 API 出错时观测是 <information>None</information>,并记 valid_action=0、成本 0。answer 动作把该样本标为 done。每个样本本轮的成本累加进 batch_completion_tokens(名字叫 tokens,实际已经乘了单价)。
397 if action == 'answer': 398 next_obs.append('') 399 dones.append(1) 400 valid_action.append(1) 401 is_route.append(0) 402 cur_completion_tokens.append(0.0) 403 elif action == 'search': 404 if route_results[0].strip().lower() == "llm name error": 405 next_obs.append(f'\n\n<information>None</information>\n\n') 406 route_results.pop(0) 407 valid_action.append(0) 408 elif route_results[0].strip().lower() == "api request error": 409 next_obs.append(f'\n\n<information>None</information>\n\n') 410 route_results.pop(0) 411 valid_action.append(0) 412 else: 413 next_obs.append(f'\n\n<information>{route_results.pop(0).strip()}</information>\n\n') 414 valid_action.append(1) 415 dones.append(0) 416 is_route.append(1) 417 cur_completion_tokens.append(completion_tokens_list.pop(0)) 418 else: 419 next_obs.append('') 420 dones.append(0) 421 valid_action.append(0) 422 is_route.append(0) 423 cur_completion_tokens.append(0.0)
execute_predictions 的核心分支。信息两侧的空行和 .strip() 是训练时的精确格式,infer_vllm.py 用的格式和它不一样。被调模型的回复 tokenize 后硬截到 max_obs_length=600(论文附录超参表里的 "Max Length for LLM API Response 600" 指的就是它;API 端的生成上限是 512)。硬截可能把 </information> 切掉,之后格式奖励会因 tag 不配对判 −1,这是 policy 无法控制的噪声。
80 def _process_next_obs(self, next_obs: List[str]) -> torch.Tensor: 81 """Process next observations from environment.""" 82 83 next_obs_ids = self.tokenizer( 84 next_obs, 85 padding='longest', 86 return_tensors='pt', 87 add_special_tokens=False, # Prevents adding special tokens 88 )['input_ids'] 89 90 if next_obs_ids.shape[1] > self.config.max_obs_length: 91 print(f"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, {next_obs_ids.shape[1]} & {self.config.max_obs_length}") 92 next_obs_ids = next_obs_ids[:, :self.config.max_obs_length] 93 94 return next_obs_ids
路由池与成本记账
池子是一个 OpenAI-compatible 端点上的若干模型 id(论文用 NVIDIA NIM,id 形如 meta/llama-3.1-70b-instruct)。policy 写出的名字经宽松的子串匹配映射到 id,然后带着 AGENT_PROMPT 发一次 chat completion;返回的文本和 completion_tokens 一起回来,成本 = completion_tokens × 该模型每百万 token 的单价(没有除以 10⁶,所以单位相当于 10⁻⁶ USD;只计输出 token,不计 prompt token,也不计延迟)。
论文说 m(·) 把模型参数量映射到每 token 成本(按 API 定价档位初始化)。代码里 m 就是一张按模型 id 查的价格表:Mixtral-8x22B(1.2)比 LLaMA-70B(0.88)贵,并不严格按参数量单调。
79API_PRICE_1M_TOKENS = { 80 "qwen/qwen2.5-7b-instruct": 0.3, 81 "meta/llama-3.1-70b-instruct": 0.88, 82 "meta/llama-3.1-8b-instruct": 0.18, 83 "mistralai/mistral-7b-instruct-v0.3": 0.2, 84 "mistralai/mixtral-8x22b-instruct-v0.1": 1.2, 85 "google/gemma-2-27b-it": 0.8, 86 "writer/palmyra-creative-122b": 1.8, 87 "nvidia/llama3-chatqa-1.5-8b": 0.18, 88}
子串匹配很宽松:任何含 "qwen" 的名字都映射到 7B 并按 0.3 计价;"llama" 必须带 8b / 70b 等尺寸,否则视为非法名。它还接受 nemotron-51b / 49b、granite 这几个不在价格表里的名字,这是一个潜在的 KeyError 崩溃点(见"坑")。
132def check_llm_name(target_llm): 133 TAU = 0 134 LLM_NAME = "" 135 if "qwen" in target_llm: 136 LLM_NAME = "qwen/qwen2.5-7b-instruct" 137 elif "palmyra" in target_llm or "creative" in target_llm: 138 LLM_NAME = "writer/palmyra-creative-122b" 139 elif "llama" in target_llm: 140 if "70b" in target_llm: 141 LLM_NAME = "meta/llama-3.1-70b-instruct" 142 elif "51b" in target_llm: 143 LLM_NAME = "nvidia/llama-3.1-nemotron-51b-instruct" 144 elif "49b" in target_llm: 145 LLM_NAME = "nvidia/llama-3.3-nemotron-super-49b-v1" 146 elif "8b" in target_llm: 147 if "chatqa" in target_llm: 148 LLM_NAME = "nvidia/llama3-chatqa-1.5-8b" 149 else: 150 LLM_NAME = "meta/llama-3.1-8b-instruct" 151 else: 152 # print("!!!!!!!!!!!LLM Name Error!!!!!!!!!!!", target_llm) 153 LLM_NAME = "" 154 elif "mistral" in target_llm: 155 LLM_NAME = "mistralai/mistral-7b-instruct-v0.3" 156 elif "mixtral" in target_llm: 157 LLM_NAME = "mistralai/mixtral-8x22b-instruct-v0.1" 158 elif "granite" in target_llm: 159 LLM_NAME = "ibm/granite-3.0-8b-instruct" 160 elif "gemma" in target_llm: 161 LLM_NAME = "google/gemma-2-27b-it" 162 TAU = 0.1 163 else: 164 # print("!!!!!!!!!!!LLM Name Error!!!!!!!!!!!", target_llm) 165 LLM_NAME = "" 166 167 return LLM_NAME, TAU
check_llm_name:Gemma 用 temperature 0.1,其余 0(greedy)。做自己的池子时应换成显式 registry(规范名 → id / 单价 / 温度 / 端点)。返回三元组 (q_id, 回复文本, 成本)。API 异常时回复是 "API Request Error"、成本 0;get_llm_response_via_api 对非超时错误最多重试 500 次、每次等 5 秒,一次限流可能卡住线程约 40 分钟。
110def request_task(data): 111 q_id, query_text, TAU, LLM_NAME, api_base, api_key = data 112 if LLM_NAME == "": 113 print("LLM Name Error") 114 return q_id, "LLM Name Error", 0.0 115 print(LLM_NAME) 116 try: 117 input_prompt = AGENT_PROMPT.format_map({"query": query_text}) 118 single_response, completion_tokens = get_llm_response_via_api(prompt=input_prompt, 119 base_url=api_base, 120 api_key=api_key, 121 TAU=TAU, 122 LLM_MODEL=LLM_NAME) 123 print(single_response, completion_tokens) 124 except Exception as e: 125 print(e) 126 single_response = "API Request Error" 127 completion_tokens = 0.0 128 129 return q_id, single_response, int(completion_tokens) * API_PRICE_1M_TOKENS[LLM_NAME]
try 之外。Reward 设计
三项规则奖励,论文 Eq.(2)–(5):
Eq. (3)Routcome = EM(ya, gt)
Eq. (4)Rcost ∝ − m(PLLM) · Tout,归一化到 [0, 1]
Eq. (5)rφ(x, y) = Rformat + (1 − α) Routcome + α Rcost
论文说这三项是分层的:格式奖励为 −1 时另外两项清零。主实验 α = 0(完全不看成本),成本实验扫 α ∈ {0.6, 0.7, 0.8, 0.9}。下面按代码把每一项落实。
① 格式奖励:五条硬规则
| 规则(Appendix A.1.2) | 代码检查 |
|---|---|
| tag 必须成对且不能嵌套 | 四种 tag 的开闭计数相等;任一 tag 内容里再出现 tag → nesting |
以 <think> 开头、以唯一的 <answer> 结尾 | strip 后前 7 个字符是 <think>,</answer> 必须是最后 9 个字符 |
| 至少一个 think、恰好一个 answer | answer_enclose_count == 1,think_enclose_count > 0,think 内容不能为空或 "..." |
| 每个 search 必须配一个 information | route_enclose_count == info_enclose_count |
| 路由内容必须是 "llm_name: query",名字合法、两边非空 | 恰好一个冒号;不含占位符;query 里不能含模型名;is_valid_llm_name |
代码里有 MAX / MEDIUM / SMALL 三档惩罚常量,但都等于 −1.0,所以实际只有 −1 / 0 两种取值,和论文一致。格式检查作用在整段 response(含被调模型返回的 information)上:被调模型的回复里只要出现 <answer>、<think> 字样,policy 就会被判 −1。
127 if think_punish: 128 return PUNISH_REWARD_MAX 129 130 if is_nesting: 131 return PUNISH_REWARD_MAX 132 133 if answer_enclose_count != 1 or think_enclose_count == 0 or route_enclose_count != info_enclose_count: 134 return PUNISH_REWARD_MAX 135 136 completion = completion.strip() 137 if not completion[:len("<think>")] == "<think>": 138 return PUNISH_REWARD_MAX 139 140 if not completion.rfind("</answer>") + 9 == len(completion): 141 return PUNISH_REWARD_MAX 142 143 if query_format_punish: 144 return PUNISH_REWARD_MEDIUM 145 146 if llm_name_punish: 147 return PUNISH_REWARD_SMALL 148 149 return 0.0
format_reward 的裁决部分(前面 L84–L124 是逐 tag 的计数与检查)。② 结果奖励:Exact Match
答案从最后一个 <answer> 块里取出,做 SQuAD 式归一化(小写、去冠词、去标点、合并空白)后与任一 gold 完全相等记 1。训练只用 EM(+reward_metric="em");代码也支持 F1(对多个 gold 取最大 token-F1),但只用于评测。
20def normalize_answer(s): 21 def remove_articles(text): 22 return re.sub(r"\b(a|an|the)\b", " ", text) 23 24 def white_space_fix(text): 25 return " ".join(text.split()) 26 27 def remove_punc(text): 28 exclude = set(string.punctuation) 29 return "".join(ch for ch in text if ch not in exclude) 30 31 def lower(text): 32 return text.lower() 33 34 return white_space_fix(remove_articles(remove_punc(lower(s)))) 35 36 37def em_check(prediction, golden_answers): 38 if isinstance(golden_answers, str): 39 golden_answers = [golden_answers] 40 normalized_prediction = normalize_answer(prediction) 41 score = 0 42 for golden_answer in golden_answers: 43 golden_answer = normalize_answer(golden_answer) 44 if golden_answer == normalized_prediction: 45 score = 1 46 break 47 return score
③ 成本奖励:sqrt + 滑动窗口分位归一化 + 取反
r 是整条轨迹所有路由调用成本之和;rmin / rmax 是最近 1000 个样本的 5 / 95 百分位;区间小于 ε 时返回中性值 0.5。于是"没有路由、成本为 0"的轨迹拿到 Rcost = 1(最大)。
36window_size = 1000 37use_percentile = True 38q_low, q_high = 0.05, 0.95 39eps = 1e-8 40 41transform = 'sqrt' 42alpha = 0.01 43 44buffer = deque(maxlen=window_size) 45 46 47def reward_preprocess(r: float) -> float: 48 if transform == 'log': 49 return np.log1p(alpha * r) 50 elif transform == 'sqrt': 51 return np.sqrt(r) 52 else: 53 return r 54 55 56def normalize_reward(r): 57 r = reward_preprocess(r) 58 buffer.append(r) 59 arr = np.array(buffer) 60 if use_percentile and len(arr) >= 2: 61 r_min = np.percentile(arr, 100 * q_low) 62 r_max = np.percentile(arr, 100 * q_high) 63 else: 64 r_min = arr.min() 65 r_max = arr.max() 66 67 denom = r_max - r_min 68 if denom < eps: 69 return 0.5 70 71 r_scaled = (r - r_min) / denom 72 return 1.0 - float(np.clip(r_scaled, 0.0, 1.0))
buffer 是模块级全局 deque:train 和 val 两个 RewardManager 共用,而且当前样本先入队再算分位数,同一成本的归一化值依赖历史和 batch 内顺序。第 42 行的 alpha=0.01 是未启用的 log 变换系数,不是论文的 α。④ 组合:代码比论文更严格
论文只说 format = −1 时后两项清零。代码(EM 模式)实际是三级门控:
r = 0,格式正确但答错
r = (1 − α)·1 + α·Rcost,答对
cost 只在答对时才计入。这是真正防 reward hacking 的机制:答错的轨迹无论多便宜都是 0,不会出现"答错但省钱"被奖励。α > 0 时答对的轨迹 reward 落在 [1−α, 1],省钱只在"答对"内部分高下。论文的 Eq.(5) 没有写出这一层。
157 if score == 0: 158 return score, api_cost, score + format_score 159 else: 160 return score, api_cost, (score + format_score) * (1.0 - cost_coe) + api_cost * cost_coe 161 else: 162 if em_check(answer, ground_truth['target']): 163 if format_score == -1.0: 164 return score, api_cost, format_score 165 else: 166 return score, api_cost, (score + format_score) * (1.0 - cost_coe) + api_cost * cost_coe 167 else: 168 if format_score == -1.0: 169 return format_score, api_cost, format_score 170 else: 171 return format_score, api_cost, format_score
compute_score_em 训练分支的 EM 部分:第 166 行是 (score + format)·(1−α) + api_cost·α,只有 em_check 为真才会走到。RewardManager 只解码 response(不含 prompt),依次算格式分、路由次数、归一化成本,再调用 compute_score_em;reward、metric、cost、route_cnt 四个标量都写在最后一个有效 token 上,之后由 GAE 沿序列向前传播。验证模式额外返回 EM 与 F1 两个 metric。
237 strict_format_score = format_reward(completion=sequences_str) 238 route_cnt = route_count(completion=sequences_str) 239 240 ground_truth = data_item.non_tensor_batch['reward_model']['ground_truth'] 241 242 # API Cost Reward 243 api_cost = normalize_reward(token_price[i]) 244 normalized_cost_reward.append(api_cost) 245 246 # select rm_score 247 data_source = data_item.non_tensor_batch['data_source'] 248 compute_score_fn = _select_rm_score_fn(data_source) 249 250 if self.state == "train": 251 metric_score, cost_score, reward_score = compute_score_fn(solution_str=sequences_str, ground_truth=ground_truth, format_score=strict_format_score, api_cost=api_cost, state=self.state, reward_metric=self.reward_metric, cost_coe=self.cost_coe) 252 metric_tensor[i, valid_response_length - 1] = metric_score 253 else: 254 metric_score_em, metric_score_f1, cost_score, reward_score = compute_score_fn(solution_str=sequences_str, ground_truth=ground_truth, format_score=strict_format_score, api_cost=api_cost, state=self.state, reward_metric=self.reward_metric, cost_coe=self.cost_coe) 255 metric_em_tensor[i, valid_response_length - 1] = metric_score_em 256 metric_f1_tensor[i, valid_response_length - 1] = metric_score_f1 257 258 reward_tensor[i, valid_response_length - 1] = reward_score 259 cost_tensor[i, valid_response_length - 1] = cost_score 260 route_tensor[i, valid_response_length - 1] = route_cnt
PPO 训练与 state masking
论文一句话:"外部回复(<info>)不进 loss,以保证训练稳定"。代码的做法是构造一份把 information token 全替换成 pad 的平行序列 responses_with_info_mask,由它派生出 info_mask,再作为 actor 的 loss_mask。
拼接 prompt、response、information 时同时维护两份张量:一份真实,一份把 information 换成 pad。之后"token ≠ pad"就是 mask。它依赖 pad id 不会出现在正文里(Qwen2.5 的 pad ≠ eos,安全;Llama-3.2 若 pad=eos 会出错)。
123 def _info_masked_concatenate_with_padding(self, 124 prompt: torch.Tensor, 125 prompt_with_mask: torch.Tensor, 126 response: torch.Tensor, 127 info: torch.Tensor = None, 128 pad_to_left: bool = True 129 ) -> torch.Tensor: 130 """Concatenate tensors and handle padding. Additionally, create a mask (info_mask) to cover the information block if it exists.""" 131 pad_id = self.tokenizer.pad_token_id 132 tensors = [prompt, response] 133 tensors_with_mask = [prompt_with_mask, response] 134 if info is not None: 135 tensors.append(info) 136 info_mask = torch.full(info.size(), pad_id, dtype=info.dtype, device=info.device) # information mask 137 tensors_with_mask.append(info_mask) 138 139 concatenated = torch.cat(tensors, dim=1) 140 concatenated_with_info = torch.cat(tensors_with_mask, dim=1) 141 mask = concatenated != pad_id if pad_to_left else concatenated == pad_id 142 sorted_indices = mask.to(torch.int64).argsort(dim=1, stable=True) 143 padded_tensor = concatenated.gather(1, sorted_indices) 144 padded_tensor_with_info = concatenated_with_info.gather(1, sorted_indices) 145 146 return padded_tensor, padded_tensor_with_info
actor.state_masking=true 时 PPO 的 clipped surrogate 与 entropy 只在 policy 自己生成的 token 上平均;KL 惩罚同样用 info_mask。但 masking 是部分的:critic 的 value loss 与 GAE 的白化统计仍用 attention_mask,information token 上也会算 value 和 advantage(γ = λ = 1 时影响不大)。
237 responses = data['responses'] 238 response_length = responses.size(1) 239 attention_mask = data['attention_mask'] 240 response_mask = attention_mask[:, -response_length:] 241 if self.config.state_masking: 242 response_mask = data['loss_mask'] 243 old_log_prob = data['old_log_probs'] 244 advantages = data['advantages'] 245 246 clip_ratio = self.config.clip_ratio 247 entropy_coeff = self.config.entropy_coeff 248 249 # all return: (bsz, response_length) 250 entropy, log_prob = self._forward_micro_batch(micro_batch=data, temperature=temperature) 251 252 pg_loss, pg_clipfrac, ppo_kl = core_algos.compute_policy_loss(old_log_prob=old_log_prob, 253 log_prob=log_prob, 254 advantages=advantages, 255 eos_mask=response_mask, 256 cliprange=clip_ratio) 257 # compute entropy loss from entropy 258 entropy_loss = verl_F.masked_mean(entropy, response_mask) 259 260 # compute policy loss 261 policy_loss = pg_loss - entropy_loss * entropy_coeff
response_mask = data['loss_mask']。KL 是 token 级惩罚加在 reward 上(kl_coef=0.001,fixed),不是 loss 项;clipped surrogate 是标准 PPO(clip ratio 来自 yaml 默认 0.2)。advantage 用 GAE(γ = λ = 1.0),需要 critic——这也是训练同时加载 actor、ref、critic 三份 3B 模型的原因。
86def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty='kl'): 87 responses = data.batch['responses'] 88 response_length = responses.size(1) 89 token_level_scores = data.batch['token_level_scores'] 90 batch_size = data.batch.batch_size[0] 91 attention_mask = data.batch['info_mask'] if 'info_mask' in data.batch else data.batch['attention_mask'] 92 response_mask = attention_mask[:, -response_length:] 93 94 # compute kl between ref_policy and current policy 95 if 'ref_log_prob' in data.batch.keys(): 96 kld = core_algos.kl_penalty(data.batch['old_log_probs'], data.batch['ref_log_prob'], 97 kl_penalty=kl_penalty) # (batch_size, response_length) 98 kld = kld * response_mask 99 beta = kl_ctrl.value 100 else: 101 beta = 0 102 kld = torch.zeros_like(response_mask, dtype=torch.float32) 103 104 token_level_rewards = token_level_scores - beta * kld
apply_kl_penalty:优先用 info_mask。185 negative_approx_kl = log_prob - old_log_prob 186 ratio = torch.exp(negative_approx_kl) 187 ppo_kl = verl_F.masked_mean(-negative_approx_kl, eos_mask) 188 189 pg_losses = -advantages * ratio 190 pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - cliprange, 1.0 + cliprange) 191 192 pg_loss = verl_F.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask) 193 pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask) 194 return pg_loss, pg_clipfrac, ppo_kl
eos_mask 就是上面的 loss_mask。训练超参(train.sh / Appendix Table 8)
| 项 | 值 | 说明 |
|---|---|---|
| policy / critic 基座 | Qwen2.5-3B-Instruct(或 Llama-3.2-3B-Instruct) | critic 也从同一权重初始化;ref 冻结 |
| 优化器 | actor lr 1e-6,critic lr 1e-5,warmup 0 | PPO + GAE(γ = λ = 1),kl_coef 0.001 |
| batch | 64 题 / step,mini-batch 32,micro-batch 8,n_agent 1 | 每题一条轨迹;GRPO 需 n_agent > 1 |
| 步数 | total_training_steps 225(实际 224 次更新),save/test 每 15 步 | 最后一个 checkpoint 是 global_step_210 |
| 长度 | prompt 4096 / response 1024 每轮 / start 2048 / obs 600 | 上下文超长时从左边截,先丢的是指令和 descriptor |
| 路由 | max_turns 4;API max_tokens 512,temperature 0(Gemma 0.1),seed 42 | 第 4 轮后再生成一次但不路由 |
| rollout | vLLM 0.6.3,TP 1,gpu_memory_utilization 0.6,temperature 1.0(train 与 eval) | 验证代码设了 greedy,但该设置在循环里丢失,实际也是采样 |
| 显存 | FSDP,actor / critic / ref 全部 param + grad + optimizer offload,gradient checkpointing,remove_padding | 4 卡 A6000;复现见实验 tab |
| reward | reward_metric em;cost_coe(α)0.0 | 成本实验 0.6 / 0.7 / 0.8 / 0.9 |
与 Search-R1 的差异
论文没有正面说自己和 Search-R1 的区别。对比两个仓库(Router-R1 HEAD 801a240 vs Search-R1 main)可以精确列出:
| 组件 | Search-R1 | Router-R1 |
|---|---|---|
| 外部动作 | <search>query</search> → 检索器 HTTP 调用(Wikipedia + E5) | <search>LLM-Name: query</search> → 模型池 API 调用(新文件 route_service.py) |
| 观测 | top-k 文档 | 被调模型的回复(≤600 token);非法名 / API 错误给 None;无效动作给空观测(Search-R1 会给纠正提示) |
| 成本 | 无 | 每条轨迹累加 completion_tokens × 单价,经 meta_info 传给 reward;_balance_batch 重排时手动同步 |
| reward | EM(+ 可选格式分) | 严格格式奖励(5 条规则)+ EM/F1 + 归一化成本奖励,α 混合,三级门控 |
| prompt | 搜索指令 | 路由指令 + 6 段 descriptor + 给被调模型的 AGENT_PROMPT |
| 配置 / 日志 | search_url、topk | do_route、+reward_metric、+cost_coe、+api_base、+api_key;新增 cost / metric / route_cnt 指标,按数据源报 EM 与 F1 |
| 未改动 | rollout 骨架 run_llm_loop、info_mask / state masking、PPO + GAE(core_algos.py 完全相同)、tensor_helper.py 逐字节相同 | |
读代码发现的坑(精选)
- 价格表 KeyError 会打崩 rollout。
check_llm_name接受 nemotron-51b/49b、granite,但API_PRICE_1M_TOKENS没有这三个 key;查价在 try 之外,policy 只要写出这类名字就会在线程池里抛异常。 - 子问题在第二个冒号处被截断。
split(':')[1]只取第一和第二个冒号之间;"At 10:30, who…" 实际发出的是 " At 10"。格式奖励也要求恰好一个冒号。 - 验证并不是 greedy。
_validate设了do_sample=False,但run_llm_loop重建 DataProto 时把 meta_info 丢了,vLLM 按 temperature 1.0 采样。单次评测有方差,论文没有报告 seed 或区间。 - cost 只在答对时生效,且归一化依赖全局 buffer。val 的 cost 会污染训练窗口;val_only 时第一个样本恒为 0.5。日志里的 cost 是归一化后的 [0,1](越大越省),不是美元。
- checkpoint 路径与步数 off-by-one。225 步实际 224 次更新,最后落盘的是
global_step_210;test.sh 的 ACTOR_PATH 要自己补上这一级目录。 - val DataLoader drop_last=True。Bamboogle 只有 125 题,
val_batch_size=500会得到 0 个 batch 直接报错;必须设成 ≤ 数据集大小。 - 每轮只能一个动作、一次路由。并行子问题、多模型聚合都需要改解析器和环境步。
- 多模态的硬阻塞。数据集只有文本 prompt(没有 pixel_values),rollout 只传 token id,池子只发纯文本 messages,成本不计图像 / prompt token,
_select_rm_score_fn只认 7 个文本 QA 数据源。
对我们项目的启示
可以直接继承的:"policy 自己做 router"的训练框架(动作协议 + 规则奖励 + state masking + veRL),descriptor 式的零重训扩池,cost 进 reward 的做法和它的归一化细节。
必须改的:
- Planner 的角色要重新设计。论文的反面证据很硬:先静态分解再逐个路由的 Prompt LLM* / KNN Router*(3B base)比不分解还差 0.04–0.16 EM,成本高十几倍。我们的 evidence-graph planner 不能是"先拆完再派",而要让分解本身也在 RL 回路里、能回看证据再决定下一步。
- 动作协议要支持并行与引用。一轮多个 search、子问题引用图像 / 区域 / 前序证据;格式奖励里"恰好一个冒号"这类规则要换成 JSON 之类的结构化动作。
- 成本要重新定义。论文只算输出 token × API 单价;我们的池子里有本地小视觉模型和 VLM,要把 prompt / 图像 token、GPU 秒、工具调用时延都计入,并做论文自己建议的 task-level normalization。
- outcome 奖励要适配 VQA。DocVQA 的 ANLS、ChartQA 的 relaxed accuracy 等连续指标下,"答错就把 cost 清零"的门控是否还合适需要重新评估;多选题可直接用 accuracy。
- 多轮能力要被真正用到。Router-R1 平均每题只调 1.0–1.4 次。要证明 planner + 多轮路由的价值,评测集必须是真正需要多条证据的组合式问题(见实验 tab 的数据集调研)。


