解锁 GPT-OSS 的智能体强化学习训练：一次实践回顾

摘要

1) 一句话总结

本文记录了在 verl 框架下为 GPT-OSS 模型解锁智能体强化学习（Agentic RL）训练的工程实践，通过修复 MoE 策略一致性、实现 Attention Sink 反向传播及优化内存分配，成功实现了稳定且高效的多步决策模型训练。

2) 关键点

实验背景与配置：使用 verl 开源框架，重点针对 GPT-OSS-20B 模型（兼容 120B），在 gsm8k、Retool（智能体编码任务）等任务上进行验证，并以 Qwen-2.5-32B 作为基准。
模板适配：首要步骤是更新训练框架以完全支持 GPT-OSS 引入的 Harmony 聊天模板，确保推演生成和轨迹构建的正确性。
修复 PPO 同策略（On-Policy）完整性：针对 MoE 架构双前向传播中非确定性路由导致的对数概率不匹配问题，通过在同策略时强制 old_log_prob = log_prob.detach()，将重要性采样比率恢复为 1。
解决训练与推理不匹配：发现 verl 硬编码的 FlashAttention v2 不支持 Attention Sink，导致学习主要由 MoE 层驱动而注意力机制贡献不足。
实现 FlashAttention V3 的 Attention Sink 反向传播：基于 vLLM 的前向传播分支，实现了 FlashAttention v3 中 Attention Sink 的反向传播（计算 Sink 参数梯度），大幅提升了模型在各类 RL 任务中的收敛速度和稳定性。
MoE 内存优化：修复了 Hugging Face Transformers 中 MoE 前向路径在 FSDP 下重复实例化专家的问题，改用 for 循环顺序处理专家，大幅降低了显存占用。
引入序列并行（Sequence Parallelism）：实现了兼容 Attention-Sink 和 FlashAttention v3 的序列并行，在注意力头级别进行拆分并使用 all-to-all 通信，有效支持了多步智能体所需的长上下文窗口训练。

3) 风险/缺口（原文明确指出的训练隐患与问题）

策略偏离风险：在 MoE 架构下，若不修正对数概率的计算差异，PPO 训练会触发错误的裁剪，将原本的同策略（On-Policy）优化变成异策略（Off-Policy）优化，导致 KL 散度和熵爆炸。
梯度爆炸与收敛失败：若底层注意力机制（如 FlashAttention）不支持 Attention Sink 的反向传播，会引发严重的训练-推理不匹配，导致梯度范数爆炸且奖励无法提升。
内存溢出（OOM）风险：在未优化的 MoE 前向路径下进行 FSDP 训练，会因重复复制隐藏状态和执行批量矩阵乘法分配过多内存，导致在 16 个 H200 节点上训练 20B 模型时反复出现 OOM 故障。

正文

智能体强化学习（Agentic RL）扩展了传统的 LLM（大语言模型）训练，它不仅优化单轮响应，还通过在训练期间与环境的直接交互来学习整个决策过程。与依赖静态数据集的传统单轮强化学习或离线偏好方法不同，智能体 RL 通过主动收集策略内（on-policy）数据来训练策略——智能体在模拟或真实环境中规划动作、调用工具、观察结果，并在多步轨迹中调整其行为。这种交互驱动的优化能够在长视野决策中分配信用，其中查询重构、工具选择和执行顺序等中间选择会直接影响最终的成功率。训练遵循一个迭代的闭环：智能体与环境交互以收集推演（rollout）轨迹，计算这些轨迹的奖励，根据观察到的结果更新策略，然后使用更新后的策略驱动下一轮的交互和数据收集（例如 GRPO 或 PPO 算法）。

LinkedIn 是一家 AI 优先的公司，致力于构建智能体来帮助职场人士取得成功。在这种背景下，模型必须在不完整的信息下进行推理，与结构化服务交互，并在多个步骤中适应不断变化的用户意图，而不是仅仅生成单一的静态响应。这些能力对于支持招聘人员、求职者、知识探索者和学习者等最终用户目标的智能体尤为关键（例如检索信息、优化查询、协调工具和执行多步工作流）。通过交互学习稳健的决策策略，智能体 RL 为通过端到端优化构建可扩展、可靠且适应性强的 AI 系统提供了原则性基础。

