GPT-OSS Agentic RL 训练：一次实践回顾

TL;DR

Hugging Face 和 LinkedIn 研究人员通过解决 PPO on-policy 完整性、注意力 sink 反向传播和 MoE 内存实例化中的关键不稳定性，成功地为 GPT-OSS 模型解锁了 agentic 强化学习（RL）。这些修复使 GPT-OSS 能够作为多步决策代理的稳定骨干，这些代理与环境和工具进行交互。

Agentic RL 和 GPT-OSS

Agentic RL 通过优化整个决策过程，区别于传统的单轮 RL。模型不是生成静态响应，而是学习在多步轨迹中规划行动、调用工具并调整行为。这需要一个闭环系统，其中代理收集 rollout 轨迹，计算奖励，并使用如 PPO 或 GRPO 之类的算法迭代更新其策略。

虽然 GPT-OSS 在性能上与 OpenAI o3-mini 和 o4-mini 相当，但其用于 agentic RL 的适用性之前未得到验证。研究人员使用了 verl 训练框架，并在 GSM8K、Retool（一个 agentic 编码任务）和可验证指令遵循任务上测试了 GPT-OSS-20B 模型。

解决 PPO On-Policy 不稳定性

初始训练运行表现出 KL 散度和熵的爆炸以及非递增奖励。团队发现这是由混合专家（MoE）架构导致的 Proximal Policy Optimization（PPO）on-policy 完整性失败。

MoE 对数概率不匹配

在纯 on-policy PPO 中，重要性采样比必须恰好为 1。然而，在像 GPT-OSS 这样的 MoE 架构中，由于浮点差异或随机性，门控网络可能在用于 rollout 生成的前向传递和用于训练的前向传递之间将输入路由到不同的专家。这导致当前对数概率和旧对数概率之间出现不匹配，错误地触发了 PPO 裁剪并违反了 on-policy 假设。

修复： 团队实施了一种对数概率替换，在环境已知为 on-policy（小批量大小等于全局批量大小时）时覆盖计算，通过设置 old_log_prob = log_prob.detach() 将重要性比强制恢复为 1。

通过注意力 Sink 修正训练-推理不匹配

即使在修复了 PPO 完整性之后，梯度范数仍然继续爆炸。研究人员识别出一个基本的训练-推理不匹配，其中推理引擎（SGLang）和训练栈（FSDP 带 FlashAttention-v2）产生了不同的 token-level 概率。

注意力 Sink 的作用

GPT-OSS 使用注意力 sink——作为 softmax 计算中的“虚拟 token”的可学习标量参数。这些 sink 使模型能够将注意力质量分配给一个学习参数，而不是强制将其分配给内容 token，从而在流式推理中提高稳定性。

在 FlashAttention v3 中的实现

团队发现 verl 硬编码了 FlashAttention v2，而 v2 并不支持注意力 sink，且 v2 和 v3 都不支持 sink 梯度所需的反向传播。为了解决这个问题，他们：

  1. 利用了来自 vLLM FlashAttention 分叉的前向传递。
  2. 实现了反向传递以计算 sink 梯度 $ rac{\partial L}{\partial S_{h}}$。

结果： 此修复在 GSM8K、VerifyIf 和 VerifyIf 任务上导致了显著更快的收敛和稳定的奖励提升。

扩展长上下文的内存效率

Agentic RL 需要随着环境反馈被追加到轨迹中而扩展上下文窗口。团队实施了两项主要的内存优化以防止 Out-of-Memory（OOM）故障。

减轻 MoE 专家实例化

研究人员发现，Hugging Face Transformers 推理前向路径会为所有专家复制隐藏状态，在 GPU 内存中实例化极大的张量（例如，尝试为 20B 模型分配 180 GiB）。他们修补了实现，以使用更内存高效的执行路径，顺序处理专家。

使用 FlashAttention v3 的序列并行

为了进一步减少每个 GPU 的激活内存，团队实施了序列并行（上下文并行）。这将输入序列划分到多个设备上，降低了峰值激活占用。

由于注意力层要求序列的所有 token 必须存在于同一 GPU 上，团队实施了一种 all-to-all 通信策略：

  • Pre-attention： 收集序列元素，在注意力头级别进行分割。
  • Post-attention： 将输出重新分配回原始序列并行布局。

此设计是 attention-sink-aware 并且与 FlashAttention v3 兼容，使模型能够处理多步代理所需的长上下文窗口。

Sources