No More Retokenization Drift: Returning Token IDs via the OpenAI Compatible API Matters in Agent RL

发表时间: 2025-10 · Blog post by The Agent Lightning (AGL) Team (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2025-10-22-agent-lightning

作者/机构：The Agent Lightning (AGL) Team (Microsoft Research)

速读

一句话结论 本文通过在 vLLM 的 OpenAI 兼容 API 中原生支持返回精确的 token ID，并结合 Agent Lightning 框架的标准化数据中间件设计，彻底消除了智能体强化学习中因文本重分词导致的训练与推理不一致问题，实现了稳定且可扩展的策略更新。

要解决什么问题 在大语言模型的强化学习中，训练算法需要模型行为策略实际采样的精确 token 序列。在单轮对话场景下，直接调用底层生成接口即可获取 token，但在智能体强化学习（Agent RL）中，现有的智能体框架（如 LangChain）高度依赖 OpenAI 兼容的高级 API 来处理聊天模板、角色分配和工具调用等复杂逻辑。这些高级 API 过去只返回字符串文本，导致系统在训练阶段必须对存储的文本进行重新分词（retokenize）。这种做法会引发严重的“重分词漂移”（Retokenization Drift）现象，即推理时生成的 token 序列与训练时重新切分出的 token 序列在底层 ID 上不一致。导致该卡点的机制主要有三点：第一是分词的不唯一性，例如模型生成时将单词切分为 H 和 AVING 两个 token，但重新分词时却按 HAV 和 ING 切分，导致目标序列错误；第二是工具调用序列化问题，工具解析器在解析和重新渲染工具调用标签时，可能会改变空格、格式，甚至自动纠正 JSON 错误，从而掩盖了模型真实的生成错误，阻止了算法通过训练消除这些错误；第三是聊天模板差异，不同框架（如 vLLM 和 HuggingFace）使用的模板微小差异会产生不同的 token。这种不一致性会导致强化学习更新严重偏离当前策略（off-policy），且这种偏差不是 token 级别的，无法通过常规的 token 级别重要性采样来修正，最终造成学习曲线极不稳定、模型发散，以及难以调试的数据分歧。

怎么做的 核心思路是让智能体在与推理引擎交互时，直接获取并保存模型底层实际采样的 token ID，从而完全绕过训练时的重新分词步骤。为了在不破坏智能体框架对高级 API 依赖的前提下实现这一点，作者与 vLLM 团队合作，在 vLLM 的 OpenAI 兼容 API 中新增了 return_token_ids 参数。当请求中携带该参数时，API 会在常规文本输出之外，额外返回 prompt_token_ids（经过聊天模板处理后的输入 token ID）和 token_ids（生成的 token ID）。在此基础上，Agent Lightning 框架被设计为一个用于智能体优化的中间件层，通过三个关键部件实现高效的数据流转。第一个部件是 Agent Runner（智能体运行器），负责接收任务并执行智能体，它在调用大语言模型时自动附加 return_token_ids 参数，并通过内置的追踪机制收集包含精确 token ID 的执行轨迹（spans），该部件可部署在 CPU 上进行大规模横向扩展。第二个部件是 Data Store（数据存储），作为中心枢纽提供标准化的数据协议，例如使用 rollouts 结构接收算法分配的任务，并接收存储 Agent Runner 收集的执行轨迹，确保异构组件间的互操作性。第三个部件是 Algorithm（算法训练器），部署在 GPU 上，负责管理强化学习循环，它通过查询 Data Store 获取精确的经验数据来执行模型更新。这种解耦设计不仅确保了训练数据与推理行为的绝对一致，还允许数据收集和模型训练异步进行，解决了智能体端需要高 CPU 并发而模型端需要密集 GPU 算力的资源错配问题，大幅降低了系统的整体耦合复杂度。

效果如何 文档通过对比强化学习的学习曲线展示了该方法对训练稳定性的决定性影响。对比基线为传统的“在训练中保存文本并重新分词”路线（对应实验中的红线和蓝线）。实验结果表明，采用基线方法时，由于重分词漂移导致的偏离策略更新，学习曲线表现出明显的震荡和发散；而采用本文方法，即直接使用推理引擎返回的 token ID（对应实验中的黄线）后，学习曲线变得平滑且稳定，成功消除了训练数据与模型实际采样数据之间的差异。该方案能以几乎零代码修改的代价无缝集成到现有的各类智能体框架中，并且没有对大语言模型的输入施加任何限制。基于这种稳定的底层支持，该框架进一步证明了其在复杂任务中的有效性，能够支持使用 verl 训练多智能体编排的 SQL 智能体、使用自动提示词优化（APO）训练房间选择智能体，以及使用 Unsloth 对结合了模型上下文协议（MCP）的数学智能体进行微调。文档未提供具体的模型参数量、训练数据规模或硬件配置等量化实验细节，也未提及开启该功能后可能带来的额外网络传输开销或存储代价。

1. 主要贡献

在智能体强化学习（Agent RL）领域，智能体通常通过兼容OpenAI的端点（endpoints）调用大语言模型（LLM）。过去，这些端点仅返回基于字符串的输入和输出。本文指出了这一机制在强化学习中引发的核心问题：重分词漂移（Retokenization Drift）。该现象发生的原因是，Token在推理期间被解码（detokenized）为字符串，随后在训练期间又被重新分词（retokenized）；即使对应的字符串完全相同，这两组Token也可能存在差异，从而导致推理与训练之间的严重不一致。

为了解决这一问题，本文提出了以下创新点与解决方案：

  1. vLLM底层特性的支持：现在可以请求vLLM的OpenAI兼容端点为提示词（prompts）和生成的响应返回精确的Token IDs。通过向/v1/chat/completions或/v1/completions传递"return_token_ids": true参数，用户可以在常规文本输出之外，接收到prompt_token_ids和token_ids。
  2. 框架级集成消除漂移：这一新特性与Agent Lightning框架完美契合。在该框架中，每次模型调用都被视为一个独立的更新样本而无需拼接；只需在启用return_token_ids的情况下记录返回的ID，即可彻底消除重分词漂移，使智能体强化学习变得高度鲁棒。

2. 背景知识与关键观察

Token ID在智能体强化学习中的重要性
大语言模型（LLM）的强化学习是在Token序列上进行训练的，因此训练器需要行为策略（behavior policy）采样出的精确Token ID。在单轮（single-turn）对话设置中，这曾经非常直接，因为调用vLLM的底层generate函数会直接返回Token。

智能体框架的API偏好与重分词漂移的产生
在智能体设置中，大多数智能体框架调用的是OpenAI风格的chat.completions或completions接口。智能体更倾向于使用这些API，因为它们提供了智能体技术栈所依赖的高级功能，例如聊天模板与角色（系统/用户/助手）、工具/函数调用、结构化输出等。然而，这些API历史上仅返回字符串，这在智能体强化学习中引发了问题。以前，存储的文本必须在训练期间重新分词（retokenized），但由于重分词漂移（retokenization drift）的存在，这种做法在实践中既不稳定也不准确。

重分词漂移在强化学习中的症状
在强化学习中观察到的症状包括：不稳定的学习曲线（如下图所示），以及你认为用于优化的数据与模型实际采样的数据之间出现难以调试的分歧。
重分词漂移导致的学习曲线不稳定
图中红线和蓝线是在相同设置下获得的（即在训练中存储文本并在训练时重新分词），而黄线直接使用了来自推理引擎的Token。

导致重分词漂移的三个核心原因
漂移可能由以下三个原因引起：

  • 非唯一的“HAVING”现象：在实践中这种情况经常出现：在生成过程中，一个单词可能是作为两个Token生成的（例如，H + AVING），但是当你后来在训练期间对文本重新分词时，你会得到一个不同的拆分方式（例如，HAV + ING）。文本看起来完全相同，但ID却不同，这使得你的学习器针对错误的序列进行了优化。
    单词“HAVING”对应不同的Token

  • 工具调用的序列化（Tool-call serialization）：生成的工具调用文本（如<tool_call>{ "name": ... }</tool_call>）会被工具调用解析器解析为聊天补全API所需的对象。随后，该对象被重新渲染回<tool_call>{ "name": ... }</tool_call>并再次重新分词。工具调用的解析和重新渲染可能会导致空格和格式的变化。在某些情况下，JSON错误甚至可能被工具调用解析器自动纠正。这掩盖了模型真正的生成错误，并阻止了通过训练消除这些错误。

  • 聊天模板的差异（Chat template difference）：不同框架中使用的聊天模板可能会有细微差别。例如，同一个LLaMA模型可以与多个聊天模板配合使用（在vLLM中有多个模板，在HuggingFace中有一个模板）。当在推理和训练中使用不同的框架时，这种分歧将产生不同的Token。

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

替代方案及其局限性
替代方案是像单轮设置中那样，保存模型生成的Token ID。这要求智能体必须在Token级别与推理引擎进行通信。然而，大多数智能体——尤其是那些使用LangChain等框架构建的智能体——依赖于兼容OpenAI的API，并且无法自行进行分词或解码。

3. 方法细节

返回Token ID的解决方案与新特性
为了彻底解决重分词漂移问题，一个更优的解决方案是使用直接返回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：为补全（completion）生成的Token ID，通过completion.choices传递。
    响应中的其他所有内容都保持与OpenAI兼容，因此现有的客户端可以继续正常工作。

Agent Lightning (v0.2) 简介与核心特性
在Agent Lightning（简称AGL）的初始版本（v0.1）中，开发团队提供了一个灵活的训练框架，适用于任何带有强化学习的智能体。它具有以下几个核心特性：

  • 与现有智能体无缝集成，几乎实现零代码更改！
  • 可使用任何智能体框架（如LangChain、OpenAI Agent SDK、Microsoft Agent Framework等）进行构建；甚至可以不使用智能体框架（即纯Python程序）。
  • 对输入到LLM的内容没有限制，允许进行灵活的编排，例如总结、多智能体协作以及其他复杂的工作流。

AGL获取Token ID的演进过程
当Agent Lightning首次发布时，团队实现了一个经过插桩的vLLM服务器（instrumented vLLM server），它对vLLM的OpenAI服务器进行了猴子补丁（monkey-patched）处理以返回Token ID。现在，AGL会自动向每个请求添加return_token_ids参数，以便引擎在其响应中包含Token ID。然后，借助AGL中内置的追踪（tracing）功能，框架能够自动收集训练器端所需的数据，包括这些精确的Token ID。

用于智能体优化的中间件架构
从更精确的数据收集角度出发，在v0.2版本中，AGL在智能体优化中的角色变得更加清晰。从概念上讲，Agent Lightning（或AGL）引入了一个可持续的中间件层和标准化的数据协议，专门用于智能体优化，特别是智能体强化学习。
Agent Lightning的概念概览

AGL的模块化组件与标准化协议
Agent Lightning设计了一组模块化、可互操作的组件，这些组件共同实现了可扩展且高效的智能体强化学习。每个组件扮演着不同的角色，同时通过标准化的数据协议和定义良好的接口进行通信：

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

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

抽象化带来的算法灵活性优势
这种抽象视图带来了显著的优势。首先，它提供了极大的算法灵活性：数据收集可以依赖于各种追踪器（tracers）或发出自定义消息，这使得定义不同的奖励或捕获任何中间变量变得简单直接。在算法端，可以通过查询（query）spans来访问所需的数据，并且定制的适配器（adapters）允许自由的数据转换。这种设计还支持深度的算法定制，如信用分配（credit assignment）、使用部分数据进行辅助模型学习、通过数据调整改善训练等。此外，在该框架内，可以扩展到更多种类的算法，例如自动提示微调（APO），以及过滤高奖励数据并通过Unsloth拟合它们。

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

4. 实验环境与结果

实验环境
本文重点阐述了系统架构设计与工程实现，并未列出特定硬件配置（如GPU型号/数量）或基准测试数据集的定量信息。在软件配置方面，系统依赖于vLLM v0.10.2（引入了支持OpenAI兼容API返回Token ID的新特性）以及Agent Lightning v0.1/v0.2框架。兼容的智能体框架包括LangChain、OpenAI Agent SDK、Microsoft Agent Framework等，并支持与Megatron-LM、Unsloth等训练框架及POML等提示词编排工具结合使用。

实验结果
本文通过对比实验（见Fig 1）直接验证了重分词漂移对强化学习训练稳定性的破坏性影响。

  • 实验内容：对比在强化学习训练中“存储文本并重新分词”与“直接使用推理引擎返回的Token ID”两种策略下的学习曲线表现。

  • 实验结果与分析：

    1. 当采用存储文本并在训练时重新分词的传统方法（红线和蓝线）时，由于非唯一的分词现象、工具调用序列化差异以及聊天模板差异，导致了严重的“离策略（off-policy）”更新，学习曲线出现剧烈波动，模型无法稳定收敛。
    2. 当直接使用来自推理引擎的精确Token ID（黄线）时，彻底消除了推理与训练间的数据不一致性，学习曲线表现出高度的稳定性和持续的收敛趋势。

  • 结论：直接通过OpenAI兼容API返回精确的Token ID是保障智能体强化学习训练稳定性的关键。

5. 结论

本文深入剖析了智能体强化学习中由于API仅返回字符串而引发的“重分词漂移”问题，指出该问题会导致严重的离策略更新并破坏训练稳定性。通过在vLLM的OpenAI兼容API中引入return_token_ids参数，本文提供了一种直接返回精确Token ID的有效解决方案。同时，结合Agent Lightning（AGL v0.2）框架，本文提出了一种解耦的中间件架构，通过数据存储（Data Store）标准化了智能体运行器与模型训练器之间的通信。这一设计不仅从根本上消除了重分词漂移，还极大地提高了算法的灵活性，降低了复杂系统的耦合度，并支持异构计算资源的独立扩展，为构建鲁棒且可扩展的智能体强化学习生态系统奠定了坚实的基础。

6. 补充细节

扩展学习资料与应用案例
为了帮助开发者深入理解和应用上述特性与框架，官方提供了以下资源：

  • 完整的项目文档。
  • 鸟瞰图（Birds Eye View）架构解析。
  • 使用verl训练SQL智能体（包含多智能体编排）的教程。
  • 使用自动提示优化（Automatic Prompt Optimization, APO）训练房间选择器智能体的教程，其中提示词编排由POML提供支持。
  • 使用Unsloth训练数学智能体（基于带有MCP的OpenAI Agents SDK构建）的教程。

致谢
Agent Lightning团队对vLLM的维护者表达了诚挚的感谢，包括Kaichao You、Nick Hill、Aaron Pham、Cyrus Leung、Robert Shaw和Simon Mo。没有他们的支持与合作，这项整合工作是不可能完成的。Agent Lightning是来自微软研究院（Microsoft Research, MSR）的开源项目，团队对MSR支持此项开源探索表示感谢。Yuge Zhang是这项工作的主要贡献者。