Tunix：解决 TPU 闲置的高吞吐量 Agent 训练库

LLM 对齐的重点已经迅速从静态的聊天机器人对齐转向动态 Agent 工作流。当今的模型不仅仅是对话——它们执行多步骤推理、调用外部 API，并与复杂的环境进行交互。

训练推理 Agent 遇到特殊的挑战和性能瓶颈。Agent RL 训练的最新演进使这个过程从单轮对齐转向涉及复杂环境交互和工具使用的多轮决策制定。这种转变在基础设施层面提出了关于 rollout 性能和效率的新挑战；当 Agent 暂停以执行代码、查询数据库或等待网络搜索时，昂贵的 AI 加速器利用率会大幅下降，因为 TPU 在等待环境步骤时处于闲置状态。

Tunix 是 Google 的后训练库，在其最新版本中原生解决了这一瓶颈，为大规模训练 LLM Agent 引入了一个高效、可组合的框架。Tunix 从两个方面保持加速器的充分利用：

异步 Rollout：一个高并发的 rollout 引擎完全将 TPU 执行与主机端环境延迟（如网络 I/O 或工具执行）解耦。

无屏障流水线：一个动态的生产者-消费者架构不断批处理和流式传输可变长度的轨迹到训练器，防止流水线停顿。

除了编排之外，Agent RL 还需要专门的可观测性。虽然 XProf 这样的标准分析器提供深层的算子级跟踪，但它们的高开销限制了它们只能进行短时间、零星的捕获。Tunix 引入了围绕域特定 RL 指标构建的持续、轻量级的植入。通过将这些高级循环指标与 TPU 时间线相关联，开发者可以获得执行效率的全局视图，快速发现和解决系统瓶颈。

总之，Tunix 旨在最大化 TPU 吞吐量、保持环境模块化，并使多轮训练效率完全透明。以下是它的工作原理。

1. 异步与解耦的 Rollout：接近零的空闲时间，最大吞吐量

实现峰值硬件吞吐量意味着让 TPU 持续忙碌。Tunix 通过结合异步 rollout 来消除执行气泡和掉队者，以及一个持续向训练器流式传输数据的解耦流水线来实现这一点。

异步 Rollout

在 Agent RL 中，轨迹生成（rollout）是耗时最长的阶段。然而，传统的同步 rollout 架构会产生两个主要问题，如下图所示。

执行气泡：当 rollout 同步等待环境初始化或返回状态和奖励时，它会在加速器中产生执行气泡并降低效率。

掉队者效应：批量生成也容易受到长尾问题的影响，整体延迟由组中最慢的轨迹决定。

Tunix 通过异步轨迹收集引擎解决了这个问题。

高并发执行：Tunix 在 RolloutOrchestrator 中利用 Python 的 asyncio，管理大量并发的 Agent-环境交互池。当一个 Agent 暂停进行主机端工具执行时，推理引擎立即转向为其他活跃轨迹生成令牌。

异步 vLLM 和 SGLang 集成：Tunix 原生集成了性能强劲的推理引擎，如 vLLM-TPU 和 SGLang-Jax。通过启用异步请求处理，引擎确保了 TPU 上的非阻塞采样和最大并发。

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

解耦的 Rollout 和训练流水线

虽然异步 rollout 解决了轨迹生成瓶颈，但端到端 RL 工作流中硬件效率的另一个关键挑战是将动态、可变长度、可能很长的 rollout 与严格同步的训练循环桥接。幼稚的方法依赖于一个同步点，强制加速器等待整个轨迹批次完成才能启动训练步骤，导致训练 TPU 饥饿。

Tunix 通过将 rollout 和训练解耦成一个持续的生产者-消费者流水线来消除这个瓶颈（如下图所示）：

生产者：异步 rollout 编排器持续将完成的轨迹输出到一个高吞吐量队列。

消费者：AgenticRLLearner 从这个队列消费。对于需要多个推理路径来计算组优势的 GRPO 这样的算法，Tunix 动态地即时对这些异步轨迹进行分组。