GPT-OSS 模型已经展现出与 OpenAI o3-mini 和 o4-mini 相当的性能，但其在智能体强化学习训练中的适用性尚未得到验证。最近的大多数工作都集中在没有工具调用的微调上。本文将探讨解锁 GPT-OSS 作为智能体应用潜在骨干模型的智能体 RL 训练之旅。

在我们的实验中，我们使用 verl 作为训练框架，因为它是开源社区中最受欢迎的框架之一。我们使用了 gsm8k、Retool 任务以及可验证指令遵循任务，这些都是 RL 训练中常用的任务。我们重点展示 GPT-OSS-20B 模型的实验结果，同时我们的 Attention-Sink 修复方案也适用于 GPT-OSS-120B。此外，我们还使用 Qwen-2.5-32B 模型作为基准，以对比 RL 训练期间的标准指标趋势。

GPT-OSS RL 训练的挑战

verl 一直是团队使用的开源框架，团队此前也曾参与协作和贡献，以帮助普及智能体强化学习训练。随着 GPT-OSS 中引入新的 Harmony 聊天模板，第一步是确保训练框架完全支持 Harmony 所需的更新消息格式和对话语义。这一步有助于在推演生成、轨迹构建和工具解析在全新模板下保持一致和正确。

团队使用 ReTool 作为代表性示例来验证代码的正确性。ReTool 是一个智能体编码任务，要求模型在代码编译器工具的辅助下解决数学问题。这种设置允许模型专注于核心推理和算法逻辑，同时将实际的算术和执行委托给工具。在一个回合（episode）中，模型多次与代码工具交互，利用执行结果作为反馈来完善其解决方案。在轨迹的最后，模型给出一个最终答案，并据此计算奖励。

在最初的训练运行中，我们观察到 KL 散度和熵爆炸，同时奖励没有增加，这表明 GPT-OSS 训练设置中存在潜在问题。与 Qwen32b 相比，GPT-OSS 20B 的奖励明显偏低，且随着训练的进行，梯度范数出现了爆炸。

verl 中的实用调试之旅：恢复 PPO 的 On-Policy 完整性

修复 MoE 对数概率不匹配

我们专注于 On-Policy（同策略）方法，因为它们提供了更高的稳定性和更可靠的收敛性。纯 On-Policy 近端策略优化（PPO）的基础要求重要性采样比率（importance sampling ratio）必须精确为 1。重要性比率的数学定义为：

$ratio = \frac{π ( a ∣ s )}{π _{old} ( a ∣ s )}$

这一要求确保策略更新仅在当前策略生成的数据上执行，从而防止意外的裁剪（clipping）。

在我们的 ReTool 训练中，即使是 On-Policy 训练，我们也观察到了非零的裁剪值。这源于两个对数概率之间的不匹配：

当前对数概率 log_prob： $lo g (π (a ∣ s))$
旧对数概率 old_log_prob： $lo g (π_{old} (a ∣ s))$

根本原因：双前向传播与 MoE 架构

在 verl 0.3.0 之前，实现依赖于两次独立的前向传播（一次计算当前的 log_prob，一次检索存储的 old_log_prob）来处理相同的状态-动作对。

在像 GPT-OSS 这样的混合专家（MoE）架构中，门控网络将输入路由到不同的专家。由于实现因素（例如微妙的浮点差异或显式的随机性），两次传播之间的专家路由可能会略有不同。

这种路由差异导致：

$lo g (π (a ∣ s)) \neq = lo g (π_{old} (a ∣ s))$

由此产生的比率偏离了 1，错误地触发了 PPO 裁剪，并违反了核心的 On-Policy 假设。

解决方案：通过对数概率替换强制 Ratio = 1

该修复方案通过在已知环境为 On-Policy 时（即当 mini-batch 大小等于全局 batch 大小时），在逻辑上覆盖有缺陷的计算来解决此问题：

if on_policy:
    old_log_prob = log_prob.detach()
else:
    old_log_prob = model_inputs["old_log_probs"]

通过将 old_log_prob 设置为等于新计算的 log_prob（使用 detach 防止梯度流过参考值），重要性比率在数学上被强制恢复为 1。这种策略绕过了由 MoE 非确定性路由引起的不稳定性，并保证了 PPO 训练期间严格的 On-Policy 行为。

纠正训练与推理的不匹配

尽管修复对数概率不匹配将重要性采样裁剪率降至零，但梯度范数继续爆炸，奖励也未能提高。为了隔离问题，我们将训练简化为 GSM8K（一个没有智能体工具使用的单步任务）。同样的不稳定性依然存在，这表明 在 verl 下使用 GPT-OSS 进行基础 RL 训练存在根本性问题。

