惊爆！GPT-OSS-20B强化学习，训练速度狂飙、超稳定！

近年来，人工智能技术发展迅猛，尤其以大型语言模型（LLMs）为代表的智能系统，正逐步从简单的对话响应，走向更复杂的决策与行动。当我们谈论AI时，我们常常想象的是一个能够独立思考、规划并执行任务的智能体，而“智能体强化学习（Agentic RL）”正是实现这一愿景的关键路径。它超越了传统的模型训练范式，不再仅仅优化单一的应答，而是通过与环境的直接互动，学习并优化一个完整的决策过程。

想象一下，一个智能体不再只是给出答案，而是像一位经验丰富的向导，能够根据实际情况，制定策略，调用合适的工具，观察结果，并在多步骤的交互中不断调整自己的行为。这种基于交互的优化方式，能够对长期的决策链进行有效的“归因”，即使是中间环节的选择，比如如何重新组织查询、选择哪种工具、执行的顺序如何，都会直接影响最终任务的成功。整个训练过程形成了一个迭代的闭环：智能体与环境互动，收集行动轨迹数据，根据这些轨迹计算奖励，然后依据观察到的结果更新决策策略，并利用更新后的策略驱动下一轮的互动和数据收集，这正是PPO等先进算法的魅力所在。

在新一代AI应用中，例如领英（美国一家知名职业社交平台）这样的“AI优先”公司，就致力于构建能帮助专业人士更成功的智能体。在这样的背景下，模型不仅要能够处理不完整的信息，与结构化服务进行交互，还需要在多步骤的流程中适应用户不断变化的需求，而不是仅仅生成一个静态的响应。这些能力对于支持招聘人员、求职者、知识探索者和学习者等终端用户的智能体来说，显得尤为关键。比如，它们可能需要检索信息、优化查询、协调多种工具，并执行多步骤的工作流。通过交互学习出稳健的决策策略，智能体强化学习为构建可扩展、可靠、适应性强的AI系统提供了坚实的基础，实现了端到端的优化。

GPT-OSS作为一款备受关注的开源大型语言模型，在单轮响应方面已经展现出与外媒提及的知名模型（如OpenAI o3-mini和o4-mini）不相上下的性能。然而，它是否同样适用于智能体强化学习的训练，此前一直没有得到充分的验证。目前大多数研究工作都集中在不涉及工具调用的微调上，比如利用GPT-OSS与某开源社区的Hugging Face Transformers框架进行微调。

本次，我们将深入探讨解锁GPT-OSS模型进行智能体强化学习训练的全过程，力求将其打造成未来智能体应用的核心骨干模型。在我们的实验中，研究团队选择了verl这一在开源社区中广泛应用的训练框架。同时，我们采用了gsm8k数学推理任务、Retool工具使用任务以及可验证指令遵循任务，这些都是强化学习训练中常见的基准。我们重点关注GPT-OSS-20B模型的实验结果，并且我们对“注意力槽”机制的修复同样适用于更大的GPT-OSS-120B模型。此外，Qwen-2.5-32B模型也被用作参照，以对比强化学习训练中各项指标的常规趋势。新媒网跨境获悉，这一系列探索旨在推动开源大模型在更复杂应用场景中的落地。

GPT-OSS强化学习训练之挑战重重

verl框架作为团队一直使用的开源工具，我们此前也曾参与并贡献，共同推动智能体强化学习训练的普及。然而，当新的“Harmony”对话模板引入GPT-OSS后，首要任务便是确保verl训练框架能够完全兼容这一更新后的消息格式和对话语义。这一步至关重要，它能确保在新的模板下，行动轨迹的生成、轨迹的构建以及工具解析都能够保持一致性和准确性。

团队选用ReTool任务作为一个具有代表性的智能体编程案例来验证代码的正确性。在ReTool任务中，模型需要在一个代码编译器的辅助下解决数学问题。这种设置让模型能够专注于核心的推理和算法逻辑，而将实际的算术运算和执行任务委托给工具完成。在一个任务回合中，模型会多次与代码工具交互，并利用执行结果作为反馈来优化其解决方案。在整个行动轨迹的末尾，模型会给出一个最终答案，并以此答案来计算奖励。

