長いコンテキストでの gpt-oss トレーニング

OpenAI gpt-oss トレーニング向けの Unsloth Flex Attention サポートを紹介できることを嬉しく思います。これにより 8倍以上の文脈長, VRAM 使用量を50%以上削減 および >1.5× 高速なトレーニング（精度低下なし） Flash Attention 3（FA3）を使用するものを含むすべての実装と比較して。Unsloth Flex Attention により、次が可能になります： 60K のコンテキスト長 BF16 LoRA の場合、80GB VRAM の H100 GPU 上で動作します。さらに：

• あなたは 今すぐエクスポート/保存 あなたの QLoRA でファインチューニングした gpt-oss モデルを llama.cpp、vLLM、Ollama、または Hugging Face に

• 我々は gpt-oss トレーニングの 損失が無限大に発散する問題を修正しました float16 GPU（T4 Colab のような）で

• 我々は gpt-oss 実装の Unsloth には無関係な問題、最も顕著なのは swiglu_limit = 7.0 が transformers の MXFP4 推論中に適切に適用されることです

🦥Unsloth Flex Attention サポートのご紹介

Unsloth の Flex Attention サポートにより、単一の 80GB VRAM H100 で QLoRA なら最大 81K コンテキスト長、BF16 LoRA なら 60K コンテキストを扱うことができます！これらの利点は次に適用されます： 両方 gpt-oss-20b と gpt-oss-120b！ 使用するコンテキスト長が長いほど、Unsloth Flex Attention から得られる利得は大きくなります：

比較すると、Unsloth 以外のすべての実装は 80GB GPU 上で最大 9K コンテキスト長に制限され、FA3 を使っても 15K コンテキストにしか到達できません。しかし、 FA3 は attention sinks の逆伝播をサポートしていないため gpt-oss トレーニングには不適切です。したがって、以前に gpt-oss トレーニングで FA3 を使用していた場合、私たちは 今は使用しないことを お勧めします。したがって、Unsloth を使わない場合に 80GB VRAM で得られる最大コンテキスト長は約 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 トークン （現在のトークンを含む）だけに注目するため、計算量は大幅に削減されます。しかし、これによりスライディングウィンドウが小さいために長いコンテキストの検索や推論が無意味になるという問題も生じます。多くの研究所はこれを、スライディングウィンドウを 2048 または 4096 トークンに拡張することで解決しています。

OpenAI は次を活用しました： Attention Sinks（アテンションシンク） Efficient Streaming Language Models with Attention Sinks からの知見 🔓 Tiled MLP：500K+の解放 この論文は、小さなスライディングウィンドウを使うことができるが、最初のトークンにグローバルアテンションを追加する必要があることを示しています！論文は以下のような良い図を提供しています：

論文は次のことを発見しました： アテンション機構は最初の数トークン（1〜4）に多くの重みを割り当てるように見える、そしてスライディングウィンドウ操作中にそれらを削除してしまうと、これらの「重要な」最初の数トークンが消え、長いコンテキストの検索が悪化します。

もし対数パープレキシティ（値が高いほど悪い）をプロットし、事前学習モデルの設定されたコンテキスト長を超えて長いコンテキスト推論を行うと、パープレキシティが急上昇する（良くない）ことが分かります。しかし赤い線（Attention Sinks を使用）は低いままで、これは非常に良いことです！

論文はまた、 Attention Is Off By One メソッド が部分的に機能することを示しています。ただし、より低いパープレキシティを得るにはいくつかの追加のシンクトークンを加える必要があります。 論文は、学習可能な単一のシンクトークンを追加することが非常にうまくいくことを示しています！ そしてそれが OpenAI が GPT-OSS に対して行ったことです！

📐Unsloth の Flex Attention 実装は

Flex Attention https://pytorch.org/blog/flexattention/ 非常に強力で、実践者にアテンション機構をカスタマイズするための 2 つのルートを提供します： スコア修飾子（f） と マスキング関数（M）.

その スコア修飾子（f） これにより、ソフトマックス操作の前にアテンションのロジットを編集でき、 マスキング関数（M） （後者）は不要な操作をスキップすることを可能にします（例：スライディングウィンドウアテンションは最後の 128 トークンのみを見る）。

ポイントは、Flex Attention が任意のスコア修飾子とマスキング関数で高速な自動生成 Triton カーネルを提供することです！

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

これは Flex Attention を使ってアテンションシンクを実装できることを意味します！単一のアテンションシンクの実装は、 OpenAI の元の GPT-OSS リポジトリ と HuggingFace の transformers 実装の両方で提供されています。

combined_logits = torch.cat([attn_weights, sinks], dim=-1)
probs = F.softmax(combined_logits, dim=-1)
scores = probs[..., :-1]

上記はシンクを Q @ K.T の末尾に連結し、ソフトマックスを行い、最後の列（シンクトークン）を取り除くことを示しています。

Flex Attention の Github リポジトリからのいくつかの可視化ユーティリティを使うことで、これを可視化できます。シーケンス長が 16、スライディングウィンドウが 5 と仮定します。左は最後のシンク列（デフォルト実装）、右はシンク位置をインデックス 0 に移動した場合（我々の実装）です。 Flex Attention の Github リポジトリを使うと可視化できます。

末尾にあるシンク位置（デフォルト）

シンク位置をインデックス 0 に移動

