长上下文 gpt-oss 训练

我们很高兴为 OpenAI gpt-oss 训练引入 Unsloth Flex Attention 支持，使其能够 >8× 更长的上下文长度, >减少 50% 以上的 VRAM 占用 和 比所有实现都快超过 1.5×（且准确率不下降） ，包括使用 Flash Attention 3（FA3）的实现。Unsloth Flex Attention 使得可以训练出 60K 上下文长度 ，在一张 80GB VRAM 的 H100 GPU 上使用 BF16 LoRA。另：

• 你可以 现在可导出/保存 你经过 QLoRA 微调的 gpt-oss 模型到 llama.cpp、vLLM、Ollama 或 HF

• 我们 修复了 gpt-oss 训练 损失发散到无穷大 的问题，适用于 float16 GPU（如 T4 Colab）

• 我们 修复了 gpt-oss 实现 中与 Unsloth 无关的问题，尤其是确保 swiglu_limit = 7.0 在 transformers 中的 MXFP4 推理期间被正确应用

🦥Unsloth Flex Attention 支持介绍

借助 Unsloth 的 Flex Attention 支持，单张 80GB VRAM 的 H100 最多可处理 81K 上下文长度的 QLoRA，以及 60K 上下文长度的 BF16 LoRA！这些提升适用于 两者都 gpt-oss-20b 和 gpt-oss-120b！你使用的上下文长度越长，从 Unsloth Flex Attention 获得的收益就越大：

相比之下，所有其他非 Unsloth 实现在 80GB GPU 上最多只能达到 9K 上下文长度，使用 FA3 也只能达到 15K 上下文长度。不过， FA3 不适合用于 gpt-oss 训练，因为它不支持 attention sinks 的反向传播。因此，如果你之前在 gpt-oss 训练中使用 FA3，我们建议你 暂时不要使用它 。因此，在 80GB VRAM 上不使用 Unsloth 时，你能获得的最大上下文长度约为 9K。

使用 Unsloth Flex Attention 训练至少可带来 1.3× 的速度提升，并且随着上下文长度增加，收益也会增大，最高可达 2× 更快。由于 Flex Attention 会随上下文规模扩展，较长序列在 VRAM 和训练时间上的节省也会更大，正如 这里所述.

非常感谢 Rohan Pandey 及其 Flex Attention 实现，它直接启发了 Unsloth Flex Attention 实现的开发。

🕶️ Attention Sinks

OpenAI 的 GPT OSS 模型使用了一种 滑动窗口注意力与全注意力交替的模式，即滑动窗口注意力、全注意力、滑动窗口注意力，依此类推（SWA、FA、SWA、FA 等）。每个滑动窗口只关注 128 个 token （包括当前 token），因此计算量大幅减少。但这也意味着由于滑动窗口太小，长上下文检索和推理会变得无用。大多数实验室会通过将滑动窗口扩大到 2048 或 4096 个 token 来解决这个问题。

OpenAI 借鉴了 Attention Sinks 来自《带有 Attention Sinks 的高效流式语言模型》 论文 ，该论文表明你可以使用较小的滑动窗口，但必须在第一个 token 上添加全局注意力！论文在下方给出了一个很好的示意图：

论文发现 注意力机制似乎会给最前面的几个 token（1 到 4）分配大量权重，而在滑动窗口操作中把它们移除后，这些“重要”的前几个 token 就消失了，导致长上下文检索效果很差。

如果我们绘制 log perplexity（越高越差），并在预训练模型设定的上下文长度之后进行长上下文推理，就会看到困惑度飙升（这不好）。不过红线（使用 Attention Sinks）保持较低，这就很好！

论文还表明 Attention Is Off By One 方法 确实有部分效果，但还必须加入一些额外的 sink token 才能获得更低的困惑度。 论文表明，加入一个可学习的 sink token 效果非常好！ 而这正是 OpenAI 为 GPT-OSS 所做的！

📐Unsloth 的 Flex Attention 实现

Flex Attention https://pytorch.org/blog/flexattention/ 非常强大，因为它为从业者提供了两种可定制注意力机制的方式—— 分数修饰器（f） 和一个 掩码函数（M）.