然而，在最初的几次训练运行中，我们观察到了令人担忧的现象：KL散度和熵值呈现爆炸式增长，而奖励却迟迟没有提升。这些都清晰地指向GPT-OSS训练设置中存在着深层次的问题，如同以下图表所示，Qwen32B在奖励表现上明显优于GPT-OSS 20B，并且随着训练的进行，梯度范数也出现了惊人的爆炸。这无疑为我们的研究带来了不小的挑战，迫使我们必须深入探究其背后的原因。

图1：左侧图显示Qwen32B模型在奖励方面显著高于GPT-OSS 20B；右侧图则展示了随着训练进程，梯度范数出现了爆炸式增长。

verl框架中的实际调试之旅：重塑PPO在线策略的完整性

重塑PPO在线策略完整性：修复MoE对数概率不匹配问题

图2：即使是在线策略训练，重要性采样剪辑值也非零。

在强化学习的深奥世界里，我们特别偏爱在线策略（on-policy）方法，因为它们以卓越的稳定性与可靠的收敛性著称。在纯粹的在线策略近端策略优化（PPO）算法中，有一个铁律般的要求：重要性采样比（importance sampling ratio）必须精确地等于1。这个比率的数学定义如下：

$ \text{ratio} = \frac{\pi(a \mid s)}{\pi_{\text{old}}(a \mid s)} $

这个严苛的要求确保了策略的更新仅基于由当前策略自身生成的数据$\pi(a \mid s) = \pi_{\text{old}}(a \mid s)$，从而有效避免了不必要的剪裁（clipping）。然而，在我们的ReTool任务训练中，却意外地观察到非零的剪辑值，如图2所示。这正是源于当前对数概率log_prob ($\log(\pi(a \mid s))$)与旧对数概率old_log_prob ($\log(\pi_{\text{old}}(a \mid s))$)之间的细微不匹配。

根本原因：双重正向传播与MoE架构的挑战

在verl框架0.3.0版本之前，其实现方式依赖于两次独立的正向传播：一次用于计算当前的log_prob，另一次则用于检索并存储old_log_prob，尽管它们都是针对同一对状态-动作。然而，在像GPT-OSS这样的“混合专家”（Mixture of Experts, MoE）架构中，其门控网络会将输入路由到不同的专家子网络。由于实现层面的一些微妙因素（例如浮点运算的微小差异，或是显式的随机性），两次正向传播的专家路由结果可能存在细微的差异。有兴趣的读者可以进一步了解《通过对齐训练和推理路由器来稳定MoE强化学习》这篇论文。正是这种路由差异，导致了：

