happy-llm 实战：基于 vLLM 实现 s1 论文的 Test-time Scaling 思考预算（Thinking Budget）

happy-llm 实战基于 vLLM 实现 s1 论文的 Test-time Scaling 思考预算Thinking Budget【免费下载链接】happy-llm 从零开始构建大模型项目地址: https://gitcode.com/GitHub_Trending/ha/happy-llm本指南以 s1-vllm-thinking-budget 实验文档 为主体完整复现并深度讲解如何用 vLLM 实现《s1: Simple test-time scaling》论文提出的思考预算Thinking Budget机制通过迭代生成 强制插入Wait!等特定词把模型的思考 token 数量控制在预算之内。读者将掌握s1.py中每一行核心代码的作用、SamplingParams的配置含义、完整的运行方式以及真实实验中出现的问题与结论。一、背景从 s1 论文看 Test-time Scaling1.1 论文核心思想《s1: Simple test-time scaling》 由李飞飞教授团队提出核心是测试时间缩放Test-time scaling在推理阶段通过调整模型的思考预算即允许模型思考/生成推理链的 token 数量来提高复杂问题的推理效率与准确性。简单来说对于需要多步推理链才能解决的问题模型的性能与思考预算存在正相关关系——当思考预算增加时模型的推理准确率会有明显提升。如上图所示在 MATH500、AIME24、GPQA Diamond 等推理密集型任务上随着平均思考时间的增加模型准确率逐步提升。这也解释了为什么让模型多想一会儿能带来更好的推理效果。1.2 论文中的两个关键技巧论文在实现中还引入了两个重要的设计本仓库实验同样沿用了这些思路1用模型回答对错来判断问题难易。论文通过让 Qwen2.5-32B-Instruct 回答问题来区分任务难度模型答对的是简单问题答错的是复杂问题。从仓库中的相关说明可以看出数据筛选层面也使用了类似策略——利用 Qwen2.5 系列模型过滤掉简单样本从而让思考预算更集中于复杂任务。2插入特定词如Wait!强制模型继续思考。论文做了消融实验探讨在未满足思考预算时插入不同的特定词对模型性能的影响。结果表明插入特定词可以有效地引导模型进行更深入的思考其中WaitWait的效果最好。上表正是论文中关于预算强制外推Budget forcing extrapolation的消融实验Table 4对比了不追加字符串No extrapolation2 倍预算不加字符串追加 Alternatively/Hmm/Wait 等特定词等多种策略。可以看到不同的字符串追加方式对准确率有明显影响这也为后续 vLLM 实现中的Wait!注入提供了理论依据。二、总体思路vLLM 如何实现思考预算vLLM 是一个高性能推理引擎支持大规模语言模型的高效推理。本仓库使用 vLLM 来实现论文中的思考预算机制。整体流程如下对比图左侧是不使用思考预算的推理过程Prompt → 构建输入 → 生成 → 响应右侧是使用思考预算的推理过程Prompt → 构建输入 → 生成 → 检查思考 token 数是否超过预算 → 未超过则在文本末尾追加Wait!重新生成 → 最终响应。可以看到使用思考预算后模型会在推理过程中插入特定词来引导自己进行更深入的思考。环境提示考虑到部分同学配置环境可能会遇到问题作者在 ucloud 平台准备了环境镜像可直接创建 ucloud 实例使用镜像链接见原文档。三、核心代码实现逐步拆解 s1.py完整的可运行代码在 s1.py下面按模块逐步拆解。3.1 依赖与工具函数代码依赖vllm与transformersfrom vllm import LLM, SamplingParams from transformers import AutoTokenizer import time其中LLMvLLM 的模型加载与推理入口SamplingParams采样参数配置对象控制temperature、max_tokens、stop等生成行为AutoTokenizer用于 token 计数与 chat template 构建。构建输入。模型使用 chat template 构造输入build_input使用tokenizer.apply_chat_template处理 system/user 消息并显式开启思考模式def build_input(prompt, tokenizer): messages [ {role: system, content: Please reason step by step, and put your final answer within \\boxed{{}}.}, {role: user, content: prompt} ] input_text tokenizer.apply_chat_template( messages, tokenizeFalse, add_generation_promptTrue, enable_thinkingTrue ) return input_text其中enable_thinkingTrue会启用模型的思考标签如think//think这是后续统计思考 token的基础system 提示要求模型逐步推理并用\boxed{}输出最终答案。token 统计。思考 token 数的统计逻辑是把 prompt 与生成结果拼接后截取think\n之后的内容进行 token 计数def count_thinking_token(outputs, tokenizer): total_token outputs[0].prompt outputs[0].outputs[0].text thinking_token total_token.split(think\n)[-1] thinking_token_id tokenizer(thinking_token)[input_ids] return total_token, len(thinking_token_id) def count_token(string, tokenizer): return len(tokenizer(string)[input_ids])3.2 主函数run_thinking_budget_sampledef run_thinking_budget_sample(llm_model, tokenizer, user_input, thinking_budget): input_text build_input(user_input, tokenizer) input_token_count count_token(input_text, tokenizer) iteration_count 0 max_token input_token_count thinking_budget sampling_params SamplingParams( temperature0.7, max_tokens4096, skip_special_tokensFalse ) think_token_count 0 while True: wait_sampling_params SamplingParams( temperature0.7, max_tokensthinking_budget - think_token_count, stop/think, skip_special_tokensFalse ) outputs llm_model.generate( input_text, wait_sampling_params ) total_token, think_token_count count_thinking_token(outputs, tokenizer) print(f第{iteration_count}次迭代思考token数{think_token_count}) if think_token_count thinking_budget: break input_text total_token \nWait!\n # \nWait a moment. Was there any loophole in my thought just now?!\n # \nWait!\n iteration_count 1 final_outputs llm_model.generate( outputs[0].prompt outputs[0].outputs[0].text \n/think\n, sampling_params ) total_content final_outputs[0].prompt final_outputs[0].outputs[0].text thinking_content total_content.split(think)[-1].split(/think)[0] print(total_content) print(f迭代次数{iteration_count}, 输入token数{input_token_count}, 思考token数{count_token(thinking_content, tokenizer)}, 总token数{count_token(total_content, tokenizer)})第一步参数准备。函数接收模型、tokenizer、用户输入和思考预算thinking_budget即允许的思考 token 上限四个参数。先构建输入文本并计算输入的 token 数量同时为最终答案生成准备好独立的sampling_params。因为max_tokens参数表示生成的最大 token 数量所以需要把思考预算转换成每次迭代的生成上限每次迭代允许模型新生成的 token 数为thinking_budget - think_token_count总预算减去已思考的 token 数。同时需要在SamplingParams中设置stop/think这样模型在生成到/think时会自动停止便于分轮统计思考内容。wait_sampling_params SamplingParams( temperature0.7, max_tokensthinking_budget - think_token_count, stop/think, skip_special_tokensFalse )第二步循环生成与 Wait! 注入。核心循环逻辑如下while True: wait_sampling_params SamplingParams( temperature0.7, max_tokensthinking_budget - think_token_count, stop/think, skip_special_tokensFalse ) outputs llm_model.generate( input_text, wait_sampling_params ) total_token, think_token_count count_thinking_token(outputs, tokenizer) print(f第{iteration_count}次迭代思考token数{think_token_count}) if think_token_count thinking_budget: break input_text total_token \nWait!\n # \nWait a moment. Was there any loophole in my thought just now?!\n # \nWait!\n iteration_count 1每次迭代中用当前已思考 token 数更新max_tokens上限调用llm_model.generate生成一段思考内容统计累计思考 token 数并打印日志若累计思考 token 数超过思考预算则跳出循环否则把已生成的完整文本作为新的输入并在末尾追加\nWait!\n引导模型继续深入思考。注释中还保留了两个可替换的提示词变体\nWait a moment. Was there any loophole in my thought just now?!\n与\nWait!\n方便读者对照论文消融实验做不同字符串的 A/B 测试。从源码结构看仓库最终选用了更简洁的\nWait!\n。第三步生成最终答案。当思考达到预算后模型还需要把思考过程总结成最终答案这一步同样需要单独的采样参数max_tokens设置为 4096。如原文档所述模型根据思考过程进行总结得出答案也需要很多 token这个值设置为多少都可以通常设置为一个较大的值即可sampling_params SamplingParams( temperature0.7, max_tokens4096, skip_special_tokensFalse )随后在已生成的思考内容末尾补上\n/think\n闭合思考标签再次调用llm_model.generate生成最终答案final_outputs llm_model.generate( outputs[0].prompt outputs[0].outputs[0].text \n/think\n, sampling_params ) total_content final_outputs[0].prompt final_outputs[0].outputs[0].text thinking_content total_content.split(think)[-1].split(/think)[0] print(total_content) print(f迭代次数{iteration_count}, 输入token数{input_token_count}, 思考token数{count_token(thinking_content, tokenizer)}, 总token数{count_token(total_content, tokenizer)})最后打印完整输出并输出关键指标迭代次数、输入 token 数、思考 token 数、总 token 数。s1.py中还额外实现了把结果写入output_{int(time.time())}.txt文件的逻辑便于留存实验记录。3.3 对照实验run_sample无思考预算为了对比思考预算的效果s1.py还提供了一个不使用思考预算的基准函数run_sample直接以max_tokens32768生成不做Wait!注入也不做思考 token 控制def run_sample(llm_model, tokenizer, user_input): input_text build_input(user_input, tokenizer) input_token_count count_token(input_text, tokenizer) sampling_params SamplingParams( temperature0.7, max_tokens32768, skip_special_tokensFalse ) final_outputs llm_model.generate( input_text, sampling_params ) total_content final_outputs[0].prompt final_outputs[0].outputs[0].text thinking_content total_content.split(think)[-1].split(/think)[0] print(total_content) print(f输入token数{input_token_count}, 思考token数{count_token(thinking_content, tokenizer)}, 总token数{count_token(total_content, tokenizer)})从源码结构看该函数与run_thinking_budget_sample的差异正是思考预算机制的全部差异前者一次生成到底后者分轮迭代 强制注入Wait!。3.4 主程序入口与模型配置if __name__ __main__: model_path /model/ModelScope/Qwen/Qwen3-14B tokenizer AutoTokenizer.from_pretrained(model_path) llm LLM( modelmodel_path, gpu_memory_utilization0.9, trust_remote_codeTrue ) print( 思考预算采样 ) run_thinking_budget_sample( llm_modelllm, tokenizertokenizer, user_inputThere are exactly three positive real numbers $ k $ such that the function\n$ f(x) \\frac{(x - 18)(x - 72)(x - 98)(x - k)}{x} $\n defined over the positive real numbers achieves its minimum value at exactly two positive real numbers $ x $. Find the sum of these three values of $ k $., thinking_budget32768 )主程序中的关键配置配置项值说明model_path/model/ModelScope/Qwen/Qwen3-14B实验所用模型为 Qwen3-14B约 14B 参数通过 ModelScope 路径加载gpu_memory_utilization0.9允许 vLLM 使用 90% 的 GPU 显存trust_remote_codeTrue信任远端代码允许加载模型自定义实现thinking_budget32768思考预算为 32768 个 token测试题目一道求三正实数 k 之和的数学题属于需要长推理链的复杂问题适合验证思考预算效果注意model_path为作者实验环境的本地路径实际使用时需替换为自己环境中模型的实际存放路径并确保 GPU 显存足够gpu_memory_utilization0.9意味着显存占用较高。四、结果分析思考预算的实际效果与局限性使用思考预算后模型在推理过程中能够更深入地思考问题从而提高推理效率和准确性。从 output 目录 中的实验记录可以看到真实运行结果输出文件迭代次数输入 token 数思考 token 数总 token 数模型给出的答案output_1754208752.txt8109327853375545 58 85 188output_1754209653.txt17109327723369744 126 152 322从运行记录看两次实验在 32768 的思考预算下分别迭代了 8 次和 17 次思考 token 数最终均超过了预算分别为 32785 和 32772总 token 数维持在 3.3 万左右。这也印证了文档中的观察模型在思考过程中可能出现重复生成相同内容导致思考 token 数量超过思考预算的情况。此外实验还发现了一些有趣的现象1Wait!未必触发真正的反思。在某些情况下就算插入了Wait!模型并不会按照论文中所示进行多种不同方式的解答尝试或是反思之前的思考过程是否正确。2重复思考导致预算超支。模型会在思考过程中重复生成相同的内容导致思考 token 数量超过思考预算。如上图所示模型的思考过程中会出现Wait a moment. Was there any loophole in my thought just now?!这类循环式自我检查语句——虽然表面上在反思但实际可能只是在重复确认已有的思路而非真正换一种思路求解。3强插特定词可能一条道走到黑。经过测试强行使用特定词如Wait!来引导模型进行更深入的思考可能会促使模型产生一条道走到黑的想法——即沿原有错误思路继续深入而不是跳出来换一个方向。当然也有一个客观原因值得考虑本实验使用的模型只有 14B 参数Qwen3-14B思考过程中的推理能力可能受到模型规模限制。五、可复现实验建议基于以上分析读者在实际使用中可以参考以下建议调整思考预算thinking_budget可以根据任务复杂度调整简单问题可以设小如 20488192复杂数学推理题可以设大如本实验的 32768。论文实验中的思考预算一般取 512、1024、2048、4096、8192 等档位。替换注入词代码注释中提供了\nWait a moment. Was there any loophole in my thought just now?!\n等变体可以对照论文消融实验对比不同特定词的效果论文结论是WaitWait效果最好。更换模型与显存主程序中gpu_memory_utilization0.9适用于显存充足的环境模型路径需替换为本地实际路径。从output记录的运行结果看14B 模型在长思考预算下仍可能出现重复思考若追求更高质量的推理可尝试更大规模的模型。保存实验记录s1.py已内置将输出与统计信息写入output_{时间戳}.txt的逻辑建议每次实验后保留输出文件便于对比不同配置的效果。六、小结本文以 s1-vllm-thinking-budget 文档 为主线结合 s1.py 完整源码系统讲解了Test-time scaling 与思考预算的理论来源源于 s1 论文通过增加思考 token 提升复杂推理任务的准确率vLLM 的完整实现方案迭代生成 stop/thinkWait!注入 预算检查的闭环流程以及两个SamplingParams各自的作用真实实验的量化结果与观察817 次迭代、约 3.3 万思考 token 的完整记录以及重复思考一条道走到黑等局限性。思考预算机制的关键价值在于它把让模型多想一会儿从不可控的直觉变成可配置、可度量的工程参数。读者可以基于本仓库代码继续调整预算、注入词和模型规模进一步探索 Test-time scaling 的上限。【免费下载链接】happy-llm 从零开始构建大模型项目地址: https://gitcode.com/GitHub_Trending/ha/happy-llm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考