Langkontext-gpt-oss-Training

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

• Sie können jetzt exportieren/speichern Ihr mit QLoRA feinabgestimmtes gpt-oss-Modell nach llama.cpp, vLLM, Ollama oder HF

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

• Wir haben die gpt-oss-Implementierungsprobleme behoben Probleme, die für Unsloth irrelevant waren, insbesondere die Sicherstellung, dass swiglu_limit = 7.0 während der MXFP4-Inferenz in transformers korrekt angewendet wird

🦥Einführung von Unsloth Flex Attention Support

Mit Unsloths Flex Attention Support kann eine einzelne H100 mit 80 GB VRAM bis zu 81K Kontextlänge mit QLoRA und 60K Kontext mit BF16 LoRA bewältigen! Diese Verbesserungen gelten für BEIDE gpt-oss-20b und gpt-oss-120b! Je mehr Kontextlänge Sie verwenden, desto größer sind die Vorteile von Unsloth Flex Attention:

Im Vergleich dazu erreichen alle anderen nicht von Unsloth stammenden Implementierungen maximal 9K Kontextlänge auf einer 80-GB-GPU und können mit FA3 nur 15K Kontext erreichen. Aber, FA3 ist für das gpt-oss-Training ungeeignet, da die Rückwärtsausbreitungsunterstützung für Attention Sinks fehlt. Wenn Sie FA3 zuvor für das gpt-oss-Training verwendet haben, würden wir Ihnen empfehlen, es vorerst nicht zu verwenden . Die maximale Kontextlänge, die Sie ohne Unsloth auf 80 GB VRAM erreichen können, beträgt also ~9K.

Das Training mit Unsloth Flex Attention liefert mindestens eine 1,3× Beschleunigung, wobei die Gewinne mit zunehmender Kontextlänge wachsen und bis zu 2× schneller werden. Da Flex Attention mit dem Kontext skaliert, bringen 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 die Entwicklung von Unsloths Flex Attention-Implementierung direkt inspiriert hat.

🕶️ Attention Sinks

OpenAIs GPT-OSS-Modell verwendet ein abwechselndes Muster aus Sliding-Window-Attention, Full Attention, Sliding-Window-Attention und so weiter (SWA, FA, SWA, FA usw.). Jedes Sliding Window beachtet nur 128 Tokens (einschließlich des aktuellen Tokens), sodass die Berechnung erheblich reduziert wird. Dies bedeutet jedoch auch, dass das Abrufen und Schlussfolgern über lange Kontexte aufgrund des kleinen Sliding Windows nutzlos wird. Die meisten Labs beheben dies, indem sie das Sliding Window auf 2048 oder 4096 Tokens erweitern.

OpenAI nutzte Attention Sinks aus dem Paper Efficient Streaming Language Models with Attention Sinks das zeigt, dass man ein kleines Sliding Window verwenden kann, jedoch eine globale Attention auf das erste Token hinzufügen muss! Das Paper bietet unten eine gute Illustration:

Das Paper stellt fest, dass der Attention-Mechanismus den ersten wenigen Tokens (1 bis 4) offenbar viel Gewicht zuweist, und indem sie während der Sliding-Window-Operation entfernt werden, verschwinden diese „wichtigen“ ersten wenigen Tokens, was zu schlechter Langkontext-Retrieval-Leistung führt.

Wenn wir die Log-Perplexity (höher ist schlechter) auftragen und Langkontext-Inferenz nach der im vortrainierten Modell festgelegten Kontextlänge durchführen, sehen wir, dass die Perplexity stark ansteigt (nicht gut). Die rote Linie (verwendet Attention Sinks) bleibt jedoch niedrig, was sehr gut ist!

Das Paper zeigt auch, dass die Methode „Attention Is Off By One“ teilweise funktioniert, allerdings muss man auch einige zusätzliche Sink-Tokens hinzufügen, um niedrigere Perplexities zu erhalten. 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 leistungsfähig, da sie dem Anwender 2 Anpassungswege für den Attention-Mechanismus bietet - einen Score-Modifikator (f) und eine Maskierungsfunktion (M).

Die Score-Modifikator (f) ermöglicht es uns, die Attention-Logits vor der Softmax-Operation zu bearbeiten, und die Maskierungsfunktion (M) ermöglicht es 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-Kernels 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 Sink ist sowohl im ursprünglichen GPT-OSS-Repo von OpenAI als auch in der Transformers-Implementierung von HuggingFace enthalten.

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

Das oben Gezeigte zeigt, dass wir den Sink ganz am Ende von Q @ K.T anhängen, die Softmax durchführen und die letzte Spalte entfernen, die das Sink-Token war.