一旦轨迹组完成，它会被后处理、评分，并直接流式传输到训练器。这个流水线确保同步训练器持续被喂养，最大化端到端吞吐量。

2. 可组合的 Agent 和环境抽象——即插即用的开源环境

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

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

Agent 层：管理 prompt 格式化、操作生成和对话历史。它自动应用策略模型的 chat 解析器，并在多轮边界处保留特殊令牌——这对确保严格的 Token-In、Token-Out (TITO) 行为至关重要。你可以通过子类化 ConversationAgentBase 来轻松自定义生成逻辑。

环境层：开箱即用，Tunix 提供预构建的 TaskEnvironment 和 ToolEnvironment 类。你也可以继承 BaseTaskEnv 来与任何外部系统对接。Tunix 自动处理多轮情节生命周期、观测路由和奖励处理。

为什么这很重要：你可以在几分钟内接入任何开源 RL 环境。因为 Agent 和环境逻辑与训练工作流完全解耦，用交互式 bash 终端替换单轮数学验证器只需要对你的训练代码做零修改。为了展示这种可组合设计的强大，我们接下来展示几个例子，说明如何轻松使用新的 Agent、模型或环境。你可以在我们的 recipes 中找到更多自定义 Agent/Env 的详细例子。

示例 1：预构建 vs. 自定义 Agent

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

from tunix.rl.agentic.agentic_grpo_learner import GRPOLearner
from tunix.rl.agentic.agents.model_agent import ModelAgent, ToolAgent

# Non tool calling single turn agent
learner = GRPOLearner(
    agent_class=ModelAgent, 
    agent_kwargs={"system_prompt": "my system prompt"},
    ...
)

# Customized tool call agent
tool_map = {"calculator": CustomizedCalculatorClass, ...}
learner = GRPOLearner(
    agent_class=ToolAgent, 
    agent_kwargs={
        "system_prompt": "my system prompt",
        "tool_parser_name": "gemma",
        "tool_map": tool_map,
    },
    ...
)

或者，你可以构建自己的自定义 Agent 并添加特定逻辑来处理模型响应。Tunix 会自动将这个 Agent 接入端到端训练工作流。例如 SWEAgent、FrozenLakeAgent

from tunix.rl.agentic.agents.base_agent import ConversationAgentBase
from tunix.rl.agentic.agents import agent_types

# Bring your own agent! 
# Notice how the agent doesn't need to know anything about the model (if it is Qwen, Llama, or Gemma)
class MyAgent(ConversationAgentBase):
    def __init__(self, args):
        ...

    def update_from_model(self, response: str, **kwargs) -> agent_types.Action:
        # Custom logic to process the raw response (e.g., extracting <answer> tags)
        ...

# Tunix automatically wires up the e2e workflow
learner = GRPOLearner(agent_class=MyAgent, agent_kwargs={...}, ...)

示例 2：接入自定义环境

与 Agent 类似，Tunix 提供了多个预构建的环境，包括 TaskEnvironment、ToolEnvironment。或者，你也可以通过实现几个主要 API 来接入自己的自定义环境，包括任何开源环境，比如下面的 Gymnasium 示例。

import gymnasium as gym
from tunix.rl.agentic.agentic_grpo_learner import GRPOLearner
from tunix.rl.agentic.environments.base_environment import BaseTaskEnv, EnvStepResult

# You only need to focus on the core logic of environment interactions, and Tunix will automatically handle the rest of the lifecycle management and function invocation.
class MyEnv(BaseTaskEnv):
    def _initial_observation(self):
        # handle env creation and initial observation
        self.env = gym.make("your_chosen_env")
        observation, info = self.env.reset(seed=42)
        return observation

    def _step_impl(self, action):
        # compute observation, reward, done, info
        action = self.env.action_space.sample() 
        obs, reward, done, info = self.env.step(action)
        return EnvStepResult(obs, reward, done, info)

    def close(self):
        self.env.close() # clean up env after trajectory is done

learner = GRPOLearner(env_class=MyEnv, ...)

3. 消除黑盒：RL 特定的轻量级分析

