Long-Context gpt-oss Training

Wir freuen uns, die Unterstützung von Unsloth Flex Attention für das OpenAI gpt-oss Training vorzustellen, die ermöglicht >8× längere Kontextlängen, >50% weniger VRAM-Verbrauch und >1,5× schnellere Ausbildung (ohne Genauigkeitsverlust) im Vergleich zu allen Implementierungen, einschließlich jener mit Flash Attention 3 (FA3). Unsloth Flex Attention macht es möglich, mit einer 60K Kontextlänge auf einer 80GB VRAM H100 GPU für BF16 LoRA zu trainieren. Außerdem:

• Du kannst jetzt exportieren/speichern Ihr QLoRA feinabgestimmtes gpt-oss Modell für llama.cpp, vLLM, Ollama oder HF

• Wir haben das gpt-oss Training wobei Verluste gegen unendlich gingen auf float16 GPUs (wie T4 Colab)

• Wir haben die gpt-oss Implementierung Fehler behoben, die für Unsloth irrelevant sind, am wichtigsten ist dabei die Gewährleistung, dass swiglu_limit = 7.0 während der MXFP4-Inferenz in transformers korrekt angewendet wird

🦥Einführung der Unsloth Flex Attention-Unterstützung

Mit Unsloths Flex Attention-Unterstützung kann eine einzelne 80GB VRAM H100 mit QLoRA bis zu 81K Kontextlänge und mit BF16 LoRA bis zu 60K Kontextlänge verarbeiten! Diese Verbesserungen gelten für BEIDE gpt-oss-20b und gpt-oss-120b! Je größer die verwendete Kontextlänge, desto mehr Vorteile erhalten Sie durch Unsloth Flex Attention:

Im Vergleich dazu erreichen alle anderen nicht-Unsloth-Implementierungen auf einer 80GB GPU maximal 9K Kontextlänge und können nur mit FA3 15K Kontext erreichen. Aber ist FA3 für gpt-oss Training ungeeignet, da es keinen Backward-Pass für Attention Sinks unterstützt. Wenn Sie also zuvor FA3 für gpt-oss Training verwendet haben, empfehlen wir Ihnen es momentan nicht zu benutzen Daher ist die maximale Kontextlänge, die Sie ohne Unsloth auf 80GB VRAM erreichen können, ~9K.

Training mit Unsloth Flex Attention liefert mindestens eine 1,3× Beschleunigung, mit wachsenden Vorteilen bei steigender Kontextlänge und erreicht bis zu 2× schnellere Ausführung. Da Flex Attention mit dem Kontext skaliert, ergeben längere Sequenzen größere Einsparungen sowohl beim VRAM als auch bei der Trainingszeit, wie hier beschrieben.

Ein großes Dankeschön an Rohan Pandey für seine Flex Attention-Implementierung, die direkt die Entwicklung von Unsloths Flex Attention-Implementierung inspiriert hat.

🕶️ Attention Sinks

Das GPT OSS Modell von OpenAI verwendet ein alternierendes Muster aus Sliding Window Attention, Full Attention, Sliding Window Attention und so weiter (SWA, FA, SWA, FA, etc.). Jedes Sliding Window beachtet nur 128 Token (inklusive des aktuellen Tokens), wodurch die Berechnung stark reduziert wird. Das bedeutet jedoch auch, dass lange Kontextwiedergewinnung und -schlussfolgerung aufgrund des kleinen Sliding Windows nutzlos werden. Die meisten Einrichtungen beheben dies, indem sie das Sliding Window auf 2048 oder 4096 Tokens erweitern.

OpenAI nutzte Attention Sinks aus den Efficient Streaming Language Models with Attention Sinks Papier , die zeigen, dass man ein kleines Sliding Window verwenden kann, allerdings muss man eine globale Attention auf das erste Token hinzufügen! Das Paper liefert dazu die folgende gute Illustration:

Das Paper stellt fest, dass der Aufmerksamkeitsmechanismus scheinbar viel Gewicht auf die ersten wenigen Tokens (1 bis 4) legt, und wenn diese während der Sliding Window-Operation entfernt werden, verschwinden diese „wichtigen“ ersten Tokens und verursachen schlechte Langzeit-Kontextwiedergewinnung.

Wenn wir die logarithmische Perplexität auftragen (höher ist schlechter) und nach der voreingestellten Kontextlänge des vortrainierten Modells Langzeit-Inferenz durchführen, sehen wir, dass die Perplexität stark ansteigt (nicht gut). Die rote Linie (verwendet Attention Sinks) bleibt jedoch niedrig, was sehr gut ist!

Das Paper zeigt auch, dass die Attention Is Off By One Methode teilweise funktioniert, allerdings muss man ein paar zusätzliche Sink-Tokens hinzufügen, um niedrigere Perplexitäten zu erzielen. Das Paper zeigt, dass das Hinzufügen eines einzigen lernbaren Sink-Tokens bemerkenswert gut funktioniert! Und genau das hat OpenAI für GPT-OSS getan!

📐Unsloths Flex Attention-Implementierung

Flex Attention https://pytorch.org/blog/flexattention/ ist äußerst mächtig, da sie dem Anwender 2 Anpassungswege für den Aufmerksamkeitsmechanismus bietet - ein Score-Modifikator (f) und ein Maskierungsfunktion (M).

Der Score-Modifikator (f) ermöglicht es uns, die Attention-Logits vor der Softmax-Operation zu bearbeiten, und das Maskierungsfunktion (M) ermöglicht uns, Operationen zu überspringen, wenn wir sie nicht benötigen (z. B. sieht Sliding Window Attention nur die letzten 128 Tokens).

Der Trick ist, dass Flex Attention schnelle automatisch generierte Triton-Kernel mit beliebigen Score-Modifikatoren und Maskierungsfunktionen bereitstellt!

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

Das bedeutet, dass wir Flex Attention verwenden können, um Attention Sinks zu implementieren! Die Implementierung eines einzelnen Attention Sinks ist sowohl in OpenAIs ursprünglichem GPT-OSS-Repo als auch in der Implementierung von HuggingFace transformers verfügbar.

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

Das Obige zeigt, dass wir das Sink ganz am Ende der Q @ K.T konkateniieren, die Softmax durchführen und die letzte Spalte entfernen, welche das Sink-Token war.

Durch die Verwendung einiger Visualisierungswerkzeuge aus Flex Attentions Github-Repo, können wir dies visualisieren. Angenommen, die Sequenzlänge war 16 und ein Sliding Window von 5. Links ist die letzte Sink-Spalte (Standardimplementierung) und rechts ist, wenn wir die Sink-Position auf Index 0 verschieben (unsere Implementierung).

Sink-Position am Ende (Standard)

Sink-Position auf Index 0 verschieben

Interessante Erkenntnis: Die offizielle Flex Attention Sliding Window-Implementierung betrachtet die Fenstergröße als die Anzahl der letzten Tokens PLUS EINS da sie das aktuelle Token einschließt. Die HuggingFace- und GPT OSS-Implementierungen sehen strikt nur die letzten N Tokens. D. h. das Folgende stammt von https://pytorch.org/blog/flexattention/ und 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

Standard Flex Attention (3+1 Tokens)

HuggingFace, GPT-OSS (3+0 Tokens)

Wir bestätigten dies auch anhand der offiziellen GPT-OSS-Implementierung von OpenAI, ob wir auf die letzten N oder N+1 Tokens achten, hier: 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
    )

Und wir sehen, dass nur die letzten 3 Tokens (nicht 3+1) beachtet werden! Das bedeutet, anstelle von <= SLIDING_WINDOW, verwenden Sie < SLIDING_WINDOW (d. h. Verwenden von weniger als, nicht gleich).

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