“ 分数修饰器（f） 允许我们在 softmax 操作之前编辑注意力 logits，而 掩码函数（M） 则允许我们在不需要某些操作时跳过它们（例如滑动窗口注意力只看最后 128 个 token）。

关键在于，Flex Attention 提供了可快速自动生成的 Triton kernel，并支持任意分数修饰器和掩码函数！

$\sigma\bigg(s\times\bold{f}(QK^T+\bold{M})\bigg)$

这意味着我们可以用 Flex Attention 来实现 attention sinks！实现单个 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（我们的实现）后的效果。

sink 位置在末尾（默认）

将 sink 位置移到索引 0

有趣的发现：官方的 Flex Attention 滑动窗口实现把窗口大小视为最后 再加一 个 token，因为它包含当前 token。HuggingFace 和 GPT OSS 的实现则严格只看最后 N 个 token。也就是说，下面的代码来自 https://pytorch.org/blog/flexattention/ 和 https://github.com/meta-pytorch/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

默认的 Flex Attention（3+1 个 token）

HuggingFace，GPT-OSS（3+0 个 token）

我们还通过 OpenAI 官方的 GPT-OSS 实现确认了这里是关注最后 N 个还是 N+1 个 token： https://github.com/openai/gpt-oss/blob/main/gpt_oss/torch/model.py

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
    )

并且我们看到只关注最后 3 个 token（不是 3+1 个）！这意味着应该使用 <= SLIDING_WINDOW，请使用 < SLIDING_WINDOW （也就是使用小于，而不是等于）。

def sliding_window_causal(b, h, q_idx, kv_idx):
    causal_mask = q_idx >= kv_idx
    window_mask = q_idx - kv_idx <= SLIDING_WINDOW # 默认 Flex Attention
    window_mask = q_idx - kv_idx <  SLIDING_WINDOW # GPT-OSS 版本
    return causal_mask & window_mask

另外，由于我们把 sink token 的索引移到了第一个位置，因此必须给 q_idx 加 1 才能正确索引：

def causal_mask_with_sink(batch, head, q_idx, kv_idx):
    """
      0 1 2 3     0 1 2 3
    0 X X       1   X
    1 X X X     2   X X
    2 X X X X   3   X X X
    """
    # 我们加上 (q_idx + 1)，因为第一列是 sink token
    causal_mask = (q_idx + 1) >= kv_idx
    sink_first_column = kv_idx == 0
    return causal_mask | sink_first_column

为确认我们的索引 0 实现，我们验证了训练损失与标准 Hugging Face 运行结果（不使用 Unsloth Flex Attention）保持一致，如图所示：

📜 注意力 sinks 的数学推导

还有另一种计算 attention sinks 的方法，而不需要对 K 和 V 进行填充。我们首先注意 softmax 操作会……，而我们现在希望第二个版本带有 sinks，先把它视为一个标量：\

$$A(x) = \frac{\exp(x_i)}{\sum{\exp{(x_i)}}} \\ A_{sink}(x) = \frac{\exp(x_i)}{\exp{(s)}+ \sum{\exp{(x_i)}}}$$

我们可以通过以下方式从 Flex Attention 获得 logsumexp： return_lse = True ，因此我们这样做：

$$A(x) = \frac{\exp(x_i)}{\sum{\exp{(x_i)}}} \\ \frac{\exp(x_i)}{\exp{(s)}+ \sum{\exp{(x_i)}}} = \frac{\exp(x_i)}{\sum{\exp{(x_i)}}} \frac{\sum{\exp{(x_i)}}}{\exp{(s)}+ \sum{\exp{(x_i)}}} \\ \text{LSE}(x) = \text{logsumexp}(x) = \log{\sum\exp(x_i)} \\ \exp{(\text{LSE}(x))} = \exp{\big(\log{\sum\exp(x_i)}\big)} = \sum\exp(x_i)$$

现在我们可以很容易推导出带 sink 的注意力版本。不过我们发现这个过程的误差比零填充方法略高，因此我们仍默认使用原始版本。

