开源 AI 食谱文档

扩展测试时计算以实现 LLM 更长时间的思考

开源 AI 食谱

加入 Hugging Face 社区

并获得增强的文档体验

在模型、数据集和 Spaces 上进行协作

通过加速推理获得更快的示例

切换文档主题

开始使用

扩展测试时计算以实现 LLM 更长时间的思考

作者: Sergio Paniego

🚨 警告：此笔记本是资源密集型的，需要大量计算能力。如果您在 Colab 中运行，它将使用 A100 GPU。

在本指南中，我们将引导您使用测试时计算来延长指令式 LLM 系统的推理时间，以解决更具挑战性的问题，例如复杂的数学问题。这种方法的灵感来源于 OpenAI o1-o3 模型，它表明推理过程中更长的推理时间可以提升模型性能。

这项技术建立在这篇博客文章中分享的实验基础上，这些实验表明，像 1B 和 3B Llama 指令模型这样的小模型，在给予足够的“思考时间”后，可以在 MATH-500 基准测试中超越大得多的模型。来自 DeepMind 的最新研究表明，通过迭代自我修正或使用奖励模型等策略，可以优化测试时计算。

该博客文章介绍了一个用于运行这些实验的新仓库。在本指南中，我们将专注于构建一个小型聊天机器人，它将通过更长时间的推理来使用小型开源模型解决更难的问题。

1. 安装依赖项

让我们首先安装 search-and-learn 仓库！🚀
这个仓库是为复现实践结果而设计的，不是一个 Python pip 包。但是，我们仍然可以用它来生成我们的系统。为此，我们需要按照以下步骤从源代码安装：

!git clone https://github.com/huggingface/search-and-learn

%cd search-and-learn
!pip install -e '.[dev]'

from huggingface_hub import notebook_login

notebook_login()

2. 设置大语言模型 (LLM) 和过程奖励模型 (PRM) 💬

如图所示，该系统由一个 LLM（根据用户输入生成中间答案）、一个PRM 模型（评估这些答案并打分），以及一个搜索策略（利用 PRM 反馈指导搜索过程的后续步骤，直到得出最终答案）组成。

让我们首先初始化每个模型。对于 LLM，我们将使用 meta-llama/Llama-3.2-1B-Instruct 模型，对于 PRM，我们将使用 RLHFlow/Llama3.1-8B-PRM-Deepseek-Data 模型。

system

import torch
from vllm import LLM
from sal.models.reward_models import RLHFFlow

model_path = "meta-llama/Llama-3.2-1B-Instruct"
prm_path = "RLHFlow/Llama3.1-8B-PRM-Deepseek-Data"

llm = LLM(
    model=model_path,
    gpu_memory_utilization=0.5,  # Utilize 50% of GPU memory
    enable_prefix_caching=True,  # Optimize repeated prefix computations
    seed=42,  # Set seed for reproducibility
)

prm = RLHFFlow(prm_path)

2.1 实例化问题、搜索策略并调用流水线

既然我们已经设置了 LLM 和 PRM，接下来我们将定义问题，选择一个搜索策略来检索相关信息，并调用流水线来处理问题。

实例化问题：在这一步中，我们定义系统将要回答的输入问题，并考虑给定的上下文。
搜索策略：系统目前支持以下搜索策略：best_of_n、beam_search 和 dvts（见图）。在本例中，我们将使用 best_of_n，但你可以根据需要轻松切换到其他任何策略。我们需要为搜索策略的配置定义一些参数。你可以在这里查看完整列表。
调用流水线：有了问题和搜索策略，我们将调用推理流水线，通过 LLM 和 PRM 处理输入，以生成最终答案。

第一步是明确定义系统将要回答的问题。这确保了我们有一个精确的任务让模型去解决。

question_text = "Convert the point $(0,3)$ in rectangular coordinates to polar coordinates.  Enter your answer in the form $(r,\theta),$ where $r > 0$ and $0 \le \theta < 2 \pi.$"
input_batch = {"problem": [question_text]}

接下来，我们定义配置，包括候选答案数量 (N) 等参数，并选择将要使用的搜索策略。搜索策略决定了我们如何探索潜在的答案。在本例中，我们将使用 best_of_n。

有了问题和配置，我们使用选定的搜索策略生成多个候选答案。这些候选答案根据其相关性和质量进行评估，并返回最终答案。

from sal.config import Config
from sal.search import beam_search, best_of_n, dvts

config = Config()
config.n = 32  # Number of answers to generate during the search

search_result = best_of_n(x=input_batch, config=config, llm=llm, prm=prm)

2.2 显示最终结果

一旦流水线通过 LLM 和 PRM 处理了问题，我们就可以显示最终结果。这个结果是模型在考虑了中间答案并使用 PRM 对它们进行评分后得出的输出。

以下是显示最终答案的方法：

search_result["pred"][0]

模型的输出可能包含特殊标记，例如 <|start_header_id|> 或 <|end_header_id|>。为了使答案更易读，我们可以在显示给最终用户之前安全地移除它们。

formatted_output = search_result["pred"][0].replace("<|start_header_id|>assistant<|end_header_id|>\n\n", "").strip()
formatted_output

移除任何特殊标记后，我们可以向用户显示最终答案。由于答案是基于 markdown 的，因此可以作为 markdown 正确地渲染显示。

from IPython.display import display, Markdown

display(Markdown(formatted_output))

3. 全部组装起来！🧑‍🏭️