Da wir den Sink-Token-Index an den ersten Platz verschoben haben, müssen wir 1 zu q_idx addieren, um korrekt zu indexieren:

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
    """
    # Wir addieren (q_idx + 1), da die erste Spalte das Sink-Token ist
    causal_mask = (q_idx + 1) >= kv_idx
    sink_first_column = kv_idx == 0
    return causal_mask | sink_first_column

Um unsere Index-0-Implementierung zu bestätigen, verifizierten wir, dass der Trainingsverlust mit den Standard-Hugging-Face-Läufen (ohne Unsloth Flex Attention) konsistent bleibt, wie in unserem Diagramm gezeigt:

📜 Mathematische Herleitung für Attention Sinks

Es gibt eine andere Möglichkeit, die Attention Sinks ohne Padding von K und V zu berechnen. Zuerst stellen wir fest, dass die Softmax-Operation das tut, und wir wollen vorerst die zweite Version mit Sinks als Skalar:\

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

Wir können das logsumexp von Flex Attention über return_lse = True erhalten, und so tun wir:

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

Und wir können nun leicht die Sink-Version der Attention ableiten. Wir stellen jedoch fest, dass dieser Prozess etwas höhere Fehler als der Null-Padding-Ansatz aufweist, daher verwenden wir standardmäßig weiterhin unsere ursprüngliche Version.

💾NEU: Speicherung als GGUF, vLLM nach gpt-oss-Training

Sie können jetzt gpt-oss mit QLoRA feinabstimmen und das Modell direkt speichern, exportieren oder zusammenführen zu llama.cpp, vLLM, oder HF - nicht nur Unsloth. Wir werden hoffentlich bald ein kostenloses Notebook veröffentlichen.

Früher war jedes QLoRA feinabgestimmte gpt-oss Modell darauf beschränkt, in Unsloth zu laufen. Wir haben diese Einschränkung aufgehoben, indem wir die Möglichkeit eingeführt haben, in MXFP4 nativen Format mit save_method="mxfp4" und On-Demand-Dequantisierung von MXFP4 Basismodelle (wie gpt-oss) zusammenzuführen, wodurch es möglich ist, Ihr feinabgestimmtes Modell im bf16-Format zu exportieren mithilfe von save_method="merged_16bit" .

Der MXFP4 Das native Merge-Format bietet im Vergleich zum bf16-Formaterhebliche Leistungsverbesserungen: Es benötigt bis zu 75% weniger Festplattenspeicher, reduziert den VRAM-Verbrauch um 50%, beschleunigt das Mergen um das 5–10-fache und ermöglicht eine viel schnellere Konvertierung in das GGUF Format.

Nach der Feinabstimmung Ihres gpt-oss-Modells können Sie es in MXFP4 Format zusammenführen mit:

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

Wenn Sie das Modell lieber mergen und ins Hugging-Face-Repo pushen möchten, verwenden Sie:

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

Um Inferenz auf dem zusammengeführten Modell auszuführen, können Sie unter anderem vLLM und Llama.cpp verwenden. OpenAI empfiehlt diese Inference-Einstellungen für beide Modelle: temperature=1.0, top_p=1.0, top_k=0

✨ Speichern für Llama.cpp

1. Holen Sie sich die neueste llama.cpp auf GitHub hier. Sie können auch den unten stehenden Build-Anweisungen folgen. Ändern Sie -DGGML_CUDA=ON zu -DGGML_CUDA=OFF wenn Sie keine GPU haben oder nur CPU-Inferenz wünschen.

   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.cp

2. Konvertieren Sie das MXFP4 zusammengeführte Modell:

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

3. Führen Sie Inferenz auf dem quantisierten Modell aus:

   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 "The meaning to life and the universe is"

✨ Speichern nach SGLang

1. SGLang aus dem Quellcode bauen:\

   # aus dem Quellcode bauen
   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-Server starten:\

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

3. Inference ausführen:\

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

♦️Direktes Fine-Tuning von gpt-oss

Wir haben auch Unterstützung für direktes Fine-Tuning von gpt-oss Modellen hinzugefügt, indem wir Patches implementiert haben, die das Laden des nativen MXFP4-quantisierten Formats ermöglichen. Das macht es möglich, das Modell 'openai/gpt-oss' mit weniger als 24GB VRAM zu laden und mit QLoRA feinabzustimmen. Laden Sie das Modell einfach mit:

model, tokenizer = FastLanguageModel.from_pretrained(
    # model_name = "unsloth/gpt-oss-20b-BF16", 
    ] # Mehr Modelle unter https://huggingface.co/unsloth
    model_name = "unsloth/gpt-oss-20b",
    dtype = dtype, # None für automatische Erkennung
    max_seq_length = max_seq_length, # Wählen Sie beliebig für lange Kontexte!
    load_in_4bit = True,  # 4-Bit-Quantisierung zur Reduzierung des Speichers
    full_finetuning = False, # [NEU!] Wir haben jetzt Full-Finetuning!
)

fügen Sie eine Peft-Schicht hinzu mit FastLanguageModel.get_peft_model und führen Sie dann SFT-Fine-Tuning über das Peft-Modell aus.

🐛Bugfixes für gpt-oss

Wir wir haben kürzlich mit Hugging Face zusammengearbeitet um Inferenzprobleme zu lösen, indem wir OpenAIs Kernel verwendeten und sicherstellten, dass swiglu_limit = 7.0 während der MXFP4-Inferenz korrekt angewendet wird.

Basierend auf Nutzerfeedback entdeckten wir, dass verlängerte QLoRA-Trainingsläufe (über 60 Schritte hinaus) dazu führen konnten, dass der Verlust divergiert und schließlich zu einem Fehler führt. Dieses Problem trat nur auf Geräten auf, die BF16 nicht unterstützen und stattdessen auf F16 zurückfallen (z. B. T4 GPUs). Wichtig ist, dass es das QLoRA-Training auf A100- oder H100-GPUs sowie LoRA-Training auf f16-GPUs nicht beeinflusste.

Nach umfangreichen Untersuchungen haben wir das Verhalten des Trainingsverlusts nun über alle GPU-Setups hinweg angeglichen, einschließlich GPUs, die auf F16 beschränkt sind. Wenn Sie zuvor Probleme deswegen hatten, empfehlen wir die Verwendung unseres neuen aktualisierten gpt-oss Notebooks!

Wir mussten viele Experimente durchführen, um die Trainingsverlustkurve von float16 so zu verschieben, dass sie der von bfloat16-Maschinen (blaue Linie) entspricht. Wir fanden Folgendes:

1. Reines float16 geht bei Schritt 50 gegen unendlich

2. Wir fanden, dass die Down-Projektionen im MoE große Ausreißer aufweisen

3. Aktivierungen müssen in bfloat16 oder float32 gespeichert werden

Nachfolgend wird die absolute Größenordnung der Aktivierungen für GPT OSS 20B gezeigt, und einige spikes — dies würde auf float16-Maschinen überlaufen, da der maximale Bereich von float16 65504 beträgt.

Wir haben das in Unsloth behoben, sodass alle float16-Trainings sofort funktionieren!

🔢 Implementierungen für Sink Attention

OpenAIs Sink-Token-Implementierung ist hier verfügbar. Wir stellen sie unten zur Verfügung:

def sdpa(Q, K, V, S, sm_scale, sliding_window=0):
    # sliding_window == 0 bedeutet kein Sliding Window
    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)

Die HuggingFace transformers-Implementierung ist hier verfügbar. Wir stellen sie ebenfalls unten zur Verfügung:

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)

    # Dies war nicht in der ursprünglichen Implementierung und beeinflusst die Ergebnisse leicht; es verhindert Überläufe in BF16/FP16
    # wenn mit bsz>1 trainiert wird, klemmen wir die Maximalwerte.

    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]  # hier verwerfen wir das 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