GPT-OSS 智能体强化学习训练实战复盘
智能体强化学习（Agentic RL）扩展了传统的大语言模型训练方式，它优化的不是单轮响应，而是通过训练期间与环境的直接交互来学习整个决策过程。与传统依赖静态数据集的单轮强化学习或基于偏好的离线方法不同，智能体强化学习通过智能体在模拟或真实环境中规划行动、调用工具、观察结果并在多步轨迹中调整行为来主动收集在线数据，从而训练策略。这种交互驱动的优化将功劳分配给长期决策，其中查询重构、工具选择和执行顺序等中间选择直接影响下游的成功。训练遵循一个迭代闭环：智能体与环境交互收集轨迹，计算这些轨迹的奖励，根据观察结果更新策略，然后使用更新后的策略驱动下一轮交互和数据收集，例如使用 GRPO 或 PPO 算法。
LinkedIn 是一家以 AI 为先的公司，构建了智能体来帮助专业人士取得更大成功。在这种场景下，模型必须基于不完整信息进行推理，与结构化服务交互，并在多步骤中适应用户不断变化的意图，而不是产生单一的静态响应。这些能力对于支持招聘人员、求职者、知识寻求者和学习者等最终用户目标的智能体尤为关键，例如检索信息、优化查询、协调工具和执行多步骤工作流。通过交互学习稳健的决策策略，智能体强化学习为通过端到端优化构建可扩展、可靠和适应性强的 AI 系统提供了原则性基础。
GPT-OSS 模型已显示出与 OpenAI o3-mini 和 o4-mini 相当的性能，但其对智能体强化学习训练的适用性尚未得到验证。最近的大多数工作都集中在没有工具调用的微调上。这篇博客探讨了将 GPT-OSS 解锁为智能体应用潜在骨干模型的智能体强化学习训练之旅。
在我们的实验中，我们使用 verl 作为训练框架，因为它是开源社区中最受欢迎的框架之一。我们使用了 gsm8k、Retool 任务和可验证指令跟随任务，这些在强化学习训练中常用。我们重点展示 GPT-OSS-20B 模型的实验结果，我们的注意力下沉修复也适用于 GPT-OSS-120B。另外使用 Qwen-2.5-32B 模型来基准化强化学习训练期间的标准指标趋势。
GPT-OSS 强化学习训练的挑战
verl 是团队使用的开源框架，团队之前曾合作并为其做出贡献，以帮助普及智能体强化学习训练。随着 GPT-OSS 中新的 Harmony 聊天模板的引入，第一步是确保训练框架完全支持 Harmony 所需的更新消息格式和对话语义。这一步有助于在新技术下保持轨迹生成、轨迹构建和工具解析的一致性和正确性。
团队使用 ReTool 作为验证代码正确性的代表性示例。ReTool 是一个智能体编码任务，要求模型在代码编译器工具的帮助下解决数学问题。这种设置允许模型专注于核心推理和算法逻辑，而将实际算术和执行委托给工具。在一个回合中，模型与代码工具交互多次，使用执行结果作为反馈来完善其解决方案。在轨迹结束时，模型产生最终答案，并据此计算奖励。
在初始训练运行中，我们观察到 KL 散度和熵爆炸，同时奖励没有增加，这表明 GPT-OSS 训练设置存在潜在问题，如图 1 所示。
平均梯度范数
平均奖励
图 1. 左：Qwen32b 的奖励显著高于 GPT-OSS 20B；右：随着训练进行，梯度范数爆炸。
verl 中的实用调试之旅：恢复 PPO 在线策略完整性
恢复 PPO 在线策略完整性：MoE 对数概率不匹配的修复
图 2. 即使在在线策略训练中，重要性采样裁剪值也非零。
我们专注于在线策略方法，因为它们提供更高的稳定性和更可靠的收敛性。纯在线策略近端策略优化（PPO）的基础要求重要性采样比必须恰好为 1。重要性比率的数学定义为：
ratio = π(a|s) / π_old(a|s)
这一要求确保策略更新仅对当前策略 π(a|s) = π_old(a|s) 生成的数据执行，防止意外的裁剪。
我们在 ReTool 训练中观察到了非零的裁剪值，如图 2 所示，这源于两个对数概率之间的不匹配：
- 当前对数概率
log_prob：log(π(a|s)) - 旧对数概率
old_log_prob：log(π_old(a|s))
根本原因：双重前向传递和 MoE 架构
在 verl 0.3.0 之前，实现依赖于对同一状态-动作对进行两次独立的前向传递（一次计算当前 log_prob，一次检索存储的 old_log_prob）。
在像 GPT-OSS 这样的混合专家（MoE）架构中，门控网络将输入路由到不同的专家。由于实现因素（例如细微的浮点差异或显式的随机性），专家路由在两次传递之间可能略有不同。感兴趣的读者可以进一步阅读《通过对齐训练和推理路由器来稳定 MoE 强化学习》。
这种路由差异导致：
log(π(a|s)) ≠ log(π_old(a|s))
由此产生的比率偏离 1，错误地触发 PPO 裁剪，违反了核心的在线策略假设。
解决方案：通过对数概率替换强制比率 = 1
该修复通过逻辑上覆盖有缺陷的计算来解决此问题，当已知环境是在线策略时（即当小批量大小等于全局批量大小时）：
if on_policy:
old_log_prob = log_prob.detach()
else:
old_log_prob = model_inputs["old_log_probs"]
通过将 old_log_prob 设置为等于新计算的 log_prob（分离以防止梯度流通过参考值），重要性比率在数学上被强制回到 1。这一策略绕过了 MoE 非确定性路由引起的不稳定性，并保证了 PPO 训练期间严格的在线策略行为。
修正训练-推理不匹配问题后，对数概率不匹配导致的采样权重裁剪比率降为零，但梯度范数仍持续爆炸，奖励也未见提升。为了定位问题，我们简化了训练任务，使用无需智能体工具调用的单步任务 GSM8K 进行测试。结果发现，同样的不稳定性依然存在（如图3绿色曲线所示），这表明在 verl 框架下对 GPT-OSS 进行基础 RL 训练存在根本性问题。
我们推测，训练-推理不匹配（Training–Inference Mismatch）可能是潜在原因：推理时，像 vLLM 和 SGLang 这样的引擎会为追求吞吐量进行激进优化；而训练时，FSDP 则优先考虑数值精度和稳定性。这种执行环境的差异，本质上可能将原本的同策略 RL 优化变成了异策略优化。
这篇博客详细解释了这种不匹配为何会导致梯度不稳定和奖励无法提升。图3对比了应用 rollout 修正前后的训练情况。应用修正后，训练动态显著改善，梯度范数保持稳定，不再爆炸。
然而，如图4左图所示，奖励提升幅度仍然有限，在简单的 GSM8K 任务上，其收敛速度相比更小的稠密模型变体依然慢得多。
为了进一步定位根本原因，我们在训练时冻结了注意力层，观察到奖励动态与未冻结时类似（图4中蓝色曲线与黄色曲线）。这表明学习主要由 MoE 层驱动，而注意力机制的有效贡献低于预期。此外，我们还观察到推理引擎与分布式训练栈之间存在显著的词元级概率不匹配，两者使用了不同的注意力内核。这些观察促使我们对注意力机制进行更深入的调查。
FlashAttentionV3 中的注意力汇支持
GPT-OSS 中使用的注意力汇是可学习的标量参数（每个注意力头一个），在 softmax 计算中充当“虚拟词元”。它允许模型将注意力权重分配给一个学习到的汇，而不是强制所有注意力都集中在内容词元上，这已被证明能改善流式推理和滑动窗口注意力训练中的注意力稳定性。
经过深入调查，我们发现了几个主要问题：
- verl 在
fsdp_worker中硬编码了 FlashAttention v2，而 v2 不支持注意力汇。 - FlashAttention v2 和 v3 均不支持注意力汇的反向传播，因此即使启用了 FlashAttention v3，它也无法按预期工作。
- 由于正向传播尚未合并到原始的 FlashAttention v3 仓库，我们利用了 vLLM 的 FlashAttention 分支中的正向传播，并实现了反向传播来计算汇的梯度。
标准注意力
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]
带汇的注意力（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] # Drop sink component
output = probs_content @ V # [B, H, N_q, d_v]
关键区别：汇参与 softmax 归一化，但不贡献给输出。
数学公式
内容词元 j 在第 i 行的注意力权重定义为：
Pij = exp(Sij) / ( Σ_{j'=1}^{N_k} exp(Sij') + exp(Sh) )
其中：
- Sij = Qi Kj⊤ / √d 是注意力分数
- Pij 是内容词元的注意力权重
- Sh 是头 h 的可学习汇参数
汇概率： 汇概率被计算但不用于输出：
Pi,h = exp(Sh) / ( Σ_{j'=1}^{N_k} exp(Sij') + exp(Sh) )
反向传播
损失 L 对汇参数 Sh 的梯度为：
∂L/∂Sh = - Σ_i Pi,h ( ∂L/∂Si,h - Σ_{j ∈ {1,…,N_k}} Pij ∂L/∂Sij )
其中：
- Pi,h 是第 i 行的汇注意力概率
- ∂L/∂Sij 是对注意力分数（包括汇）的梯度
简化梯度： 由于汇被计算但不用于输出，其梯度 ∂L/∂Si,h = 0。
因此，反向传播方程简化为：
∂L/∂Sh = - Σ_i Pi,h ( Σ_{j ∈ {1,…,N_k}} Pij ∂L/∂Sij )
正向传播改编自 vLLM 的 FlashAttention 分支，我们实现了反向传播来计算汇参数的梯度。该实现将在内部审查流程后发布。
结果
在 FlashAttention v3 中应用修复后，我们观察到 GPT-OSS-20B 在一系列强化学习任务上的收敛速度显著加快。这些任务包括数学推理的单轮 RL（GSM8K——图5中的红色曲线）、指令遵循（VerifyIf，在域外多 if 基准上评估——图6），以及使用工具调用的多轮智能体 RL（ReTool——图7）。
在所有设置中，训练变得稳定，并展现出持续的奖励提升。
内存高效训练
缓解因重复 MoE 专家实例化导致的 FSDP 内存爆炸
我们反复遇到的一个问题是 FSDP 前向传播期间内存分配过高，导致在 16 个 H200 节点上训练 GPT-OSS-20B bf16 模型时反复出现内存不足故障（最大响应长度：16k，提示长度：8k）。对于一个 200 亿参数的 MoE 模型来说，这种行为极不寻常。
2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m File "/home/jobuser/.local/lib/python3.10/site-packages/transformers/models/gpt_oss/modeling_gpt_oss.py", line 123, in forward
2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m hidden_states = hidden_states.repeat(num_experts, 1)
2025-11-27T11:15:27.927Z [36m(TaskRunner pid=32081)[0m torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 180.00 GiB. GPU 0 has a total capacity of 139.72 GiB of which 110.94 GiB is free. Process 685851 has 24.88 GiB memory in use. Process 692458 has 3.87 GiB memory in use. Of the allocated memory 23.28 GiB is allocated by PyTorch, and 84.43 MiB is reserved by PyTorch but unallocated.
我们发现问题源于 Hugging Face Transformers 中 MoE 前向传播路径的两种不同实现。其他用户也报告了此问题。当 verl 在 FSDP 下计算对数概率时，会触发推理前向传播路径。在当前 Hugging Face 的实现中，该路径会为所有专家复制隐藏状态并执行批处理矩阵乘法，从而在 GPU 内存中实例化极大的张量。相比之下，训练前向传播路径使用 for 循环依次处理每个专家，然后合并结果。虽然速度较慢，但这种方法的内存效率要高得多。
@GPUMemoryLogger(role="dp actor", logger=logger)
def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:
"""
....
"""
# set to eval, this essentially prioritizes parallelism at the cost of memory efficiency
self.actor_module.eval()
...
我们修补了 Hugging Face 的实现，使用更内存高效的执行路径，避免重复实例化专家。
序列并行与 Flash Attention v3
智能体强化学习（Agentic RL）要求智能体在多步交互中持续扩展上下文。每一步的环境观察和反馈都会追加到上下文，作为后续决策的输入，这在训练时带来了巨大的内存效率和可扩展性挑战。
在全分片数据并行（FSDP）下，模型参数、优化器状态和梯度会在整个训练集群的所有 GPU 间分片。每个 GPU 只存储和更新自己负责的参数分片，但 rollout 数据会在所有 GPU 上复制——这意味着每个 GPU 都要处理每个 rollout 的完整交互历史。
前向传播时，如果计算到某个层的参数不在本地，就会触发一次 all_gather 操作，在所有 GPU 间重建完整参数。反向传播时，对应的 reduce_scatter 操作会聚合梯度，确保每个 GPU 只保留自己的本地分片。这提供了一定的扩展性：GPU 数量增加，每个 GPU 的内存占用就会减少。
FSDP 通过分片模型参数、梯度和优化器状态，实现了模型层面的扩展。序列并行（或称上下文并行）则更进一步，它将输入序列在设备间分区，从而降低每个 GPU 的峰值激活内存。
随着序列并行维度的增加，每个 GPU 的最大激活内存会相应减少。我们实现的序列并行方案是注意力池感知的，并且兼容 FlashAttention v3（图 8，右）。
图 8。左：无序列并行的推理。右：带序列并行的推理，在注意力层前后执行额外的 all-to-all 通信。这会将序列在并行工作节点间分区，将注意力计算的峰值内存占用降低一个与序列并行度成比例的量。
序列并行沿着序列维度扩展，以减少每个 GPU 的激活内存占用。所有序列的输入 token 会被打包成一个连续的列表（移除填充 token），同时使用位置 ID 来区分属于不同序列的 token。这种设计天然受益于 FlashAttention 的变长支持。对于序列并行，注意力层以外的层没有位置间依赖，因此它们不需要每个 GPU 持有完整的序列分片，这些层无需额外通信。
但注意力层要求属于同一序列的所有 token 都在同一个 GPU 上，才能正确计算注意力权重。为了满足这个约束，会执行一次 all-to-all 通信来收集序列元素，分区在注意力头级别进行。这种设计避免了在注意力计算内部进行通信，否则通信开销会变得难以承受。注意力层之后，一次 all-to-all 通信将输出重新分发回原始的序列并行布局，之后剩余的非注意力层就可以在没有进一步同步的情况下继续执行。
结论
我们为 GPT-OSS 骨干模型启用智能体强化学习训练的旅程，是一次务实的回顾。它表明，要解锁开源大语言模型的高级能力，需要细致入微、深入底层的工程实践。
我们的贡献提升了 GPT-OSS 在智能体应用中的可行性，具体包括：
- 稳定 PPO： 我们贡献了一个修复，恢复了在线策略的完整性，覆盖了 MoE 架构非确定性导致的 log-probability 不匹配问题（图 2）。
- 启用注意力池支持： 我们成功实现并将注意力池反向传播集成到 FlashAttention v3 中，纠正了之前导致训练不稳定和收敛缓慢的灾难性训练-推理不匹配问题（图 5、6、7）。
- 扩展内存效率： 我们引入了关键的内存优化，包括修补 MoE 物化过程，以及将序列并行与新的注意力池支持集成，从而能够训练多步智能体所需的长上下文窗口（图 8）。
这些工程努力验证了 GPT-OSS 作为一个可扩展、高性能的骨干模型，适合构建下一代智能、多步决策智能体。
致谢
感谢 Deepak Agarwal、Bee-Chung Chen、Animesh Singh、Gungor Polatkan、Balaji Krishnapuram 和 Jitendra Agarwal 的领导支持。
参考文献
- Feng, Jiazhan, et al. Retool: Reinforcement Learning for Strategic Tool Use in LLMs. arXiv preprint arXiv:2504.11536 (2025).
- Xiao, Guangxuan, et al. Efficient Streaming Language Models with Attention Sinks. arXiv preprint arXiv:2309.17453 (2023).
- When Speed Kills Stability: Demystifying RL Collapse from the Training–Inference Mismatch.
https://yingru.notion.site/When-Speed-Kills-Stability-Demystifying-RL-Collapse-from-the-Training-Inference-Mismatch-271211a558b7808d8b12d403fd15edda
觉得有用？分享给更多人