Google 推出 Tunix：基于 JAX 的高吞吐智能体后训练库

大型语言模型（LLM）对齐的重点已迅速从静态聊天机器人对齐转向动态自治工作流。如今的模型不仅会对话——它们还能执行多步推理、调用外部 API 并与复杂环境互动。

训练推理型智能体面临特殊的挑战和瓶颈。智能体强化学习（agentic RL）训练的最新发展将过程从单轮对齐转向多轮决策，其中涉及复杂的环境互动和工具使用。这一转变在基础设施方面提出了新的挑战，影响展开性能和效率；当智能体暂停执行代码、查询数据库或等待网络搜索时，昂贵的 AI 加速器利用率会急剧下降，因为 TPU 在等待环境步骤时处于空闲状态。

Tunix——谷歌的后训练库——在其最新发布版本中原生地解决了这一瓶颈，引入了一个高效、可组合的框架，用于大规模训练 LLM 智能体。Tunix 在两个方面保持加速器的充分利用：

除了编排之外，智能体强化学习还需要专门的可观测性。虽然标准分析器如 XProf 提供深度操作级别的轨迹，但其高开销限制了只能进行短时、零散的捕获。Tunix 引入了围绕特定领域 RL 指标直接构建的连续轻量级检测工具。通过将这些高层循环指标与 TPU 时间线关联，开发者可以获得执行效率的全局视图，从而快速发现并解决系统瓶颈。

最终，Tunix 的设计目标是最大化 TPU 吞吐量，保持环境模块化，并使多轮训练效率完全透明。其内部工作机制如下。

实现硬件峰值吞吐量意味着需要让 TPU 不断保持繁忙状态。Tunix 通过结合异步 rollout 以消除执行空洞和滞后者，以及一个连续向训练器流式传输数据的解耦流水线来实现这一目标。

在智能体强化学习中，轨迹生成（rollout）是最耗时的阶段。然而，传统的同步 rollout 架构会产生两个主要问题，如下图所示。

Tunix 通过异步轨迹收集引擎（Asynchronous Trajectory Collector Engine）解决了这个问题。

这种架构将模型推理、工具执行和奖励计算完全重叠，从而保持高硬件利用率。

尽管异步 rollout 解决了轨迹生成的瓶颈，但端到端 RL 工作流中硬件效率的另一个关键挑战是将动态、长度可变并且可能是长尾的 rollout 与严格同步的训练循环相连接。一个简单的做法是依赖一个同步点，这会强制加速器等待整个批次的轨迹完成后才能启动训练步骤，从而使训练 TPU 处于空闲状态。Tunix 通过将 rollout 和训练解耦成连续的生产者-消费者管道（如下图所示）消除了这一瓶颈：

轨迹组一旦完成，就会被后处理、评分，并直接流入训练器。这个管道确保同步训练器不断得到数据输入，从而最大化端到端吞吐量。

RL 框架中的一个主要摩擦点是算法与环境循环的严格耦合。修改代码库以支持新的开源软件（OSS）基准，例如 SWE-bench、WebArena 或自定义游戏引擎，通常需要大规模重写。

Tunix 通过解耦、可组合的架构解决了这个问题。通过公开干净的 API 边界，Tunix 自动化步骤调用和生命周期管理，使你可以完全专注于核心交互逻辑。

重要性：你可以在几分钟内接入任何开源 RL 环境。由于代理和环境逻辑与训练工作流完全解耦，用一个单步数学验证器替换为交互式 bash 终端无需修改任何训练代码。为了展示这种可组合设计的强大功能，我们接下来展示几个新代理、模型或环境能够多么轻松地使用的示例。你可以在我们的示例配方中找到更多定制化 Agent/Env 的详细示例：https://github.com/google/tunix/tree/main/examples。

Tunix 提供了内置类，如 ModelAgent 和 ToolAgent，可通过配置立即工作。

或者，您可以构建自己的自定义Agent，并添加具体逻辑来处理模型响应。Tunix会在端到端训练工作流中自动连接此Agent。例如：SWEAgent，FrozenLakeAgent

类似于Agents，Tunix提供了许多预构建的环境，包括TaskEnvironment，ToolEnvironment。或者，您也可以通过简单地实现几个主要API来引入自己的自定义环境，包括任何开源环境，例如下面的Gymnasium示例。

在大规模运行异步agentic训练时，传统日志记录显得不足。您需要细粒度但特定领域的可见性来识别效率问题：瓶颈出现在生成阶段吗？工具调用耗时过多吗？还是数据加载器过慢？

像XProf这样的标准分析器：https://openxla.org/xprof 提供详细的操作级追踪，以了解如内核和模型执行等微观性能。然而，使用这些工具捕获长跨度追踪通常成本高昂，并且在低级数据噪声中识别宏观瓶颈仍然困难。对于agentic RL的复杂工作流，开发者需要一种基于领域特定指标的轻量级宏观视图，这些指标能够直接映射到RL阶段。

Tunix通过仔细跟踪一组关键的RL特定指标提供了这一整体视图，这些指标既代表全球管道（如rollout、训练和权重同步阶段的交互），也代表重要子步骤（每次模型调用、环境交互等）。由于这些指标轻量，它们可以在整个训练任务中持续运行。用户可以快速识别工作流的全局停滞位置，然后部署像XProf这样的工具进行有针对性的进一步调试。

上图显示了从一个多轮agentic训练任务中捕获的Perfetto追踪，详细展示了CPU线程和TPU设备的分阶段执行时间线。如追踪所示，TPU设备的利用率远高于CPU线程，其空闲时间主要由于环境执行延迟。

这种分阶段RL管道的宏观追踪使您能够：

Tunix 将分布式多轮强化学习（RL）的“黑箱”变成了训练任务可透明、可优化的时间线。

如果你正在评估用于智能体强化学习（Agentic RL）的框架，Tunix 的优势如下：

无论你是在复现最先进的推理模型、微调 Gemma 或 Qwen 系列模型以“思考”，还是部署复杂的多智能体系统，Tunix 都提供了下一代推理智能体所需的高性能基础。从今天开始构建吧！

Tunix 正在由 Google 和更广泛的社区积极进行开源开发。如果你正在构建下一代推理智能体，欢迎访问我们的 GitHub Issues，让我们知道你正在接入哪些环境。