gpt-oss 强化学习

您现在可以对 OpenAI 进行训练 gpt-oss 通过 RL 和 GRPO 使用 Unsloth。Unsloth 现在提供 最快的推理 （快 3 倍）， 最低的显存使用 （少 50%）和 最长的上下文 （长 8 倍），在 gpt-oss RL 相对于任何实现上均无准确性退化。 由于 gpt-oss 的强化学习（RL）尚未与 vLLM 兼容，我们不得不将 Transformers 的推理代码重写，以在 gpt-oss 上实现约 21 令牌/秒的 3 倍更快推理。对于 BF16，Unsloth 也实现了最快的推理（约 30 令牌/秒），尤其是在显存使用方面相对更高效，使用的显存比任何其他 RL 实现少 50%。我们计划在 vLLM 与 RL 兼容后支持我们的 50% 权重共享功能 。

• 免费笔记本： gpt-oss-20b GRPO Colab 笔记本 此笔记本会自动创建 更快的矩阵乘法内核 并使用 4 个新的 Unsloth 奖励函数。我们还展示了如何 对抗奖励作弊 ，这是 RL 面临的最大挑战之一。\

  

使用 Unsloth，您可以在 15GB 显存上使用 GRPO 训练 gpt-oss-20b，并且 免费 在 Colab 上训练。我们引入了嵌入下沉（embedding offloading），通过 offload_embeddings来减少约 1GB 的使用。Unsloth 的新推理在 任何 GPU（包括 A100、H100 及旧的 T4）上运行得更快。gpt-oss-120b 可以很好地适配到 120GB 显存的 GPU 上。

Unsloth 是唯一支持 gpt-oss 的 4 位 RL 的框架。所有性能提升都归功于 Unsloth 独特的 权重共享, 灵活注意力（Flex Attention）, 待机 和自定义内核。

提醒： Flash Attention 3（FA3） 不适合 gpt-oss 训练 ，因为它当前不支持 attention sinks 的反向传播，导致 训练损失不正确。如果您 没有 使用 Unsloth，FA3 可能会被 默认启用，因此请务必再次确认它未被使用！ 禁用 FA3 将导致 O(N^2) 的内存使用增加，因此 Unsloth 是唯一通过我们 Flex attention 实现为 gpt-oss 提供 O(N) 内存使用的 RL 框架。

⚡ 使推理快得多

推理在 RL 训练中至关重要，因为我们需要它在最大化某个奖励函数之前生成候选解（请参见此处 有关更详细的解释）。为了在没有 vLLM 的情况下为 gpt-oss 实现最快的推理速度，我们重写了 Transformers 的推理代码并集成了许多创新，包括像 Unsloth 这样的自定义算法 灵活注意力（Flex Attention），在 torch.compile 中使用特殊标志（如组合内核）。我们为 gpt-oss 编写的新推理代码已针对一个已经优化过的基线（比原生 Transformers 快 2 倍）进行了评估。

vLLM 不支持 gpt-oss 的 RL，因为它缺乏对 gpt-oss 的 BF16 训练和 LoRA 支持。没有 Unsloth，只有通过全精度 BF16 的训练可行，这使得 内存使用 提高 800% 以上。大多数框架默认启用 FA3（Flash Attention 3）（它可以减少显存使用并提高速度） ，但这会导致训练损失不正确。请参见 FA3 仓库的 Issue 1797 。您必须禁用 FA3，因为它会阻止长上下文训练，因为 FA3 使用 O(N) 内存，而朴素 attention 会以 O(N^2) 规模膨胀。为了使 attention sinks 可微，我们实现了 Unsloth Flex Attention.

我们通过对 BitsandBytes 的 4 位基准测试以及单独的 BF16 测试来评估 gpt-oss 的 RL 推理。Unsloth 的 4 位推理约快 4 倍，BF16 在显存使用方面也更高效。