$ \log(\pi(a \mid s)) \neq \log(\pi_{\text{old}}(a \mid 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（并进行detach()操作，以防止梯度流经这个参考值），我们从数学上强制将重要性比率恢复为1。这一策略巧妙地绕过了MoE架构非确定性路由所带来的不稳定性，从而确保了PPO训练过程中严格的在线策略行为。这个看似简单的改动，却为GPT-OSS模型在智能体强化学习领域迈出了坚实的一步，为后续的稳定训练奠定了基础。

纠正训练与推理不一致问题

尽管修复了对数概率不匹配的问题，并将重要性采样剪辑比率成功降至零，但我们仍然面临着严峻的挑战：梯度范数持续爆炸，奖励也未能如期提升。为了更精确地定位问题，我们进一步将训练任务简化为GSM8K，这是一个不涉及智能体工具使用的单步骤任务。然而，即使在这样的简化场景下，同样的不稳定性依然存在，如图3中绿色曲线所示，这表明GPT-OSS在verl框架下的基础强化学习训练存在根本性问题。

我们猜测，训练与推理之间存在不匹配可能是潜在的根源。在推理阶段，像外媒提及的vLLM和SGLang这样的引擎为了追求吞吐量会进行激进的优化；而在训练阶段，FSDP（全分片数据并行）则更注重数值精度和稳定性。这种差异，实际上可能将原本应是在线策略的强化学习，悄然转变为一种离线策略优化。之前有文章详细阐述了为何此类不匹配会导致梯度不稳定和奖励难以提升。图3对比了有无“轨迹校正”（rollout correction）的训练运行情况（关于轨迹校正的更多细节，可参考verl团队的另一篇文章）。

应用轨迹校正后，训练动态确实得到了显著改善，梯度范数保持稳定，不再出现爆炸。然而，如图4左图所示，奖励的提升幅度仍然不大，并且在简单的GSM8K任务上，其收敛速度与更小的密集模型相比，依然显著缓慢。这提醒我们，虽然取得了部分进展，但深层次的问题仍需进一步探索。

图3：不同训练配置下的梯度范数行为。绿色曲线：未进行轨迹校正的训练，梯度不稳定；红色曲线：为将问题隔离到注意力机制，冻结了注意力层，结果部分稳定；蓝色曲线：启用轨迹校正（序列级重要性采样）的训练，梯度范数稳定。



图4：左图：即使应用了轨迹校正，GSM8K任务上的奖励提升依然缓慢，性能与训练中冻结注意力层的情况相近。右图：在推理引擎（SGLang采用支持注意力槽前向传播的Triton内核）和训练栈（FSDP采用FlashAttention-v2）之间观察到显著的对数困惑度不匹配，表明存在大的训练-推理不一致。

为了进一步深入探究问题的根源，我们尝试在训练过程中冻结注意力层，结果观察到奖励动态与未冻结注意力层时的表现（图4中的蓝色曲线与黄色曲线）相似。这暗示着模型的学习主要由MoE（混合专家）层驱动，而注意力机制的贡献却远未达到预期效果。此外，我们还在推理引擎与分布式训练栈之间观察到了显著的令牌级别概率不匹配，这主要是因为两者采用了不同的注意力核（attention kernels）。这些观察结果共同促使我们对注意力机制进行了更深层次的调查。

FlashAttentionV3对注意力槽的支持

GPT-OSS模型中使用的“注意力槽”（Attention Sinks）是一种巧妙的设计：它们是可学习的标量参数，每个注意力头都对应一个，在softmax计算中扮演着“虚拟令牌”的角色。这一机制允许模型将一部分注意力权重分配给这些可学习的“槽”，而不是强制所有注意力都集中在内容令牌上。实践证明，这能有效提升流式推理和滑动窗口注意力训练中的稳定性。

在经过一番深入细致的调查之后，我们识别出了几个关键问题：

• verl框架在fsdp_worker中硬编码了FlashAttention v2，但该版本并不支持注意力槽。
• FlashAttention v2和v3版本对注意力槽的反向传播（backward pass）支持不足，即使启用了FlashAttention v3，其效果也未能达到预期。
• 由于注意力槽的前向传播（forward pass）尚未合并到原始的FlashAttention v3代码库中，我们巧妙地借鉴了vLLM FlashAttention分支（PR #75）中的前向传播实现，并在此基础上成功实现了注意力槽的反向传播计算，从而解决了梯度计算的问题。

标准注意力机制回顾

传统的注意力计算模式如下所示：

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]

这里，Q、K、V分别代表查询、键和值，d是维度，softmax函数将分数转换为概率分布，确保所有内容令牌的注意力权重之和为1。

GPT-OSS中的带注意力槽的注意力机制

而在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]

核心区别在于，注意力槽参数sink_param会参与softmax归一化计算，但它本身并不直接贡献于最终的输出值。这意味着，注意力槽能够像一个“虚拟的”令牌一样，吸纳一部分注意力权重，从而在保持整体概率和为1的同时，避免所有注意力都过度集中在内容令牌上。

数学公式深入解读

对于行i中内容令牌j的注意力权重P_{ij}定义为：

$ P_{ij} = \frac{\exp(S_{ij})} {\sum_{j'=1}^{N_k} \exp(S_{ij'}) + \exp(S_h)} $

其中：

• $S_{ij} = Q_i K_j^\top / \sqrt{d}$ 是注意力分数。
• $P_{ij}$ 是内容令牌的注意力权重。
• $S_h$ 是注意力头h的可学习注意力槽参数。

