开源 AI 食谱文档

扩展 LLM 的测试时计算以进行更长时间的思考

开源 AI 食谱

加入 Hugging Face 社区

并获得增强的文档体验

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

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

切换文档主题

开始使用

扩展 LLM 的测试时计算以进行更长时间的思考

作者：Sergio Paniego

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

在本食谱中，我们将指导您使用测试时计算扩展 Instruct LLM 系统的推理时间，以解决更具挑战性的问题，例如复杂的数学问题。这种方法受到 OpenAI o1-o3 模型 的启发，表明在推理过程中更长的推理时间可以提高模型性能。

这项技术建立在这篇博文中分享的实验基础上，该博文表明，较小的模型（如 1B 和 3B Llama Instruct 模型）在 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 组成，该 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

我们将首先使用 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. 多样化验证树搜索 (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 进行推理后训练→