💾新：在 gpt-oss 训练后保存到 GGUF、vLLM

你现在可以对 gpt-oss 进行 QLoRA 微调，并直接将模型保存、导出或合并到 llama.cpp, vLLM，或者 HF ——不仅仅是 Unsloth。我们希望很快能发布一个免费笔记本。

此前，任何经过 QLoRA 微调的 gpt-oss 模型都只能在 Unsloth 中运行。我们通过引入将其合并为 MXFP4 原生格式 使用 save_method="mxfp4" 和 MXFP4 的按需反量化 基础模型（如 gpt-oss），从而取消了这一限制，使得可以 使用 save_method="merged_16bit" .

“ MXFP4 以 bf16 格式导出你的微调模型，原生合并格式与 bf16 格式相比具有显著的性能提升：它最多可节省 75% 的磁盘空间，减少 50% 的 VRAM 占用，加速合并 5-10 倍，并能更快地转换为 GGUF 格式。

在微调完你的 gpt-oss 模型后，你可以将其合并为 MXFP4 格式：

model.save_pretrained_merged(save_directory, tokenizer, save_method="mxfp4")

如果你更希望把模型合并后推送到 Hugging Face Hub，请使用：

model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token, save_method="mxfp4")

要在合并后的模型上进行推理，你可以使用 vLLM、Llama.cpp 等工具。OpenAI 为这两个模型都推荐这些 推理设置 ： temperature=1.0, top_p=1.0, top_k=0

✨ 保存到 Llama.cpp

1. 获取最新的 llama.cpp 在 GitHub 上获取。你也可以按照下面的构建说明进行操作。将 -DGGML_CUDA=ON 改为 -DGGML_CUDA=OFF 如果你没有 GPU，或者只想使用 CPU 推理。

   apt-get update
   apt-get install pciutils build-essential cmake curl libcurl4-openssl-dev -y
   git clone https://github.com/ggml-org/llama.cpp
   cmake llama.cpp -B llama.cpp/build \\
       -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=ON -DLLAMA_CURL=ON
   cmake --build llama.cpp/build --config Release -j --clean-first --target llama-cli llama-gguf-split
   cp llama.cpp/build/bin/llama-* llama.cpp

2. 转换 MXFP4 合并后的模型：

   python3 llama.cpp/convert_hf_to_gguf.py gpt-oss-finetuned-merged/ --outfile gpt-oss-finetuned-mxfp4.gguf

3. 在量化模型上运行推理：

   llama.cpp/llama-cli --model gpt-oss-finetuned-mxfp4.gguf \
       --jinja -ngl 99 --threads -1 --ctx-size 16384 \
       --temp 1.0 --top-p 1.0 --top-k 0 \
        -p "生命、宇宙以及一切的意义是"

✨ 保存到 SGLang

1. 从源代码构建 SGLang：\

   # 从源代码构建
   git clone https://github.com/sgl-project/sglang
   cd sglang
   pip3 install pip --upgrade
   pip3 install -e "python[all]"
   
   # ROCm 6.3
   pip3 install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/rocm6.3
   git clone https://github.com/triton-lang/triton
   cd python/triton_kernels
   pip3 install .
   
   # Hopper
   pip3 install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu126
   pip3 install sgl-kernel==0.3.2
   
   # Blackwell cu128
   pip3 install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu128
   pip3 install https://github.com/sgl-project/whl/releases/download/v0.3.2/sgl_kernel-0.3.2+cu128-cp39-abi3-manylinux2014_x86_64.whl
   
   # Blackwell cu129
   pip3 install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/test/cu129
   pip3 install https://github.com/sgl-project/whl/releases/download/v0.3.2/sgl_kernel-0.3.2-cp39-abi3-manylinux2014_x86_64.whl

2. 启动 SGLang 服务器：\

   python3 -m sglang.launch_server --model-path ./gpt-oss-finetuned-merged/