注意力槽概率： 注意力槽的概率会被计算，但不会用于最终输出：

$ P_{i,h} = \frac{\exp(S_h)} {\sum_{j'=1}^{N_k} \exp(S_{ij'}) + \exp(S_h)} $

反向传播的实现

损失L对注意力槽参数S_h的梯度计算如下：

$ \frac{\partial L}{\partial S_h} = - \sum_i P_{i,h} \left( \frac{\partial L}{\partial S_{i,h}} - \sum_{j \in {1,\ldots,N_k}} P_{ij} \frac{\partial L}{\partial S_{ij}} \right) $

其中：

• $P_{i,h}$ 是行i的注意力槽概率。
• $\partial L/\partial S_{ij}$ 是损失对注意力分数的梯度，包括注意力槽的梯度。

简化的梯度计算： 由于注意力槽仅参与计算但不直接贡献于输出，其梯度 $\partial L/\partial S_{i,h} = 0$。因此，反向传播方程简化为：

$ \frac{\partial L}{\partial S_h} = - \sum_i P_{i,h} \left( \sum_{j \in {1,\ldots,N_k}} P_{ij} \frac{\partial L}{\partial S_{ij}} \right) $

新媒网了解到，这一关键的前向传播部分借鉴了vLLM的FlashAttention分支，而其反向传播则由团队独立实现，以计算注意力槽参数的梯度。目前，该实现正在内部审查流程中，待审核完成后便会对外发布。

显著的成果展示

在将FlashAttention v3中的修复方案应用之后，我们欣喜地看到GPT-OSS-20B模型在多种强化学习任务中展现出显著加速的收敛速度。这些任务涵盖了数学推理领域的单轮强化学习（GSM8K任务，如图5中的红色曲线，其收敛速度远超其他未修复版本），可验证指令遵循任务（VerifyIf，在域外多条件判断基准上进行评估，如图6所示），以及涉及工具使用的多轮智能体强化学习任务（ReTool，如图7所示）。

在所有这些场景中，训练过程变得异常稳定，并持续展现出稳健的奖励提升。这无疑证明了我们的修复工作取得了突破性进展，为GPT-OSS模型在复杂智能体应用中发挥其全部潜力奠定了坚实基础。

图5：单轮GSM8K任务，红色曲线（修复后）比其他未修复版本收敛速度快得多。

图6：在可验证指令遵循任务中，未修复版本（蓝色曲线）训练崩溃，而修复版本则展现出稳定的奖励提升。



图7：在Retool任务中，修复后的版本展现出稳定的奖励提升，且梯度没有爆炸（fa2代表未修复的Flash Attention 2，而fa3代表修复后的Flash Attention 3）。修复后，验证准确率得分显著提高。

内存高效训练：大模型时代的基石

缓解FSDP中因MoE专家重复实例化导致的显存激增

在我们的GPT-OSS-20B bf16模型训练过程中，一个反复出现的顽固问题是：在使用16个H200节点（最大响应长度16k，提示长度8k）时，FSDP（全分片数据并行）正向传播期间会发生异常的显存分配，并多次导致“内存不足”（OOM）错误。对于一个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: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正向路径的两种不同实现上。值得一提的是，这一问题也曾被其他用户报告过，例如在GitHub issue 40073中就有所提及。

当verl框架在FSDP环境下计算对数概率时，会触发推理路径的正向传播。在当前的Hugging Face实现中，这条路径会将隐藏状态为所有专家进行重复复制，并执行批量矩阵乘法，这导致在GPU内存中实例化了极其庞大的张量。相比之下，训练路径则采用了一种更为内存高效的方式：它通过一个循环顺序处理每个专家，然后再将结果组合起来。虽然这种方法速度稍慢，但在内存效率方面却表现出显著优势。

@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的序列并行技术

智能体强化学习要求智能体在多步骤交互过程中，持续维护并不断扩展上下文信息。每一步的观察结果和环境反馈都会被追加到上下文中，作为后续决策的输入。这无疑给训练过程中的内存效率和可扩展性带来了巨大的挑战。