现在，让我们创建一个封装整个流水线的方法。这将使我们能够在未来的应用程序中轻松重用该过程，使其高效且模块化。

通过结合 LLM、PRM、搜索策略和结果显示，我们可以简化工作流程，并确保其可用于其他任务或问题。

我们简化了工作流程，确保它可以重复用于不同的任务或问题。此外，我们将跟踪每个方法所花费的时间，以便我们能够理解使用每种策略和配置的实际影响。

以下是我们构建该方法的方式：

import time


def generate_with_search_and_learn(question, config, llm, prm, method="best_of_n"):
    """
    Generate an answer for a given question using the search-and-learn pipeline.

    Args:
    - question (str): The input question to generate an answer for.
    - config (Config): Configuration object containing parameters for search strategy.
    - llm (LLM): Pretrained large language model used for generating answers.
    - prm (RLHFFlow): Process reward model used for evaluating answers.
    - method (str): Search strategy to use. Options are 'best_of_n', 'beam_search', 'dvts'. Default is 'best_of_n'.

    Returns:
    - str: The formatted output after processing the question.
    """
    batch = {"problem": [question]}

    start_time = time.time()
    if method == "best_of_n":
        result = best_of_n(x=batch, config=config, llm=llm, prm=prm)
    elif method == "beam_search":
        result = beam_search(examples=batch, config=config, llm=llm, prm=prm)
    elif method == "dvts":
        result = dvts(examples=batch, config=config, llm=llm, prm=prm)

    elapsed_time = time.time() - start_time
    print(f"\nFinished in {elapsed_time:.2f} seconds\n")

    tokenizer = llm.get_tokenizer()
    total_tokens = 0
    for completion in result["completions"]:
        for comp in completion:
            output_tokens = tokenizer.encode(comp)
            total_tokens += len(output_tokens)

    print(f"Total tokens in all completions: {total_tokens}")

    formatted_output = result["pred"][0].replace("<|start_header_id|>assistant<|end_header_id|>\n\n", "").strip()
    return formatted_output

⏳ 3.1 比较每种策略的思考时间

让我们比较三种方法的思考时间：best_of_n、beam_search 和 dvts。每种方法都使用相同数量的答案进行评估，测量以秒为单位的思考时间和生成的令牌数。

在下面的结果中，best_of_n 方法的思考时间最短，而 dvts 方法耗时最长。然而，由于其搜索策略更简单，best_of_n 生成了更多的令牌。

方法	搜索过程中的答案数量	思考时间 (秒)	生成的令牌数
best_of_n	8	3.54	3087
beam_search	8	10.06	2049
dvts	8	8.46	2544

这种比较说明了不同策略之间的权衡，平衡了思考时间和搜索过程的复杂性。

1. Best of n (最优 n 选)

我们首先使用 best_of_n 策略。以下是如何跟踪此方法的思考时间：

>>> question = "Convert the point $(0,3)$ in rectangular coordinates to polar coordinates.  Enter your answer in the form $(r,\theta),$ where $r > 0$ and $0 \le \theta < 2 \pi.$"

>>> config.n = 8

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="best_of_n"
... )

Finished in 3.54 seconds

Total tokens in all completions: 3087

display(Markdown(formatted_output))

2. Beam Search (束搜索)

现在，让我们尝试使用 beam_search 策略。

>>> config.n = 8
>>> # beam search specific
>>> config.sort_completed = True
>>> config.filter_duplicates = True

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="beam_search"
... )

Finished in 10.06 seconds

Total tokens in all completions: 2049

display(Markdown(formatted_output))

3. Diverse Verifier Tree Search (DVTS) (多样性验证树搜索)

最后，让我们尝试 dvts 策略。

>>> config.n = 8
>>> # dvts specific
>>> config.n_beams = config.n // config.beam_width

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="dvts"
... )

Finished in 8.46 seconds

Total tokens in all completions: 2544

display(Markdown(formatted_output))

🙋 3.2 使用简单问题测试系统

在最后一个例子中，我们将使用一个简单的问题来测试系统，观察它在简单情况下的表现。这可以帮助我们验证系统即使对于基本查询也能按预期工作。

让我们尝试以下问题：

>>> question = "What's the capital of Spain?"

>>> config.n = 32

>>> formatted_output = generate_with_search_and_learn(
...     question=question, config=config, llm=llm, prm=prm, method="best_of_n"
... )

Finished in 1.03 seconds

Total tokens in all completions: 544

display(Markdown(formatted_output))

即使我们设置了更多的候选答案 (N)，思考时间仍然相对较短 (1.03 秒，生成 544 个令牌)。这表明系统能够高效地处理较简单的问题，花费较少的时间，同时利用其增强的能力来解决更复杂的问题。

🏆 我们现在有了一个功能齐全的流水线，它利用了测试时计算，使系统能够为更复杂的查询“思考更长时间”，同时也为简单问题保持快速的响应时间。

这种方法确保系统可以根据任务的复杂性调整其思考时间，为简单和具有挑战性的问题提供高效且响应迅速的解决方案。

4. 继续探索与资源 🧑‍🎓️

如果你渴望继续探索，请务必查看原始的实验性博客以及其中提到的所有参考文献。这些资源将加深你对测试时计算、其优势及其在 LLM 中应用的理解。

祝您学习和实验愉快！🚀

< > 在 GitHub 上更新

←HF Spaces 上的 Phoenix 可观测性仪表板在 TRL 中使用 GRPO 对 LLM 进行推理后训练→