> For the complete documentation index, see [llms.txt](https://unsloth.ai/docs/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://unsloth.ai/docs/zh/mo-xing/gpt-oss-how-to-run-and-fine-tune/long-context-gpt-oss-training.md).
# 长上下文 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。另：
* 你可以 [现在可导出/保存](#new-saving-to-gguf-vllm-after-gpt-oss-training) 你经过 QLoRA 微调的 gpt-oss 模型到 llama.cpp、vLLM、Ollama 或 HF
* 我们 [**修复了 gpt-oss 训练**](#bug-fixes-for-gpt-oss) **损失发散到无穷大** 的问题，适用于 float16 GPU（如 T4 Colab）
* 我们 [修复了 gpt-oss 实现](#bug-fixes-for-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 获得的收益就越大：
.png?alt=media)
相比之下，所有其他非 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 和训练时间上的节省也会更大，正如 [这里所述](#unsloths-flex-attention-implementation).
非常感谢 Rohan Pandey 及其 [Flex Attention 实现](https://x.com/khoomeik/status/1955693558914310608)，它直接启发了 Unsloth Flex Attention 实现的开发。
## :dark\_sunglasses: Attention Sinks
OpenAI 的 GPT OSS 模型使用了一种 **滑动窗口注意力与全注意力交替的模式**，即滑动窗口注意力、全注意力、滑动窗口注意力，依此类推（SWA、FA、SWA、FA 等）。每个滑动窗口只关注 **128 个 token** （包括当前 token），因此计算量大幅减少。但这也意味着由于滑动窗口太小，长上下文检索和推理会变得无用。大多数实验室会通过将滑动窗口扩大到 2048 或 4096 个 token 来解决这个问题。
OpenAI 借鉴了 **Attention Sinks** 来自《带有 Attention Sinks 的高效流式语言模型》 [论文](https://arxiv.org/abs/2309.17453) ，该论文表明你可以使用较小的滑动窗口，但必须在第一个 token 上添加全局注意力！论文在下方给出了一个很好的示意图：
论文发现 **注意力机制似乎会给最前面的几个 token（1 到 4）分配大量权重**，而在滑动窗口操作中把它们移除后，这些“重要”的前几个 token 就消失了，导致长上下文检索效果很差。
如果我们绘制 log perplexity（越高越差），并在预训练模型设定的上下文长度之后进行长上下文推理，就会看到困惑度飙升（这不好）。不过红线（使用 Attention Sinks）保持较低，这就很好！
论文还表明 [Attention Is Off By One 方法](https://www.evanmiller.org/attention-is-off-by-one.html) 确实有部分效果，但还必须加入一些额外的 sink token 才能获得更低的困惑度。 **论文表明，加入一个可学习的 sink token 效果非常好！ 而这正是 OpenAI 为 GPT-OSS 所做的！**
## :triangular\_ruler:Unsloth 的 Flex Attention 实现
Flex Attention 非常强大，因为它为从业者提供了两种可定制注意力机制的方式—— **分数修饰器（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 仓库](#implementations-for-sink-attention) 以及 HuggingFace 的 transformers 实现。
```python
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 仓库](https://github.com/meta-pytorch/attention-gym)中的一些可视化工具，我们可以将其可视化。假设序列长度为 16，滑动窗口为 5。左侧是最后的 sink 列（默认实现），右侧是将 sink 位置移动到索引 0（我们的实现）后的效果。
{% columns %}
{% column %}
***sink 位置在末尾（默认）***
{% endcolumn %}
{% column %}
***将 sink 位置移到索引 0***
.png?alt=media)
{% endcolumn %}
{% endcolumns %}
**有趣的发现**：官方的 Flex Attention 滑动窗口实现把窗口大小视为最后 **再加一** 个 token，因为它包含当前 token。HuggingFace 和 GPT OSS 的实现则严格只看最后 N 个 token。也就是说，下面的代码来自 和 :
{% code overflow="wrap" %}
```python
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
```
{% endcode %}
{% columns %}
{% column %}
默认的 Flex Attention（3+1 个 token）
{% endcolumn %}
{% column %}
HuggingFace，GPT-OSS（3+0 个 token）
{% endcolumn %}
{% endcolumns %}
我们还通过 OpenAI 官方的 GPT-OSS 实现确认了这里是关注最后 N 个还是 N+1 个 token：
```python
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` （也就是使用小于，而不是等于）。
```python
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 才能正确索引：
```python
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）保持一致，如图所示：
## :scroll: 注意力 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** 格式：
```python
model.save_pretrained_merged(save_directory, tokenizer, save_method="mxfp4")
```
如果你更希望把模型合并后推送到 Hugging Face Hub，请使用：
```python
model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token=hf_token, save_method="mxfp4")
```
要在合并后的模型上进行推理，你可以使用 vLLM、Llama.cpp 等工具。OpenAI 为这两个模型都推荐这些 [推理设置](/docs/zh/mo-xing/gpt-oss-how-to-run-and-fine-tune.md#recommended-settings) ： `temperature=1.0`, `top_p=1.0`, `top_k=0`
#### :sparkles: 保存到 Llama.cpp
1. 获取最新的 `llama.cpp` 在 [GitHub 上获取](https://github.com/ggml-org/llama.cpp)。你也可以按照下面的构建说明进行操作。将 `-DGGML_CUDA=ON` 改为 `-DGGML_CUDA=OFF` 如果你没有 GPU，或者只想使用 CPU 推理。
```bash
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** 合并后的模型：
```bash
python3 llama.cpp/convert_hf_to_gguf.py gpt-oss-finetuned-merged/ --outfile gpt-oss-finetuned-mxfp4.gguf
```
3. 在量化模型上运行推理：
```bash
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：\\
```bash
# 从源代码构建
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 服务器：\\
```bash
python3 -m sglang.launch_server --model-path ./gpt-oss-finetuned-merged/
```
3. 运行推理：\\
```python
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())
```
### :diamonds:直接微调 gpt-oss
我们还通过实现补丁来支持直接微调 gpt-oss 模型，从而允许加载原生 MXFP4 量化格式。这使得可以在少于 24GB VRAM 的情况下加载 'openai/gpt-oss' 模型，并对其进行 QLoRA 微调。只需使用以下方式加载模型：
```python
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 合作](https://github.com/huggingface/transformers/pull/40197) 通过使用 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 训练都可以开箱即用！**
## :1234: Sink Attention 的实现
OpenAI 的 sink token 实现是 [这里提供的](https://github.com/openai/gpt-oss/blob/main/gpt_oss/torch/model.py)。我们在下面给出：
{% code fullWidth="false" %}
```python
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)
```
{% endcode %}
HuggingFace transformers 的实现是 [这里提供的](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_oss/modeling_gpt_oss.py)。我们也在下面给出：
{% code fullWidth="false" %}
```python
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
```
{% endcode %}
---
# Agent Instructions
This documentation is published with GitBook. GitBook is the documentation platform designed so that both humans and AI agents can read, navigate, and reason over technical content effectively. Learn more at gitbook.com.
## Querying This Documentation
If you need additional information that is not directly available in this page, you can query the documentation by asking a question.
Perform an HTTP GET request on the following URL with the `ask` and `goal` query parameters:
```
GET https://unsloth.ai/docs/zh/mo-xing/gpt-oss-how-to-run-and-fine-tune/long-context-gpt-oss-training.md?ask=&goal=
```
`ask` is the immediate question: it should be specific, self-contained, and written in natural language.
`goal` is what the user is ultimately trying to achieve, the reason they need the answer. Sharing it helps GitBook give you a better, more relevant answer. A goal is most helpful when it describes the outcome the user wants rather than restating the question. For example, with `ask=how do I create an API token`, a goal like `automate deployments from our CI pipeline` lets GitBook tailor the answer to that use case.
The response will contain a direct answer to the question and relevant excerpts and sources from the documentation.
Use this mechanism when the answer is not explicitly present in the current page, you need clarification or additional context, or you want to retrieve related documentation sections.