在全分片数据并行（FSDP）模式下，模型参数、优化器状态和梯度都会在整个集群（即训练集群中的所有GPU）中进行分片。每个GPU只存储和更新其分配到的参数分片，而行动轨迹数据则在所有GPU上进行复制——这意味着每个GPU都要处理每个轨迹的完整智能体交互历史。在正向传播过程中，当计算到达某个层，其参数不在本地GPU上时，就会触发一次all_gather操作，将完整的参数实例化到所有GPU上。而在反向传播过程中，相应的reduce_scatter操作会聚合梯度，并确保每个GPU只保留其本地分片。这种机制提供了一定程度的扩展性：随着GPU数量的增加，每个GPU的内存占用会相应减少。FSDP通过分片模型参数、梯度和优化器状态，实现了模型级别的扩展。

而“序列并行”（Sequence Parallelism，或称上下文并行）则更进一步，通过在设备间划分输入序列，从而降低每个GPU的峰值激活内存。随着序列并行维度的增加，每个GPU的最大激活内存也会相应减少。我们已经成功实现了对注意力槽敏感并兼容FlashAttention v3的序列并行功能（如图8右图所示）。

图8：左图：未采用序列并行时的推理模式。右图：采用序列并行时的推理模式，在注意力层前后会进行额外的all-to-all通信。这种方式将序列在并行工作器之间进行划分，并以与序列并行度成比例的因子降低注意力计算的峰值内存占用。

序列并行技术沿序列维度进行扩展，以减少每个GPU的激活内存占用。通过移除填充令牌，来自所有序列的输入令牌被打包成一个连续的列表，同时使用位置ID来区分属于不同序列的令牌。这种设计天然地受益于FlashAttention对可变长度序列的支持。

对于序列并行而言，除了注意力层之外的其他层不具备位置间依赖性；因此，它们不需要每个GPU都持有一个完整的序列分片，也无需额外的通信。然而，注意力层为了正确计算注意力权重，要求属于同一序列的所有令牌必须位于同一个GPU上。为了满足这一约束，在注意力层前后会执行一次all-to-all通信来收集序列元素，并在注意力头级别进行拆分。这种设计避免了在注意力计算内部进行通信，否则这将是极其昂贵的。在注意力层之后，通过一次简单的all-to-all通信，将输出重新分布回其原始的序列并行布局，之后剩余的非注意力层便无需进一步同步即可继续进行。

结语

我们为GPT-OSS骨干模型解锁智能体强化学习训练的旅程，是一次充满实践意义的回顾。它清晰地表明，要充分发挥开源大型语言模型的先进能力，需要细致入微、深入骨髓的工程创新。我们团队的贡献，极大地提升了GPT-OSS在智能体应用领域的实用性，具体体现在以下几个关键方面：

1. 稳定PPO训练： 我们贡献了一项关键修复，成功恢复了在线策略的完整性，巧妙解决了由MoE架构非确定性引起的对数概率不匹配问题（如图2所示）。
2. 全面支持注意力槽： 我们成功实现了注意力槽的反向传播功能，并将其集成到FlashAttention v3中。这一突破性进展，彻底纠正了此前灾难性的训练与推理不一致问题，解决了导致模型不稳定和收敛缓慢的根本原因（如图5、图6和图7所示）。
3. 内存效率优化： 我们引入了至关重要的内存优化措施，包括修补MoE专家层的实例化过程，并创新性地将序列并行与新的注意力槽支持相结合。这些努力使得模型能够在长上下文窗口下进行高效训练，而这对于多步骤智能体来说是必不可少的能力（如图8所示）。

新媒网跨境认为，这些卓越的工程努力，不仅验证了GPT-OSS作为下一代智能、多步骤决策智能体骨干模型的可扩展性和高性能，更标志着我们在追求更强大、更实用AI的道路上迈出了坚实的一步。

致谢

感谢Deepak Agarwal、Bee-Chung Chen、Animesh Singh、Gungor Polatkan、Balaji Krishnapuram和Jitendra Agarwal等团队领导与核心专家的鼎力支持。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/gpt-oss-20b-rl-train-stable-fast-fix.html