我们假设 训练-推理不匹配 可能是潜在原因：推理时执行（vLLM 和 SGLang 等引擎积极优化吞吐量）与 FSDP 下的训练时执行（优先考虑数值精度和稳定性）之间的差异，实际上可能将原本的 On-Policy RL 变成了 Off-Policy 优化。

在应用推演校正（rollout correction，即序列级重要性采样）后，训练动态显著改善，梯度范数保持稳定而不再爆炸。然而，奖励仅有适度增加，并且在简单的 GSM8K 任务上的收敛速度与较小的密集模型变体相比仍然慢得多。

为了进一步隔离根本原因，我们在训练期间冻结了注意力层，并观察到与未冻结运行相似的奖励动态。这表明学习主要由 MoE 层驱动，而注意力机制的贡献不如预期。此外，我们观察到推理引擎（支持 Attention-Sink 前向传播的 Triton 内核）和分布式训练栈（FlashAttention-v2）之间存在巨大的 token 级概率不匹配。这些观察结果促使我们对注意力机制进行更深入的调查。

FlashAttentionV3 中的 Attention Sink 支持

GPT-OSS 中使用的 Attention Sinks 是可学习的标量参数（每个注意力头一个），在 softmax 计算中充当“虚拟 token”。它们允许模型将注意力质量分配给学习到的 Sink，而不是强制将所有注意力集中在内容 token 上，这已被证明可以提高流式推理和滑动窗口注意力训练中的注意力稳定性。

经过深入调查，我们发现了几个主要问题：

verl 在 fsdp_worker 中硬编码了 FlashAttention v2，而它不支持 Attention Sinks。
FlashAttention v2 和 v3 均不支持 Attention Sink 的反向传播，因此即使启用了 FlashAttention v3，它也无法按预期工作。
由于前向传播尚未合并到原始的 FlashAttention v3 仓库中，我们利用了 vLLM FlashAttention 分支的前向传播，并实现了反向传播以计算 Sink 梯度。

标准注意力机制：

scores = QK^T / sqrt(d)               # [B, H, N_q, N_k]
probs = softmax(scores, dim=-1)       # Σ_j P_ij = 1
output = probs @ V                   # [B, H, N_q, d_v]

带 Sinks 的注意力机制 (GPT-OSS)：

scores = QK^T / sqrt(d)                               # [B, H, N_q, N_k]
combined = concat([scores, sink_param], dim=-1)       # [B, H, N_q, N_k+1]
probs = softmax(combined, dim=-1)                     # Σ_j P_ij + P_sink = 1
probs_content = probs[..., :-1]                       # 丢弃 sink 组件
output = probs_content @ V                           # [B, H, N_q, d_v]

关键区别： Sink 参与 softmax 归一化，但不贡献于输出。

数学公式： 第 $i$ 行中内容 token $j$ 的注意力权重定义为：

$P_{ij} = \frac{e x p ( S _{ij} )}{\sum _{j^{'} = 1}^{N_{k}} e x p ( S _{i j^{'}} ) + e x p ( S _{h} )}$

其中：

$S_{ij} = Q_{i} K_{j}^{⊤} / d$ 是注意力分数
$P_{ij}$ 是内容 token 的注意力权重
$S_{h}$ 是头部 $h$ 的可学习 Sink 参数

Sink 概率： Sink 概率被计算但不用于输出：

$P_{i, h} = \frac{e x p ( S _{h} )}{\sum _{j^{'} = 1}^{N_{k}} e x p ( S _{i j^{'}} ) + e x p ( S _{h} )}$

反向传播： 损失 $L$ 相对于 Sink 参数 $S_{h}$ 的梯度为：

$\frac{\partial L}{\partial S _{h}} = - \sum_{i} P_{i, h} (\frac{\partial L}{\partial S _{i, h}} - \sum_{j \in {1, \dots, N_{k}}} P_{ij} \frac{\partial L}{\partial S _{ij}})$

由于 Sink 被计算但不用于输出，其梯度 $\partial L / \partial S_{i, h} = 0$ 。因此，反向方程简化为：

$\frac{\partial L}{\partial S _{h}} = - \sum_{i} P_{i, h} (\sum_{j \in {1, \dots, N_{k}}} P_{ij} \frac{\partial L}{\partial S _{ij}})$