Unsloth 的 gpt-oss RL 最棒的一点是它可以在任何 GPU 上运行，即使是不支持 BF16 的 GPU。我们的免费 gpt-oss-20b Colab 笔记本使用的是较旧的 15GB T4 GPU，因此推理示例也能很好地运行！

🛠️ gpt-oss Flex Attention 的问题和怪癖

我们不得不按照 此处所述 更改对 attention sinks 的实现，以允许左填充（left padding）情况下的生成工作。我们必须获取 logsumexp 并应用 sigmoid 激活以如下方式改变注意力权重：

$$A(X) = \sigma \bigg( \frac{1}{\sqrt{d}}QK^T \bigg)V \\ A(X) = \frac{\exp{\frac{1}{\sqrt{d}}QK^T}}{\sum{\exp{\frac{1}{\sqrt{d}}QK^T}}}V \\ \text{LSE} = \log{\sum{\exp{\frac{1}{\sqrt{d}}QK^T}}} \\ A_{sinks}(X) = A(X) \odot \sigma (\text{LSE} - \text{sinks})$$

在 gpt-oss 中，推理时的左填充掩码也是一个棘手问题。我们发现不仅要在生成令牌时考虑 KV Cache 的预填充，还要考虑批量生成中每个提示（prompt）中不同数量的填充符，这会改变我们存储块掩码的方式。此类示例可以见于下面：

普通因果掩码：

   k0 k1 k2 k3 k4   <-- 键
q0  X
q1  X  X
q2  X  X  X
q3  X  X  X  X
q4  X  X  X  X  X   <-- 最后一行查询（对解码最重要）

在一般情况下的推理（解码）

    k0 k1 k2 k3 k4
q0
q1
q2
q3
q4   X  X  X  X  X

如果我们天真地使用相同的掩码策略，这将失败：

    k0 k1 k2 k3 k4
q0
q1
q2
q3
q4   X   （注意 q4 在当前设置中有 q_idx=0，因为这是当前的第一个查询）

在生成（解码）阶段，我们通常只关心注意力矩阵的最后一行，因为只有一个查询令牌会关注所有先前的键令牌。如果我们天真地应用因果掩码（q_idx ≥ k_idx），这会失败，因为我们的单个查询的索引为 0，而有 n_k 个键令牌。为了解决这个问题，我们需要在掩码创建中加入偏移来决定应关注哪些令牌。但天真的方法很慢，因为偏移在每一步都会改变，迫使掩码和内核重新生成。我们通过缓存和编译优化解决了这个问题。

更困难的部分是批量生成。序列长度不同，因此填充使掩码创建复杂。Flex Attention 面临许多 挑战 且动态掩码很难处理。更糟的是，如果未编译，它会回退到非融合的 eager attention 实现，这很慢且占用大量内存（在序列长度上为二次复杂度，而不是线性）。

引用自 https://github.com/meta-pytorch/attention-gym/issues/15#issuecomment-2284148665

您需要使用 _compile=True 调用此函数。我们本质上将您的块掩码映射到一个完整的 Q_LEN x KV_LEN 矩阵以生成块掩码。没有编译，我们需要具体化这个完整矩阵，这可能会在长序列上导致 OOM。

此外，您需要运行 flex_attention = torch.compile(flex_attention)。如果不编译，flex 会回退到一个非融合的 eager 实现，这在调试时很有用，但它要慢得多并会具体化完整的分数矩阵。

最终，掩码必须动态地处理带有 KV Cache 的预填充与解码、批处理和每个序列的填充符，保持 torch.compile 友好，并支持滑动窗口。

🔍 Flash Attention 调查

我们探索的另一个有趣方向是尝试集成 Flash Attention。其优点被广泛认可，但一个限制是它在 gpt-oss 的反向传播中不支持 attention sinks。为了解决这个问题，我们重构了注意力机制，使其仅作用于 attention 输出和 FlashAttention 易于提供的 logsumexp 值。鉴于这些好处，这似乎是一个显而易见的选择。

