告别重分词偏移:在智能体强化学习中通过 OpenAI 兼容 API 返回 Token ID 的重要性
摘要。 智能体(Agent)通常通过 OpenAI 兼容的端点调用 LLM,而这些端点此前仅返回基于字符串的输入和输出。在智能体强化学习(Agent RL)中,由于我们称之为重分词偏移(Retokenization Drift)的现象,这可能导致训练和推理之间的不一致。这种现象的发生是因为 Token 在推理期间被去分词(detokenized),随后在训练期间又被重新分词(retokenized);即使对应的字符串完全相同,这两组 Token 也可能存在差异。现在,您可以要求 vLLM 的 OpenAI 兼容端点返回 Prompt 和生成响应的精确 Token ID。只需将 "return_token_ids": true 传递给 /v1/chat/completions 或 /v1/completions,您就会在常规文本输出旁收到 prompt_token_ids 和 token_ids。这使得智能体 RL 更加鲁棒,因为不再会发生偏移。这与 Agent Lightning 完美契合,在 Agent Lightning 中,每次模型调用都被视为独立的更新样本而无需拼接;只需在启用 return_token_ids 的情况下记录返回的 ID 即可。
相关链接
- 文档:OpenAI 兼容服务器
- 可尝试此功能的项目:Agent Lightning (GitHub, 文档)
为什么 Token ID 对智能体强化学习很重要
LLM 的强化学习(RL)是在 Token 序列上进行训练的,因此训练器需要行为策略采样的精确 Token ID。在单轮设置中,这通常很直接,因为调用 vLLM 的底层 generate 会直接返回 Token。
在智能体场景中,大多数智能体框架调用 OpenAI 风格的 chat.completions / completions。智能体更青睐这些 API 而非原始的 generate,因为它们提供了智能体技术栈赖以构建的高级功能,如聊天模板与角色(系统/用户/助手)、工具/函数调用、结构化输出等。这些 API 在历史上仅返回字符串,这可能会在智能体 RL 中引起问题。以前,存储的文本必须在训练期间重新分词,但在实践中,由于重分词偏移,这种做法不稳定且准确性较低。
您在 RL 中会看到的症状:学习曲线不稳定(如下图所示),以及难以调试的数据偏差——即您认为正在优化的数据与模型实际采样的数据之间存在差异。
红线和蓝线是在相同设置下获得的(即在训练中存储文本并重新分词),而黄线则是直接使用推理引擎生成的 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)RL 更新。同策略(On-policy)对于稳定的 RL 训练至关重要,微妙的变化也会产生巨大影响。由重分词偏移引起的离线策略效应甚至不在 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) 中,我们为任何智能体提供了灵活的 RL 训练框架。它具有几个核心特性:
- 与现有智能体无缝集成,几乎“零代码更改”!
- 可配合任何智能体框架构建(LangChain, OpenAI Agent SDK, Microsoft Agent Framework, ...);甚至无需智能体框架(直接使用 Python 程序)。
- 对 LLM 的输入没有限制,允许灵活的编排,如摘要、多智能体协作和其他复杂工作流。
在 Agent Lightning 最初发布时,我们实现了一个插桩式 vLLM 服务器,通过对 vLLM 的 OpenAI 服务器进行猴子补丁(monkey-patch)来返回 Token ID。现在,AGL 会自动向每个请求添加 return_token_ids,以便引擎在响应中包含 Token ID。随后,利用 AGL 内置的追踪(tracing)能力,我们会自动收集训练端所需的数据,包括这些 Token ID。
智能体优化的中间件
从更精确的数据收集角度出发,在 v0.2 中,我们进一步明确了 AGL 在智能体优化中的角色。从概念上讲,Agent Lightning (AGL) 为智能体优化(尤其是智能体强化学习)引入了一个可持续的中间件层和标准化的数据协议。
Agent Lightning 的概念概览。
Agent Lightning 设计了一套模块化、可互操作的组件,共同实现可扩展且高效的智能体强化学习。每个组件扮演不同角色的同时,通过标准化的数据协议和定义良好的接口进行通信。
- Agent Runner —— 负责执行智能体以完成指派的任务。它接收任务,委托给智能体执行,收集结果和中间数据,并将这些信息报告给数据存储中心。Agent Runner 与 LLM 端独立运行,因此可以使用不同的资源(如 CPU)进行托管,并可以水平扩展以支持大量并发的智能体实例。
- Algorithm (Model Trainer) —— 托管用于推理和训练的大语言模型(LLM)。该组件编排整个 RL 循环,包括任务采样、展开(rollout)管理以及基于收集的经验数据进行模型更新。它通常运行在 GPU 资源上,并通过共享数据协议与 Agent Runner 异步交互。
- Data Store —— 作为管理智能体 RL 生态系统中所有数据交换和存储的中央仓库。它提供标准化的接口和统一的数据架构,以确保异构组件之间的互操作性。通过这种设计,Algorithm 和 Agent Runner 可以间接但有效地进行通信,实现灵活且可扩展的协作。例如,使用标准化的
rollouts,Algorithm 可以异步地将任务委托给 Agent Runner,后者执行任务并通过spans数据结构报告执行轨迹。
Agent Lightning 中的训练循环。
在这种以数据存储为中心的设计理念下,所有智能体训练迭代都被抽象为两个步骤。第一步是收集智能体运行数据(AGL 中的 spans)并将其存储在数据存储中心;第二步是从存储中心检索所需数据并将其发送到算法端进行训练。
这种抽象视图带来了多项优势。首先,它提供了更大的算法灵活性:数据收集可以依赖于各种追踪器或发送自定义消息,使得定义不同的奖励或捕获任何中间变量变得非常简单。在算法端,可以通过 查询(query) spans 来访问所需数据,而自定义的 适配器(adapters) 则允许自由地进行数据转换。
这种设计还支持算法自定义,如信用分配、使用部分数据进行辅助模型学习、通过数据调整改进训练等。此外,在此框架内,我们可以扩展到更多种类的算法,如自动 Prompt 调优 (APO),以及通过 Unsloth 筛选高奖励数据并进行拟合。
该设计的第二个主要优势在于其能够通过模块化分离降低系统的整体复杂性,同时允许不同组件利用各自的资源和优化策略。智能体 RL 训练系统本质上是复杂的,因为它利用动态的、环境驱动的交互来使模型从经验数据中持续学习。一个典型的智能体 RL 技术栈由几个关键组件组成,包括智能体框架(如 LangChain、MCP)、LLM 推理引擎(如 vLLM)和训练框架(如 Megatron-LM)。如果没有解耦的架构,这些组件的独立性和异构性会导致巨大的系统复杂性。相比之下,解耦设计允许系统适应不同的资源需求:例如,智能体端可能需要更高的 CPU 容量,而 LLM 推理和训练通常是 GPU 密集型的。这种模块化结构还促进了每个组件的独立水平扩展,提高了效率和可维护性。
更多资料
- 完整文档
- 鸟瞰图 (Birds Eye View)
- 使用 verl 训练 SQL 智能体(具备多智能体编排)
- 使用自动 Prompt 优化训练房间选择智能体,其中 Prompt 编理由 POML 驱动。
- 使用 Unsloth 训练数学智能体(配合 OpenAI Agents SDK 和 MCP 构建)
祝训练愉快!⚡
致谢
我们要向 vLLM 的维护者们表示衷心的感谢,包括 尤凯超、Nick Hill、Aaron Pham、Cyrus Leung、Robert Shaw 和 Simon Mo。如果没有他们的支持与合作,这种集成是不可能实现的。
Agent Lightning 是微软研究院(Microsoft Research)的一个开源项目。我们衷心感谢 MSR 对此开源探索的支持。张宇歌是这项工作的主要贡献者。