在大规模运行异步 Agent 训练时，传统的日志记录显得力不从心。你需要粒度化但又域特定的可见性来识别效率问题：瓶颈是在生成阶段吗？工具调用花费了太多时间吗？还是数据加载器太慢了？

XProf 这样的标准分析器提供详细的算子级跟踪来理解微观级性能，如内核和模型执行。然而，使用这些工具捕获长跨度的跟踪通常是高成本的，并且在低级数据的噪声中识别宏观级瓶颈仍然很困难。对于 Agent RL 的复杂工作流，开发者需要一个轻量级的、构建在直接映射到 RL 阶段的域特定指标之上的宏观级视图。

Tunix 通过仔细跟踪代表全局流水线（rollout、训练和权重同步阶段如何交互）和重要子步骤（每个模型调用、环境交互等）的一组最小关键 RL 特定指标来提供这个大局观。因为它们是轻量级的，这些指标在整个训练工作期间持续运行。用户可以快速识别工作流在全局范围内在哪里停顿，然后部署像 XProf 这样的工具进行有针对性的进一步调试。

上面的图展示了从多轮 Agent 训练工作中捕获的 Perfetto 跟踪，详细说明了跨 CPU 线程和 TPU 设备的分阶段执行时间线。如跟踪所示，TPU 设备的利用率远高于 CPU 线程的利用率，CPU 线程的空闲时间主要是由于环境执行延迟。

这个分阶段 RL 流水线的宏观级跟踪使你能够：

精准定位 TPU 饥饿：直观显示 Python 工具调用或环境执行阻止异步流水线的确切时刻。反过来说，它也让你确认并行 rollout 成功地重叠以保持加速器饱和。

验证流水线对齐：追踪宏观阶段的精确时间以确保它们对齐而不引入隐藏的延迟气泡。你可以轻松验证训练器没有在等待 rollout 生成，或者权重同步没有导致严重的执行延迟。

优化训练配置：使用指标数据动态调整性能。例如，你可以通过将线程池与 TPU 空闲时间相关联来调整最大 rollout 并发，或者基于数据生成速率和 HBM 约束优化训练微批次大小。

Tunix 将分布式多轮 RL 的"黑盒"转变成了训练工作的透明、可优化的时间线。

Tunix 与生态系统的比较

如果你在评估 Agent RL 框架，以下是 Tunix 如何脱颖而出的方式：

vs. OpenRLHF / veRL：OpenRLHF 和 veRL 在使用 Ray + vLLM 上取得了显著进展。然而，它们主要是为 PyTorch 生态系统构建的。Tunix 将这一能力原生带到 JAX/TPU 生态系统。完全构建在 JAX、Flax 和 Optax 之上，Tunix 提供原生 Pathways 多主机分布式训练，利用 XLA 的编译器优化。

vs. Hugging Face TRL：TRL 很适合标准单轮 SFT（监督微调）和 RLHF（人类反馈强化学习）。然而，编排复杂的、多轮异步循环通常需要大量的自定义粘合代码。Tunix 使多轮、工具使用环境成为开箱即用的一等公民。

vs. Ray RLlib：RLlib 是一个全面、通用的 RL 强大工具。然而，原生地将现代 LLM 映射到在加速器上共享权重而没有重大开销是复杂的。Tunix 反其道而行：它是一个 LLM 优先的库，将高性能 RL 直接带到原生 LLM 服务基础设施。

开始构建你的 Agent

无论你是在复现 SOTA 推理模型、微调 Gemma 或 Qwen 系列来"思考"，或者部署复杂的多 Agent 系统，Tunix 都提供了下一代推理 Agent 所需的高性能基础。立即开始构建！

🌟 Star 仓库 & 探索代码：github.com/google/tunix

📖 Recipes：SWE 编码 Agent、数学、游戏 Agent。

📖 阅读文档：深入了解我们的 Agent RL 架构在 tunix.readthedocs.io

🚀 尝试快速开始：跳转到我们的 /examples 文件夹，探索各种 recipes 并立即运行你的第一个训练工作！

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