菜单
主题
vLLM Logo

彻底消除重分词漂移：通过兼容 OpenAI 的 API 返回 Token ID 对智能体强化学习至关重要

9 分钟阅读
Agent Lightning (AGL) 团队

TL;DR（摘要）。 智能体通常通过 OpenAI 兼容的端点调用大语言模型（LLM），这些端点此前仅返回基于字符串的输入和输出。在智能体强化学习（Agent RL）中，由于我们称之为重分词漂移（Retokenization Drift）的现象，这可能导致训练和推理之间的不一致。该现象产生的原因是：推理时文本被反分词（Detokenized），而在训练时又被重新分词（Retokenized）；即使对应的字符串相同，两组 Token 可能并不一致。现在，你可以请求 vLLM 的 OpenAI 兼容端点返回提示词和生成的回复的精确 Token ID。只需将 "return_token_ids": true 传递给 /v1/chat/completions 或 /v1/completions，你就会在常规文本输出之外收到 prompt_token_ids 和 token_ids。这使得智能体强化学习更加稳健，不再发生漂移。这与 Agent Lightning 完美契合，在 Agent Lightning 中，每次模型调用都被视为独立的更新样本，无需拼接；只需通过启用的 return_token_ids 记录返回的 ID 即可。

链接

为什么 Token ID 对智能体强化学习（Agent RL）至关重要

大模型的强化学习是在 Token 序列上进行训练的，因此训练器需要行为策略采样出的精确 Token ID。在单轮设置中，这通常很简单，因为调用 vLLM 的底层 generate 方法会直接返回 Token。

在智能体设置中，大多数智能体框架都会调用 OpenAI 风格的 chat.completions / completions 接口。智能体更青睐这些 API 而非原始的 generate，因为它们提供了智能体架构所需的高级功能，例如聊天模板与角色（系统/用户/助手）、工具/函数调用、结构化输出等。这些 API 过去仅返回字符串，这可能会在智能体强化学习中引发问题。此前，存储的文本必须在训练期间重新分词，但这在实践中是不稳定且不够精确的，原因就是重分词漂移。

在强化学习中，你会看到的症状包括：不稳定的学习曲线（如下图所示），以及你认为优化所基于的数据与模型实际采样的数据之间存在难以调试的偏差。


红色和蓝色线条是在相同设置下（即在训练中存储文本并进行重分词）得到的，而黄色线条直接使用了来自推理引擎的 Token。

这种漂移可能由以下三个原因引起。

  • 非唯一的 "HAVING"。在实践中这会持续出现：一个单词在生成时可能被拆分为两个 Token（例如 H + AVING），但当你稍后在训练期间对该文本进行重新分词时，你可能会得到不同的拆分（例如 HAV + ING）。文本看起来一样，但 ID 不同，导致学习器针对错误的序列进行了优化。

"HAVING" 这个词对应着不同的 Token。
  • 工具调用序列化。生成的工具调用文本（如 <tool_call>{ "name": ... }</tool_call>）会被工具调用解析器解析为聊天完成 API 所需的对象。随后，该对象被渲染回 <tool_call>{ "name": ... }</tool_call> 并再次进行重分词。工具调用解析和重新渲染可能会导致空白字符和格式的变化。在某些情况下，JSON 错误甚至会被工具调用解析器自动纠正。这掩盖了模型真实的生成错误，阻碍了通过训练消除这些错误。
  • 聊天模板差异。不同框架中使用的聊天模板可能略有不同。例如，同一个 LLaMA 模型可以配合多个聊天模板使用（vLLM 中有多个，而 HuggingFace 中有一个）。当在推理和训练中使用不同的框架时，这种差异会产生不同的 Token。

这三个因素导致了重分词漂移，进而导致训练不稳定，这可能是因为它们导致了推理和训练之间的不一致，进而造成了离策略（Off-policy）强化学习更新。同策略（On-policy）对于稳定的强化学习训练至关重要，微妙的变化也会产生巨大的影响。由重分词漂移引起的离策略效应甚至不是在 Token 级别上的，因此无法通过 Token 级别的重要度采样来纠正。

另一种选择是保存模型生成的 Token ID，正如在单轮设置中所做的那样。这要求智能体必须在 Token 级别与推理引擎进行通信。然而，大多数智能体——尤其是那些使用 LangChain 等框架构建的智能体——依赖于兼容 OpenAI 的 API，无法自行进行分词或反分词。关于这一部分的更多讨论，请参见此处。

解决方案与新功能

更好的解决方案是使用直接返回 Token ID 的 OpenAI 兼容 API。Agent Lightning 和 vLLM 团队合作将此功能直接加入到了 vLLM 核心代码中。从 vLLM v0.10.2 开始，OpenAI 兼容 API 包含了一个 return_token_ids 参数，允许在请求聊天消息时同时请求 Token ID。当你在请求中将其设置为 true 时，响应将包含两个额外字段

  • prompt_token_ids：输入的 Token ID（经过任何聊天模板处理后），以及
  • token_ids：为补全生成的 Token ID，通过 completion.choices 传递。

