第41章 Unsloth 支持 gpt-oss 长上下文微调：Flex Attention 提升 8 倍效能

Unsloth 长上下文 gpt-oss 微调 2025年8月28日 · 作者：Daniel 和 Michael 我们很高兴宣布，Unsloth Flex Attention 现已支持 OpenAI gpt-oss 训练。相较于包括使用 Flash Attention 3 (FA3) 在内的所有实现方案，它能让上下文长度延长 8 倍以上，显存占用减少 50% 以上，训练速度提升 1.5 倍以上。借助 Unsloth Flex Attention，只需 80GB 显存即可通过 BF16 LoRA 训练 60K 长度的上下文。 其他更新如下： * 现在可将经过 QLoRA 微调的 gpt-oss 模型导出并保存为 llama.cpp、vLLM、Ollama 或 Hugging Face 格式。 * 修复了在 float16 GPU（如 T4 Colab）上 gpt-oss 训练损失变为无穷大的问题。 * 修复了 gpt-oss 实现中与 Unsloth 无关的问题，主要确保在 transformers 中的 MXFP4 推理期间正确应用 $swiglu\_limit = 7.0$。 gpt-oss-20b 在 1 张 H100 80GB 显卡上： * 上下文长度延长 8 倍以上 * 速度提升 1.7 倍以上 * 显存占用减少 50% 以上 🦥 介绍 Unsloth Flex Attention 使用 Unsloth Flex Attention，单张 80GB 显存的 H100 显卡即可支持 QLoRA 下高达 81K 的上下文长度，以及 BF16 LoRA 下 60K 的上下文长度。使用的上下文长度越长，从 Unsloth Flex Attention 获得的收益越大： 相比之下，所有非 Unsloth 的实现方案在 80GB 显存下最大上下文长度仅为 9K，使用 FA3 时也仅能达到 15K。但 FA3 不适合 gpt-oss 训练，因为它缺乏对 attention sinks 反向传播的支持。因此，如果您之前使用 FA3 训练 gpt-oss，我们建议暂时停止使用。这意味着在不使用 Unsloth 的情况下，80GB 显存能支持的最大上下文长度约为 9K。 使用 Unsloth Flex Attention 进行训练至少能带来 1.3 倍的速度提升，且随着上下文长度增加，收益进一步增长，最高可达 2 倍。由于 Flex Attention 具有可扩展性，更长的序列能带来显存和训练时间上的更大节省。 特别感谢 Rohan Pandey 的 Flex Attention 实现，它直接启发了 Unsloth Flex Attention 的开发。 测试配置：Alpaca 数据集，批大小 2，梯度累积步数 4，秩 rank = 32，并在所有线性层（q, k, v, o, gate, up, down）上应用 LoRA。 📐 Unsloth 的 Flex Attention 实现 Flex Attention 功能极其强大，它为实践者提供了两条自定义注意力机制的路径——分数修饰器 ($f$) 和掩码函数 ($M$)。 分数修改器（f）允许我们在 softmax 之前编辑 attention logits，而掩码函数（M）让我们可以跳过不必要的计算（例如滑动窗口注意力只关注最后 128 个 token）。
妙处在于 Flex Attention 能通过自动生成的快速 Triton kernel，支持任意的分数修改器和掩码函数！

这意味着我们可以用 Flex Attention 来实现 attention sink！单个 attention sink 的实现可以在 OpenAI 的原始 gpt-oss 仓库和 HuggingFace transformers 的实现中找到。combined_logits = torch.cat([attn_weights, sinks], dim=-1)
probs = F.softmax(combined_logits, dim=-1)
scores = probs[..., :-1]上面的代码展示了我们把 sink 拼接在 Q @ K.T 的末尾，做 softmax，然后删掉作为 sink token 的最后一列。

借助 Flex Attention GitHub 仓库里的一些可视化工具，我们可以把这个过程可视化。假设序列长度为 16，滑动窗口大小为 5。左图是最后一列的 sink（默认实现），右图是把 sink 移到索引 0 的位置（我们的实现）。一个有趣的发现：Flex Attention 官方的滑动窗口实现把窗口大小当作"最近 token 数再加一"，因为它把当前 token 也算进去了。而 HuggingFace 和 GPT OSS 的实现则严格只看最近 N 个 token。也就是说，下面这段代码来自 Flex Attention 和 Attention Gym：def sliding_window_causal(b, h, q_idx, kv_idx):
causal_mask = q_idx >= kv_idx
window_mask = q_idx - kv_idx <= SLIDING_WINDOW
return causal_mask & window_mask我们还通过 OpenAI 官方的 GPT-OSS 实现确认了到底是关注最近 N 个还是 N+1 个 token：mask = torch.triu(Q.new_full((n_tokens, n_tokens), -float("inf")), diagonal=1)
if sliding_window > 0:
mask += torch.tril(
mask.new_full((n_tokens, n_tokens), -float("inf")), diagonal=-sliding_window
)🕶️ Attention SinksOpenAI 的 GPT OSS 模型采用滑动窗口注意力和全注意力交替的模式（SWA、FA、SWA、FA 等）。每个滑动窗口只关注 128 个 token（包括当前 token），因此计算量大幅减少。但这也意味着由于滑动窗口太小，长上下文的检索和推理基本失效。大多数实验室的做法是把滑动窗口扩大到 2048 或 4096 个 token 来解决这个问题。
OpenAI 借鉴了《Efficient Streaming Language Models with Attention Sinks》论文中提出的 Attention Sinks（注意力池）技术。该研究表明，可以使用较小的滑动窗口，但必须对首个 Token 保持全局注意力！论文提供了如下图示：