Mit einigen Visualisierungswerkzeugen aus dem Github-Repo von Flex Attentionkönnen wir das visualisieren. Nehmen wir an, die Sequenzlänge beträgt 16 und das Sliding Window 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 das aktuelle Token eingeschlossen ist. Die HuggingFace- und GPT-OSS-Implementierungen sehen strikt nur die letzten N Tokens. Also stammt das Folgende aus 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 haben auch anhand der offiziellen GPT-OSS-Implementierung von OpenAI bestätigt, ob wir hier die letzten N- oder N+1-Tokens beachten: 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, dass wir statt <= SLIDING_WINDOW, verwende < SLIDING_WINDOW verwenden sollten (also kleiner 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 außerdem den Index des Sink-Tokens an den Anfang verschoben haben, müssen wir 1 zu q_idx addieren, um korrekt zu indizieren:

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 Implementierung mit Index 0 zu bestätigen, haben wir überprüft, 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 noch eine andere Möglichkeit, die Attention Sinks zu berechnen, ohne K und V zu paddieren. Wir stellen zunächst fest, dass die Softmax-Operation das tut, und wir wollen die zweite Version mit Sinks vorerst 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 aus Flex Attention über return_lse = True erhalten, und daher machen 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 die Sink-Version der Attention leicht herleiten. Wir stellen jedoch fest, dass dieser Prozess etwas höheren Fehler aufweist als der Zero-Padding-Ansatz, daher verwenden wir weiterhin standardmäßig unsere ursprüngliche Version.

💾NEU: Speichern in GGUF, vLLM nach dem gpt-oss-Training

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

Zuvor war jedes mit QLoRA feinabgestimmte gpt-oss-Modell darauf beschränkt, nur in Unsloth ausgeführt zu werden. Wir haben diese Einschränkung entfernt, indem wir die Möglichkeit eingeführt haben, im MXFP4 nativen Format unter Verwendung von save_method="mxfp4" und On-Demand-Dekquantisierung von MXFP4 Basismodellen (wie gpt-oss) zusammenzuführen, wodurch es möglich wird, Ihr feinabgestimmtes Modell im bf16-Format zu exportieren mit save_method="merged_16bit" .

Die MXFP4 Das native Merge-Format bietet gegenüber dembf16-Format deutliche Leistungsverbesserungen: Es benötigt bis zu 75 % weniger Speicherplatz, reduziert den VRAM-Verbrauch um 50 %, beschleunigt das Zusammenführen um das 5-10-Fache und ermöglicht eine deutlich schnellere Konvertierung in das GGUF

Neu: Das Speichern oder Zusammenführen von mit QLoRA feinabgestimmten Modellen zu GGUF wird jetzt für die Verwendung in anderen Frameworks unterstützt (z. B. Hugging Face, llama.cpp mit GGUF). MXFP4 Nachdem du dein gpt-oss-Modell feinabgestimmt hast, kannst du es mit Folgendem in das

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

Wenn Sie das Modell lieber zusammenführen und auf den Hugging-Face-Hub hochladen 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 Inferenz-Einstellungen für beide Modelle: temperature=1.0, top_p=1.0, top_k=0

✨ Speichern für Llama.cpp

1. Hole dir die neueste llama.cpp auf GitHub hier. Du kannst auch den untenstehenden Build-Anweisungen folgen. Ändere -DGGML_CUDA=ON zu -DGGML_CUDA=OFF wenn du keine GPU hast oder nur CPU-Inferenz möchtest.

   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. Konvertiere 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ühre 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 "Der Sinn des Lebens und des Universums ist"

✨ Speichern in SGLang

1. Bauen Sie SGLang aus dem Quellcode:\

   # 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. Starten Sie den SGLang-Server:\

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

3. Inferenz 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())

♦️gpt-oss direkt feinabstimmen

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

model, tokenizer = FastLanguageModel.from_pretrained(
    # model_name = "unsloth/gpt-oss-20b-BF16", 
    model_name = "unsloth/gpt-oss-20b",
    dtype = dtype, # None für automatische Erkennung
    max_seq_length = max_seq_length, # Wähle beliebig für langen Kontext!
    load_in_4bit = True,  # 4-Bit-Quantisierung zur Speicherreduzierung
    full_finetuning = False, # [NEU!] Wir haben jetzt Full-Finetuning!
    # token = "hf_...", # verwende einen, wenn du geschützte Modelle nutzt
)

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

🐛Fehlerbehebungen für gpt-oss

Wir kürzlich mit Hugging Face zusammengearbeitet um Inferenzprobleme zu beheben, indem OpenAIs Kernel verwendet und sichergestellt wird, dass swiglu_limit = 7.0 während der MXFP4-Inferenz korrekt angewendet wird.

Basierend auf Nutzerfeedback haben wir entdeckt, dass längere QLoRA-Trainingsläufe (über 60 Schritte hinaus) dazu führen konnten, dass die Verluste divergierten und schließlich einen Fehler auslösten. Dieses Problem trat nur auf Geräten auf, die BF16 nicht unterstützen und stattdessen auf F16 zurückgreifen (z. B. T4-GPUs). Wichtig ist, dass es weder QLoRA-Training auf A100- oder H100-GPUs noch LoRA-Training auf f16-GPUs beeinträchtigte.

Nach umfangreichen Untersuchungen haben wir nun das Verhalten des Trainingsverlusts über alle GPU-Setups hinweg angeglichen, einschließlich GPUs, die auf F16 beschränkt sind. Wenn Sie zuvor deshalb Probleme hatten, empfehlen wir Ihnen, unser neues aktualisiertes gpt-oss-Notebook zu verwenden!

Wir mussten viele, viele Experimente durchführen, um die Trainingsverlustkurve von float16 der von bfloat16-Maschinen (blaue Linie) anzugleichen. Wir haben Folgendes festgestellt:

1. Reines float16 geht bei Schritt 50 gegen unendlich

2. Wir haben festgestellt, dass die Down-Projektionen im MoE enorme Ausreißer aufweisen

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

Unten sind die absoluten Aktivierungsbeträge für GPT OSS 20B gezeigt, und einige spiken wirklich - dies führt auf float16-Maschinen zu einem Überlauf, da der maximale Bereich von float16 65504 beträgt.

Wir haben dies in Unsloth behoben, sodass das gesamte float16-Training sofort funktioniert!

🔢 Implementierungen für Sink Attention

OpenAIs Implementierung des Sink-Tokens ist hier bereitgestellt. Wir stellen sie unten bereit:

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 bereitgestellt. Wir stellen sie auch unten bereit:

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 enthalten und beeinflusst die Ergebnisse leicht; es verhindert Überläufe in BF16/FP16
    # wenn mit bsz>1 trainiert wird, begrenzen 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]  # wir verwerfen hier den 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