响应中的其他所有内容保持与 OpenAI 兼容，因此现有的客户端可以继续正常工作。

Agent Lightning (v0.2) 简介

在 Agent Lightning（简称 AGL）的初始版本（v0.1）中，我们为任何带有强化学习的智能体提供了一个灵活的训练框架。它具有几个核心特性。

  • 与现有智能体无缝集成，且无需修改代码（几乎如此）！
  • 可以使用任何智能体框架构建（LangChain、OpenAI Agent SDK、Microsoft Agent Framework 等）；甚至无需智能体框架（直接使用 Python 程序）。
  • 对 LLM 的输入没有约束，允许灵活的编排，如摘要、多智能体协作和其他复杂工作流。

当 Agent Lightning 首次发布时，我们实现了一个 插桩 vLLM 服务器，该服务器通过猴子补丁（monkey-patch）方式让 vLLM 的 OpenAI 服务器返回 Token ID。现在，AGL 会自动为每个请求添加 return_token_ids，以便引擎在响应中包含 Token ID。然后，利用 AGL 中嵌入的 追踪（Tracing）功能，我们自动收集训练端所需的数据，包括这些 Token ID。

智能体优化的中间件

从更精确的数据收集这一视角出发，在 v0.2 中，我们明确了 AGL 在智能体优化中的作用。从概念上讲，Agent Lightning（即 AGL）引入了一个可持续的中间件层和标准化的数据协议，用于智能体优化，特别是智能体强化学习。


Agent Lightning 概念概览。

Agent Lightning 设计有一组模块化、可互操作的组件，共同实现可扩展且高效的智能体强化学习。每个组件通过标准化的数据协议和定义良好的接口进行通信，发挥独特作用。

  • 智能体运行器（Agent Runner） — 负责执行智能体以完成指定的任务。它接收任务，将其委派给智能体执行，收集结果和中间数据，并将这些报告回数据存储。智能体运行器与 LLM 侧分开运行，因此可以托管在不同的资源（例如 CPU）上，并可水平扩展以支持大量并发智能体实例。
  • 算法（模型训练器） — 托管用于推理和训练的大语言模型（LLM）。该组件编排整个强化学习循环，包括任务采样、展开（Rollout）管理和基于收集到的经验数据进行模型更新。它通常在 GPU 资源上运行，并与智能体运行器通过共享数据协议进行异步交互。
  • 数据存储（Data Store） — 作为智能体强化学习生态系统中管理所有数据交换和存储的中央仓库。它提供标准化接口和统一数据模式，以确保异构组件之间的互操作性。通过这种设计，算法和智能体运行器可以间接而有效地进行通信，从而实现灵活且可扩展的协作。例如，使用标准化的 rollouts，算法可以异步地将任务委派给智能体运行器，由运行器执行这些任务并通过 spans 数据结构报告执行轨迹。

Agent Lightning 中的训练循环。

在这种以数据存储为中心的设计理念下，所有智能体训练迭代被抽象为两个步骤。第一步是收集智能体运行数据（AGL 中的 spans）并将其存储在数据存储中；第二步是从存储中检索所需数据并将其发送到算法侧进行训练。

这种抽象视角带来了几个优势。首先，它提供了更高的算法灵活性：数据收集可以依赖各种追踪器或发出自定义消息，使得定义不同奖励或捕获任何中间变量变得简单。在算法侧，可以通过 query spans 访问所需数据，自定义 适配器（adapters）可实现自由的数据转换。

该设计还支持算法自定义，如信用分配、使用部分数据的辅助模型学习、通过数据调整进行训练改进等。此外，在此框架内，我们可以扩展到更多种类的算法，例如自动提示词优化（APO）、筛选高奖励数据并通过 Unsloth 进行拟合。

这种设计的第二个主要优势在于它能够通过模块化分离降低整体系统复杂性，同时使不同组件能够利用不同的资源和优化策略。智能体强化学习系统本质上是复杂的，因为它利用动态的、环境驱动的交互来使模型从经验数据中持续学习。一个典型的智能体强化学习栈包含几个关键组件，包括智能体框架（例如 LangChain、MCP）、LLM 推理引擎（例如 vLLM）和训练框架（例如 Megatron-LM）。如果没有解耦架构，这些组件的独立性和异构性会导致巨大的系统复杂性。相比之下，解耦设计允许系统适应多样的资源需求：例如，智能体侧可能需要更高的 CPU 容量，而 LLM 推理和训练通常是 GPU 密集型的。这种模块化结构还有助于每个组件的独立水平扩展，从而提高效率和可维护性。

更多资料

训练愉快！⚡

致谢

我们要向 vLLM 维护者表示诚挚的感谢，包括 Kaichao You、Nick Hill、Aaron Pham、Cyrus Leung、Robert Shaw 和 Simon Mo。没有他们的支持与协作，此次集成是不可能实现的。

Agent Lightning 是来自微软研究院的一个开源项目。我们衷心感谢 MSR 对这项开源探索的支持。Yuge Zhang 是这项工作的主要贡献者。