長文脈 gpt-oss 学習

OpenAIのgpt-oss学習向けに、Unsloth Flex Attention対応を導入できることをうれしくお知らせします。これにより >8倍長いコンテキスト長, >50%少ない VRAM 使用量 および 1.5倍超の高速学習（精度低下なし） Flash Attention 3（FA3）を使う実装を含むすべての実装と比べて実現します。Unsloth Flex Attentionにより、 6万トークンのコンテキスト長 を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でQLoRAなら最大8.1万トークン、BF16 LoRAなら6万トークンのコンテキスト長を扱えます。これらの改善は 両方の gpt-oss-20b と gpt-oss-120b! 使用するコンテキスト長が長いほど、Unsloth Flex Attentionによる恩恵も大きくなります。

比較すると、Unsloth以外の実装はすべて80GB GPUでは最大9Kのコンテキスト長にとどまり、FA3でも15Kまでしか到達できません。しかし、 FA3はattention sinkのバックワードパスをサポートしていないため、gpt-oss学習には不向きです。そのため、これまで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トークン （現在のトークンを含む）だけなので、計算量は大幅に削減されます。しかし、その一方でウィンドウが小さいため、長文コンテキストの検索や推論はほぼ使い物になりません。多くの研究室ではこれを、スライディングウィンドウを2048または4096トークンに拡張することで解決しています。

OpenAIは Attention Sinks Efficient Streaming Language Models with Attention Sinks 論文の手法を活用しました。この論文では、小さなスライディングウィンドウを使いつつ、最初のトークンにグローバル注意を加える必要があることを示しています。下図はその良い例です：

論文では、 注意機構が最初の数トークン（1〜4）に大きな重みを割り当てるように見えることが示されており、スライディングウィンドウ処理中にそれらを取り除くと、これらの「重要な」最初の数トークンが消えてしまい、長文コンテキストの検索性能が悪化します。

log perplexity（高いほど悪い）を描き、事前学習モデルの設定コンテキスト長を超えて長文コンテキスト推論を行うと、perplexityが急上昇するのが分かります（良くありません）。しかし、赤線（Attention Sinksを使用）は低いままで、これは非常に良いことです！

論文ではさらに、 Attention Is Off By One手法も部分的には有効だが、より低いperplexityを得るには追加のsinkトークンも必要だと示しています。 論文では、学習可能な単一のsinkトークンを追加するだけで驚くほどうまくいくことが示されています！ そしてOpenAIはGPT-OSSでこれを実施しました！

📐UnslothのFlex Attention実装

Flex Attention https://pytorch.org/blog/flexattention/ は非常に強力です。というのも、注意機構に対して実務者に2つのカスタマイズ手段を提供するからです。すなわち スコア修飾子（f） と マスキング関数（M）.

この スコア修飾子（f） で、softmax演算の前にattention logitsを編集できます。また マスキング関数（M） では、必要ない演算をスキップできます（例えばスライディングウィンドウ注意は最後の128トークンだけを見る）。

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

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

つまり、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トークンである最後の列を取り除いていることを示しています。

次のような可視化ユーティリティを使うと Flex AttentionのGitHubリポジトリのこれを可視化できます。系列長が16、スライディングウィンドウが5だったとします。左は最後のsink列（デフォルト実装）、右はsinkの位置をインデックス0に移動した場合（私たちの実装）です。

末尾にsinkを置く（デフォルト）

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

興味深い発見：公式のFlex Attentionのスライディングウィンドウ実装では、ウィンドウサイズを最後のトークン数＋1 として扱います。これは現在のトークンを含むためです。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トークン）

また、OpenAIの公式GPT-OSS実装を通じて、ここで最後のNトークンを見るのか、それともN+1トークンを見るのかも確認しました： 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

また、sinkトークンのインデックスを最初に移動したので、正しくインデックスするには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
    """
    # 最初の列はsinkトークンなので(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なし）と一致することを、次のグラフで検証しました：

📜 attention sinkの数学的導出

KとVをパディングせずにattention sinkを計算する別の方法があります。まずsoftmax演算が何をするかに注目し、今のところsink付きの第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)

これでattentionのsink版を簡単に導出できます。ただし、この手順はゼロパディング方式よりやや誤差が大きいことが分かったため、引き続き元の方式をデフォルトにしています。

💾新機能: gpt-oss 学習後の GGUF、vLLM への保存

これで gpt-oss を QLoRA でファインチューニングし、モデルを直接 llama.cpp, vLLM、または HF へ保存、エクスポート、またはマージできます。Unsloth だけではありません。無料ノートブックをできるだけ早く公開する予定です。

これまでは、QLoRAで微調整したgpt-ossモデルはUnsloth内でしか実行できませんでした。そこで、 MXFP4 ネイティブ形式 でマージする機能を追加し、 save_method="mxfp4" および MXFP4 のオンデマンド逆量子化 gpt-ossのようなベースモデルを使って 微調整済みモデルをbf16形式で書き出せるようにしました。以下を使います： save_method="merged_16bit" .

この MXFP4 ネイティブマージ形式は、 bf16 形式と比較して大きな性能向上を提供します。ディスク容量を最大 75% 節約し、VRAM 消費を 50% 削減し、マージを 5〜10 倍高速化し、さらに GGUF 形式への変換を大幅に高速化します。

gpt-oss モデルのファインチューニング後、それを MXFP4 形式に次のようにマージできます:

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

モデルをマージしてhugging-face hubにpushしたい場合は、次を使ってください：

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 への保存

最新の 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

次を変換します MXFP4 マージ済みモデル:

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 \
     -p "人生と宇宙の意味は"

✨ SGLangへの保存

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

SGLangサーバーを起動：\

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

推論を実行：\

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を直接ファインチューニング

ネイティブ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_...", # gated モデルを使う場合はこれを使用
)

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マシン（青線）と同等にするために、何度も何度も実験を行う必要がありました。次のことが分かりました：

純粋なfloat16は50ステップ目で無限大になります
MoEのダウンプロジェクションに非常に大きな外れ値があることが分かりました
活性化はbfloat16またはfloat32で保存する必要があります

以下はGPT OSS 20Bの絶対値活性化を示しており、非常に大きなスパイクがいくつかあります。float16の最大範囲は65504なので、float16マシンではオーバーフローします。

これをUnslothで修正したので、すべてのfloat16学習がそのまま動作します！

🔢 Sink Attentionの実装

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

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]  # ここで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

最終更新 2 か月前

役に立ちましたか？