然而，我们很快就注意到问题。尽管前几层表现如预期，但后面的层，尤其是第 18 层到第 24 层，产生的输出与 transformers 的 eager 模式实现相比出现了显著偏差。重要的是，这种差异不能归因于误差累积，因为每一层的输入在两种方法中都是相同的。为进一步验证，我们还将结果与 Unsloth 的 FlexAttention.

进行了比较。这需要进一步调查为什么只有最后几层在 flash attention 实现与其他实现之间表现出如此剧烈的差异。

Flash Attention 3 不支持 attention sinks 的反向传播

FA3 在大多数训练包中通常默认启用（但 Unsloth 除外），但这对于 gpt-oss 来说是不正确的。使用 FA3 会使训练损失完全错误，因为 FA3 不支持 gpt-oss 的 attention sinks 的反向传播。许多人仍然不知道这一点，所以请务必谨慎！

⚠️ 我们能否对抗奖励作弊？

RL 的最终目标是最大化某个奖励（比如速度、收入或某个指标）。但 RL 可能会 作弊。 当 RL 算法学会某个技巧或利用某种方法来提高奖励，而实际上并没有完成最终任务时，这被称为“奖励作弊".

”。这就是模型学会修改单元测试以通过编程挑战的原因，这些是现实部署中的关键阻碍。一些其他很好的示例可以在 维基百科.

中找到。 在我们的 免费 gpt-oss RL 笔记本

🏆奖励作弊

中，我们探讨了如何在代码生成场景中对抗奖励作弊，并展示了针对常见错误模式的切实可行的解决方案。我们看到模型会修改计时函数、外包给其他库、缓存结果，并彻底作弊。对抗之后，结果是我们的模型生成真正优化的矩阵乘法内核，而不是巧妙的作弊手段。

在 RL 过程中常见的一些奖励作弊示例包括：

懒惰

RL 学会使用 Numpy、Torch 等其他库，这些库调用了优化的 CUDA 内核。我们可以通过检查生成的代码是否导入了其他非标准的 Python 库来阻止 RL 算法调用优化代码。

缓存与作弊

RL 学会缓存输出结果，并通过检查 Python 全局变量来找到实际输出。

我们可以通过用一个大的伪矩阵清空缓存来阻止 RL 算法使用缓存数据。我们还必须通过多次循环和轮次仔细基准测试。

作弊 RL 学会编辑计时函数，使其输出经过的时间为 0。我们可以通过限制其 以及 局部变量全局变量 来阻止 RL 算法使用全局或缓存变量。我们还将使用 exec 来创建函数，因此我们必须将输出保存到一个空字典中。我们还通过types.FunctionType(f.__code__, {})

来禁止全局变量访问。

\ 教程：如何用 RL 训练 gpt-oss 大型语言模型在涉及复杂环境的任务中经常表现不佳。然而，通过应用 奖励函数强化学习

（RL）并设计定制的 gpt-oss ，这些挑战可以被克服。 GRPO RL 可以适配用于自动内核或策略创建等任务。本教程展示了如何使用

和 Unsloth 来自主地击败 2048。

2048 笔记本 （OpenAI 官方示例）	内核生成 笔记本

我们的笔记本已包含有关如何完成整个过程的逐步指南。

• 您将构建的内容：

• 训练 gpt-oss-20b，使模型能够自动赢得 2048

• 创建模型可以交互的最小 2048 环境 定义 奖励函数

  1. ，其：

  2. 检查生成的策略是否能编译并运行，

  3. 防止奖励作弊（禁止外部导入），并且

• 奖励真实的游戏胜利

运行推理并导出模型（MXFP4 4 位或合并的 FP16） 硬件： 2048 示例可以在免费 Colab T4 上运行，但训练会很慢。A100/H100 要快得多。4 位加载 + LoRA 使您能够在适度的显存中装入 20B 模型