興味深い発見：公式の Flex Attention スライディングウィンドウ実装はウィンドウサイズを「最後のトークン数として」考慮します プラスワン （現在のトークンを含むため）。HuggingFace と GPT OSS の実装は厳密に最後の N トークンのみを見ます。つまり下記は 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 トークン）

HuggingFace、GPT-OSS（3+0 トークン）

我々はまた、最後の N トークンに注目するのか N+1 トークンに注目するのかを確認するために OpenAI の公式 GPT-OSS 実装を通じて確認しました： 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 トークン（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

また、シンクトークンのインデックスを最初に移動したため、正しくインデックス付けするには 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) を加えます
    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 を使用しない）とトレーニング損失が一致することを我々のグラフで検証しました：

📜 アテンションシンクの数学的導出

K と V をパディングせずにアテンションシンクを計算する別の方法があります。まずソフトマックス操作が行うことに注意し、ここではスカラーとしてシンク付きの 2 番目のバージョンを求めます：\

$$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)$$

これでアテンションのシンク版を容易に導出できます。ただしこの方法はゼロパディング方式よりも多少誤差が大きいため、我々は依然として元のバージョンをデフォルトとしています。

Unslothは現在gpt-ossのRLをサポートしています！我々は2つのノートブックを作成しました。詳細はgpt-oss RL用の特定のブログをご覧ください：💾

新機能：gpt-ossトレーニング後のGGUF、vLLMへの保存 llama.cpp, vLLM、または QLoRAでgpt-ossをファインチューニングした後、直接モデルを保存、エクスポート、またはマージして HF

以前は、QLoRA でファインチューニングした gpt-oss モデルは Unsloth 内でしか実行できませんでした。我々は次をマージする機能を導入することでその制限を取り除きました： を使用して ネイティブフォーマット を使用して で保存するか、ネイティブの および ファインチューニング済みモデルを保存するには、bf16 形式でエクスポートすることもできます、 （gpt-oss のような）ベースモデルをマージすることで、 次のコマンドを使って bf16 フォーマットでファインチューニング済みモデルをエクスポートできるようになりました： オンデマンドの MXFP4 の逆量子化 .

その を使用して MXFP4 Safetensors 形式でsave_method="mxfp4" GGUF を使って保存できます。

形式への変換をはるかに高速にします。 を使用して 新機能：QLoRA でファインチューニングしたモデルを GGUF に保存またはマージして他のフレームワーク（例：Hugging Face、llama.cpp の GGUF）で使用できるようになりました。

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

✨ モデルをマージして直接 Hugging Face ハブへプッシュすることを好む場合：

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
   model.push_to_hub_merged(repo_name, tokenizer=tokenizer, token= hf_token, save_method="mxfp4")

2. Llama.cpp への保存 を使用して cp llama.cpp/build/bin/llama-* llama.cp

   マージ済みモデルを変換する：

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

   量子化モデルで推論を実行する：
       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 \

✨ 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": "What is the capital of France?"}],
   }
   
   response = requests.post(url, json=data)
   print_highlight(response.json())

♦️gpt-oss を直接ファインチューニングする

我々はまた、ネイティブな MXFP4 量子化フォーマットをロードできるようにするパッチを実装することで、gpt-oss モデルの直接ファインチューニングのサポートを追加しました。これにより 'openai/gpt-oss' モデルを 24GB 未満の VRAM でロードし、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_...", # ゲート付きモデルを使う場合は指定してください
)

Peft レイヤーを追加するには次を使用します FastLanguageModel.get_peft_model そして Peft モデル上で SFT ファインチューニングを実行します。

🐛 gpt-oss のバグ修正

我々は 最近 Hugging Face と協力して OpenAI のカーネルを使用し、MXFP4 推論中に次が正しく適用されるようにすることで推論問題を解決しました。 swiglu_limit = 7.0 が正しく適用されること。

ユーザーからのフィードバックに基づき、長時間の QLoRA トレーニング（60 ステップを超える）を行うと 損失が発散して最終的にエラーになることがあると判明しました。これは BF16 をサポートせず代わりに F16 にフォールバックするデバイス（例：T4 GPU）でのみ発生しました。重要なことに、これは A100 や H100 GPU 上の QLoRA トレーニングや f16 GPU 上の LoRA トレーニングには影響しませんでした。

徹底的な調査の後、我々は F16 に限定された GPU を含むすべての GPU 構成でトレーニング損失の挙動を揃えました。もし以前このために問題を経験していたなら、新しい更新された gpt-oss ノートブックの利用をお勧めします！

float16 のトレーニング損失曲線を bfloat16 マシン（青線）と同等にするために多くの実験を行いました。その結果、次のことが分かりました：

1. 純粋な float16 は 50 ステップで無限大に発散する

2. 我々は MoE のダウンプロジェクションに非常に大きな外れ値があることを発見しました

3. アクティベーションは bfloat16 または float32 で保存する必要がある

下図は GPT OSS 20B の絶対値のアクティベーションを示しており、いくつかが大きくスパイクしています—float16 の最大範囲は 65504 なので float16 マシンではオーバーフローします。

我々はこれを Unsloth で修正したため、すべての float16 トレーニングはそのまま動作します！

🔢 Sink Attention（シンクアテンション）の実装

OpenAI のシンクトークン実装は ここで提供されています。以下に提供します：

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
    if attention_mask is not 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]  # ここでシンクを落とします
    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