研究发现，注意力机制似乎会赋予前几个 Token（第 1 到第 4 个）较高的权重。而在滑动窗口操作中，由于我们只能看到最近 N 个（例如 128 个）Token，如果移除窗口，这些“重要”的前几个 Token 就会消失。

如果我们绘制对数困惑度（数值越高表示效果越差），并在预训练模型的上下文长度之后进行长上下文推理，会发现困惑度急剧上升（效果不佳）。然而，红线（Attention Sinks）保持在较低水平，效果非常好！
论文还展示了 Attention Is Off By One 方法在一定程度上有效，但需要额外添加几个 Sink Token 才能降低困惑度。研究表明，添加一个可学习的单个 Sink Token 效果显著！OpenAI 正是在 GPT-OSS 中采用了这种做法！💾新功能：gpt-oss 训练后保存为 GGUF 及 vLLM 格式
现在，您可以对 gpt-oss 进行 QLoRA 微调，并直接保存、导出或合并模型至 llama.cpp、vLLM 或 Hugging Face，不再局限于 Unsloth。我们即将发布免费 Notebook。

过去，任何经过 QLoRA 微调的 gpt-oss 模型都只能在 Unsloth 中运行。通过引入在 LoRA 合并过程中对 MXFP4 基础模型（如 gpt-oss）进行即时反量化，我们移除了这一限制。这使得以 bf16 格式导出微调模型成为可能。

模型微调完成后，你可以通过一条指令将其合并为 16-bit 格式：model.save_pretrained_merged(save_directory, tokenizer) 如果你更喜欢合并模型并直接推送到 hugging-face hub，可以使用：model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token) ✨ 直接微调 gpt-oss 我们还实现了加载原生 MXFP4 量化格式的补丁，支持直接微调 gpt-oss 模型。这意味着你可以使用少于 24GB 的 VRAM 加载 'openai/gpt-oss' 模型，并对其进行 QLoRA 微调。只需按以下方式加载模型： model, tokenizer = FastLanguageModel.from_pretrained(
#model_name = "unsloth/gpt-oss-20b-BF16",
model_name = "unsloth/gpt-oss-20b",
dtype = dtype, # None 为自动检测
max_seq_length = max_seq_length, # 长上下文可选任意值！
load_in_4bit = True, # 4-bit 量化以减少内存占用
full_finetuning = False, # [新特性!] 现在支持全量微调！
# token = "hf_...", # 若使用受限模型需填写 token
) 🐛 gpt-oss 修复项 我们最近与 Hugging Face 合作，通过使用 OpenAI 的 kernel 并确保在 MXFP4 推理期间正确应用 swiglu_limit = 7.0，解决了推理问题。

根据用户反馈，我们发现延长 QLoRA 训练（超过 60 步）会导致 loss 发散并最终报错。此问题仅出现在不支持 BF16 而回退至 F16 的设备上（例如 T4 GPU）。重要的是，这并未影响 A100 或 H100 GPU 上的 QLoRA 训练，也未影响 f16 GPU 上的 LoRA 训练。

经过深入调查，我们已统一了所有 GPU 配置下的训练 loss 行为，包括仅限 F16 的 GPU。如果你之前因此遇到问题，我们建议使用更新后的 gpt-oss notebook！ 为了将 float16 的训练 loss 曲线调整为与 bfloat16 机器（蓝线）等价，我们进行了大量实验。发现如下：
纯 float16 在第 50 步会变为无穷大
我们发现 MoE 中的 down projections 存在巨大的离群值
Activations 必须保存为 bfloat16 或 float32
下图展示了 GPT OSS 20B 的绝对幅度 activations，其中一些出现了尖峰 - 由于 float16 的最大范围仅为 65504，这在 float16 机器上会导致溢出。
我们已在 Unsloth 中修复了这个问题，现在所有 float16 训练都能开箱即用！📈gpt-oss-20b 基准测试 我们测试了 gpt-oss-20b，在所有线性层（Q、K、V、O、gate、up 和 down）上做 LoRA，rank = 32，batch size 为 1。所有序列都填充到指定的最大序列长度，以模拟长上下文微调的工作负载。gpt-oss-20b BF16 LoRA - 上下文长度 vs/ GPU 显存上下文长度Unsloth
(+ Flex A)官方 Cookbook + FA3官方 Cookbook102445.246.647.3204845.9449.751.3409647.0756.171.1819249.2768.7OOM16,38454OOMOOM32,76863.73OOMOOM61,23480OOMOOM💕 感谢！一如既往，非常感谢大家使用并分享 Unsloth，我们真的非常感激。🙏

如往常一样，欢迎加入我们的 Reddit 页面和 Discord 服务器寻求帮助，或者只是来表示支持！你也可以在 Twitter 上关注我们，并在 Substack 上订阅我们的newsletter。感谢阅读！Daniel & Michael Han 🦥
2025年8月28日立即微调 gpt-oss！免费开始使用 加入我们的 Discord