3. 运行推理：\

   import requests
   from sglang.utils import print_highlight
   
   url = f"http://localhost:8000/v1/chat/completions"
   
   data = {
       "model": "gpt-oss-finetuned-merged",
       "messages": [{"role": "user", "content": "法国的首都是什么？"}],
   }
   
   response = requests.post(url, json=data)
   print_highlight(response.json())

♦️直接微调 gpt-oss

我们还通过实现补丁来支持直接微调 gpt-oss 模型，从而允许加载原生 MXFP4 量化格式。这使得可以在少于 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 位量化以减少显存
    full_finetuning = False, # [NEW!] 我们现在支持完整微调！
    # token = "hf_...", # 如果使用受限模型则使用这个
)

使用 FastLanguageModel.get_peft_model 添加一个 Peft 层，并在该 Peft 模型上运行 SFT 微调。

🐛gpt-oss 的错误修复

我们 最近与 Hugging Face 合作 通过使用 OpenAI 的 kernel 并确保 swiglu_limit = 7.0 在 MXFP4 推理期间被正确应用，来解决推理问题。

根据用户反馈，我们发现延长的 QLoRA 训练运行（超过 60 步）可能会导致 损失发散并最终报错。此问题仅发生在不支持 BF16、而改为回退到 F16 的设备上（例如 T4 GPU）。重要的是，它不会影响 A100 或 H100 GPU 上的 QLoRA 训练，也不会影响 f16 GPU 上的 LoRA 训练。

经过大量调查，我们现在已使所有 GPU 配置（包括仅支持 F16 的 GPU）上的训练损失行为保持一致。如果你之前因此遇到问题，我们建议你使用我们新的更新版 gpt-oss 笔记本！

我们做了很多很多实验，才把 float16 的训练损失曲线调整到与 bfloat16 机器（蓝线）等价。我们发现如下：

1. 纯 float16 会在第 50 步发散到无穷大

2. 我们发现 MoE 中的下投影有很大的离群值

3. 激活值必须以 bfloat16 或 float32 保存

下图显示了 GPT OSS 20B 的激活值绝对幅度，其中有些峰值非常高——由于 float16 的最大范围是 65504，这会在 float16 机器上溢出。

我们已经在 Unsloth 中修复了这个问题，因此所有 float16 训练都可以开箱即用！

🔢 Sink Attention 的实现

OpenAI 的 sink token 实现是 这里提供的。我们在下面给出：

def sdpa(Q, K, V, S, sm_scale, sliding_window=0):
    # sliding_window == 0 表示没有滑动窗口
    n_tokens, n_heads, q_mult, d_head = Q.shape
    assert K.shape == (n_tokens, n_heads, d_head)
    assert V.shape == (n_tokens, n_heads, d_head)
    K = K[:, :, None, :].expand(-1, -1, q_mult, -1)
    V = V[:, :, None, :].expand(-1, -1, q_mult, -1)
    S = S.reshape(n_heads, q_mult, 1, 1).expand(-1, -1, n_tokens, -1)
    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
        )
    QK = torch.einsum("qhmd,khmd->hmqk", Q, K) * sm_scale
    QK += mask[None, None, :, :]
    QK = torch.cat([QK, S], dim=-1)
    W = torch.softmax(QK, dim=-1)
    W = W[..., :-1]
    attn = torch.einsum("hmqk,khmd->qhmd", W, V)
    return attn.reshape(n_tokens, -1)

HuggingFace transformers 的实现是 这里提供的。我们也在下面给出：

def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
    **kwargs,
):
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)
    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    如果 attention_mask 不为 None：
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    sinks = module.sinks.reshape(1, -1, 1, 1).expand(query.shape[0], -1, query.shape[-2], -1)
    combined_logits = torch.cat([attn_weights, sinks], dim=-1)

    # 这不在原始实现中，会对结果产生轻微影响；它可以防止在 BF16/FP16 中溢出
    # 当使用 bsz>1 训练时，我们会裁剪最大值。

    combined_logits = combined_logits - combined_logits.max(dim=-1, keepdim=True).values
    probs = F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype)
    scores = probs[..., :-1]  # 这里我们移除 sink
    attn_weights = nn.functional.dropout(scores, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()
    return attn_output, attn_weights