gpt-oss 强化学习

你现在可以训练 OpenAI gpt-oss 使用 RL 和 GRPO 通过 Unsloth。Unsloth 现在提供 最快的推理 （快 3 倍）， 最低的 VRAM 占用 （减少 50%）以及 最长的上下文 （长 8 倍），用于 gpt-oss RL，相比任何实现都更优——且没有精度下降。 由于 gpt-oss 上的强化学习（RL）目前还不兼容 vLLM，我们不得不将推理代码从 Transformers 代码重写，以便为 gpt-oss 提供快 3 倍的推理速度，达到约 21 tokens/s。对于 BF16，Unsloth 也实现了最快的推理（约 30 tokens/s），尤其是在 VRAM 占用方面表现最佳，相比任何其他 RL 实现可减少 50% 的 VRAM。我们计划支持我们的 50% 权重共享功能 一旦 vLLM 变得与 RL 兼容。

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

  

使用 Unsloth，你可以在 15GB VRAM 上用 GRPO 训练 gpt-oss-20b，并且 免费 在 Colab 上运行。我们引入了嵌入卸载（embedding offloading），通过 offload_embeddings将使用量再减少 1GB。Unloth 的新推理在 任何 GPU 上运行更快，包括 A100、H100 和老旧的 T4。gpt-oss-120b 可以很好地装入 120GB VRAM 的 GPU 中。

Unsloth 是唯一支持 gpt-oss 4-bit RL 的框架。所有性能提升都归功于 Unsloth 独有的 权重共享, Flex Attention, 待命 和自定义内核。

提醒： Flash Attention 3（FA3） 不适合 gpt-oss 训练 ，因为它目前不支持 attention sink 的反向传播，这会导致 错误的训练损失。如果你 没有 使用 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）（这会减少 VRAM 占用并提高速度） 但这会导致错误的训练损失。参见 FA3 仓库中的问题 1797 。不过你必须禁用 FA3，因为它会阻止长上下文训练，因为 FA3 使用 O(N) 内存占用，而朴素 attention 会膨胀为 O(N^2) 占用。因此，为了使 attention sink 可微分，我们实现了 Unsloth Flex Attention.

我们通过对 BitsandBytes 4-bit 进行了基准测试来评估 gpt-oss RL 推理，并且还对 BF16 做了单独测试。Unsloth 的 4-bit 推理速度约快 4 倍，而 BF16 也更高效，尤其是在 VRAM 使用方面。

Unsloth 的 gpt-oss RL 最棒的一点是，它可以在任何 GPU 上工作，甚至包括那些不支持 BF16 的 GPU。我们免费的 gpt-oss-20b Colab 笔记本使用的是较老的 15GB T4 GPU，因此这些推理示例运行良好！

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

我们不得不修改 attention sink 的实现， 这里所述 以便生成能够与左侧填充一起工作。我们不得不获取 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 中一个棘手的问题。我们发现，不仅要在生成 token 时考虑 KV Cache 预填充，还要考虑批量生成中每个提示里的不同填充 token 数量，这会改变我们存储 block mask 的方式。下面可以看到这样的一个例子：

普通因果掩码：

   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   <-- 最后一行 query（对解码最重要）

对于一般情况的推理（解码）

    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，因为这是当前设置中的第一个 query）

对于生成（解码阶段），我们通常只关心 attention 矩阵的最后一行，因为只有一个 query token 在关注所有之前的 key token。如果我们天真地应用因果掩码（q_idx ≥ k_idx），这会失败，因为我们的单个 query 的索引是 0，而有 n_k 个 key token。为了解决这个问题，我们需要在构建掩码时加入偏移量，以决定应该关注哪些 token。但朴素的方法很慢，因为偏移量每一步都会变化，迫使我们重新生成掩码和内核。我们通过缓存和编译优化解决了这个问题。

更困难的部分是批量生成。序列长度不同，因此填充会使掩码构建复杂化。Flex Attention 有很多 挑战 ，而且动态掩码很棘手。更糟的是，如果不编译，它会回退到 eager attention，这很慢且内存占用高（序列长度上是二次方而非线性）。

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

你需要用 _compile=True 来调用它。我们本质上是将你的 block mask 映射到一个完整的 Q_LEN x KV_LEN 矩阵上，以生成 block mask。如果不编译，我们就需要把这个完整矩阵物化出来，而这在长序列上可能导致 OOM。

此外，你还需要运行 flex_attention = torch.compile(flex_attention)。如果不编译，flex 会回退到一种非融合的 eager 实现，这在调试时很有用，但速度要慢得多，而且会物化完整的分数矩阵。

最终，这个掩码必须通过 KV Cache 动态处理 prefill 与 decode、按序列处理 batch 和 padding tokens，保持 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 sink 的反向传播

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

⚠️ 我们能对抗奖励黑客吗？

RL 的最终目标是最大化某个奖励（例如速度、收入、某个指标）。但 RL 可以 作弊。 当 RL 算法学会了一种技巧，或者利用某些东西来提高奖励，但实际上最后并没有完成任务时，这就叫做“奖励黑客".

这也是模型学会修改单元测试以通过编程挑战的原因，而这些是现实部署的关键障碍。还有一些其他很好的例子来自 维基百科.

在我们的 免费 gpt-oss RL 笔记本中 我们探索了如何在代码生成场景中对抗奖励黑客，并展示了针对常见错误模式的切实解决方案。我们看到模型会修改计时函数、外包给其他库、缓存结果，甚至直接作弊。对抗之后，结果是我们的模型生成了真正优化过的矩阵乘法内核，而不是巧妙的作弊手段。

🏆奖励黑客

RL 期间奖励黑客的一些常见例子包括：

懒惰

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

缓存与作弊

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

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

作弊

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

教程：如何使用 RL 训练 gpt-oss

LLM 通常很难处理涉及复杂环境的任务。不过，通过应用 强化学习 （RL）并设计一个自定义的 奖励函数，这些挑战可以被克服。

RL 可适用于自动生成内核或策略创建等任务。本教程展示如何训练 gpt-oss 使用 GRPO 和 Unsloth，自主击败 2048。

我们的笔记本已经包含了逐步指南，教你如何完成整个流程。

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

你将构建什么：

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

• 创建一个模型可以交互的最小化 2048 环境

• 定义 奖励函数 ，用于：

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

  2. 防止奖励黑客（禁止外部导入），以及

  3. 奖励实际的游戏成功

• 运行推理并导出模型（MXFP4 4 位或合并 FP16）

硬件： 2048 示例可以在免费的 Colab T4 上运行，但训练会很慢。A100/H100 会快得多。4-bit 加载 + LoRA 让你可以把一个 20B 模型装入适中的 VRAM