前向传播改编自 vLLM 的 FlashAttention 分支，我们实现了反向传播来计算 Sink 参数的梯度。

结果： 在 FlashAttention v3 中应用修复后，我们观察到 GPT-OSS-20B 在一系列强化学习任务中的收敛速度大幅加快。这些任务包括数学推理上的单轮 RL（GSM8K）、指令遵循（VerifyIf）以及带工具使用的多轮智能体 RL（ReTool）。在所有设置下，训练变得稳定并表现出稳定的奖励提升。

内存高效训练

缓解重复实例化 MoE 专家导致的 FSDP 内存爆炸

我们经常遇到的一个问题是 FSDP 前向传播期间分配了过多的内存，这导致在 16 个 H200 节点上训练 GPT-OSS-20B bf16 模型时反复出现内存不足（OOM）故障。对于一个 20B 参数的 MoE 模型来说，这种行为是非常出乎意料的。

我们将问题定位为 Hugging Face Transformers 中 MoE 前向路径的两种不同实现。当 verl 在 FSDP 下计算对数概率时，会触发推理前向路径。在当前的 Hugging Face 实现中，该路径会为所有专家复制隐藏状态并执行批量矩阵乘法，从而在 GPU 内存中实例化极大的张量。相比之下，训练前向路径使用 for 循环顺序处理每个专家，然后组合结果。虽然速度较慢，但这种方法的内存效率要高得多。

我们对 Hugging Face 的实现进行了补丁修复，以使用更节省内存的执行路径，避免重复实例化专家。

结合 Flash Attention V3 的序列并行

智能体 RL 要求智能体在保持不断扩展的上下文的同时，在多个步骤中与环境交互。每个步骤的观察和环境反馈都会附加到上下文中，并用作后续决策的输入，这给训练期间的内存效率和可扩展性带来了重大挑战。

在完全分片数据并行（FSDP）下，模型参数、优化器状态和梯度在整个世界大小（即训练集群中的所有 GPU）上进行分片。每个 GPU 仅存储和更新其分配的参数分片，而推演数据则复制到所有 GPU 上——这意味着每个 GPU 都要处理每次推演的完整智能体交互历史。

序列并行（或上下文并行）通过跨设备划分输入序列，进一步降低了每块 GPU 的峰值激活内存。随着序列并行维度的增加，每块 GPU 的最大激活内存相应减少。我们实现了感知 Attention-Sink 并兼容 FlashAttention v3 的序列并行。

序列并行沿序列维度扩展，以减少每块 GPU 的激活占用。通过移除填充 token，来自所有序列的输入 token 被打包成一个连续的列表，同时使用位置 ID 来区分属于不同序列的 token。这种设计自然受益于 FlashAttention 对可变长度的支持。对于序列并行，除注意力层外的其他层没有位置间的依赖关系；因此，它们不需要每个 GPU 持有完整的序列分片，这些层也不需要额外的通信。

然而，注意力层要求属于同一序列的所有 token 都存在于同一 GPU 上，以便正确计算注意力权重。为了满足这一约束，我们在注意力头级别进行拆分，执行 all-to-all 通信来收集序列元素。这种设计避免了注意力计算本身的内部通信。在注意力层之后，单次 all-to-all 通信将输出重新分配回其原始的序列并行布局，之后剩余的非注意力层可以继续进行而无需进一步同步。

结论

我们为 GPT-OSS 骨干模型启用智能体 RL 训练的旅程是一次实践回顾，它突显了在开源 LLM 中解锁高级功能需要细致、深入的工程设计。

我们做出的贡献改变了 GPT-OSS 在智能体应用中的可行性，具体包括：

稳定 PPO： 我们贡献了一个修复方案来恢复 On-Policy 完整性，覆盖了由 MoE 架构的非确定性引起的对数概率不匹配。
启用 Attention Sink 支持： 我们成功实现并将 Attention Sink 反向传播集成到 FlashAttention v3 中，纠正了以前导致不稳定和收敛缓慢的灾难性训练-推理不匹配。
扩展内存效率： 我们引入了关键的内存优化，包括修补 MoE 实例化过程，以及将序列并行与新的 Attention Sink 支持相集成，从而实现了多步智能体必不可少的长上下文窗口训练。

这些工程努力验证了 GPT-OSS 可以作为一个可扩展且高性能的骨干网络，用于构建下一代智能、多步决策智能体。

wsl-docs

探索