Entraînement gpt-oss à contexte long

Nous sommes ravis de présenter la prise en charge d’Unsloth Flex Attention pour l’entraînement OpenAI gpt-oss, ce qui permet >8× plus longues longueurs de contexte, >50 % d’utilisation de VRAM en moins et un entraînement >1,5× plus rapide (sans dégradation de la précision) par rapport à toutes les implémentations, y compris celles utilisant Flash Attention 3 (FA3). Unsloth Flex Attention permet de s’entraîner avec une longueur de contexte de 60K sur un GPU H100 avec 80 Go de VRAM pour BF16 LoRA. Aussi :

• Vous pouvez désormais exporter/enregistrer votre modèle gpt-oss affiné en QLoRA vers llama.cpp, vLLM, Ollama ou HF

• Nous avons corrigé l’entraînement gpt-oss les pertes qui divergeaient vers l’infini sur des GPU float16 (comme le T4 Colab)

• Nous avons corrigé l’implémentation gpt-oss des problèmes sans rapport avec Unsloth, notamment en garantissant que swiglu_limit = 7.0 est correctement appliqué lors de l’inférence MXFP4 dans transformers

🦥Présentation de la prise en charge d’Unsloth Flex Attention

Avec la prise en charge de Flex Attention par Unsloth, un seul H100 avec 80 Go de VRAM peut gérer jusqu’à 81K de longueur de contexte avec QLoRA et 60K de contexte avec BF16 LoRA ! Ces gains s’appliquent à LES DEUX gpt-oss-20b et gpt-oss-120b! Plus vous utilisez une grande longueur de contexte, plus vous gagnerez avec Unsloth Flex Attention :

En comparaison, toutes les autres implémentations non-Unsloth plafonnent à 9K de longueur de contexte sur un GPU de 80 Go, et ne peuvent atteindre que 15K de contexte avec FA3. Mais, FA3 ne convient pas à l’entraînement de gpt-oss car il ne prend pas en charge la rétropropagation pour les puits d’attention. Donc, si vous utilisiez auparavant FA3 pour l’entraînement de gpt-oss, nous vous recommandons de ne pas l’utiliser pour le moment. Ainsi, la longueur de contexte maximale que vous pouvez obtenir sans Unsloth sur 80 Go de VRAM est d’environ 9K.

L’entraînement avec Unsloth Flex Attention offre au moins un accélération de 1,3×, avec des gains qui augmentent à mesure que la longueur de contexte augmente, jusqu’à 2× plus rapide. Comme Flex Attention évolue avec le contexte, les séquences plus longues entraînent des économies plus importantes en VRAM et en temps d’entraînement, comme décrit ici.

Un grand merci à Rohan Pandey pour son implémentation de Flex Attention, qui a directement inspiré le développement de l’implémentation Flex Attention d’Unsloth.

🕶️ Puits d’attention

Le modèle GPT OSS d’OpenAI utilise un pattern alterné d’attention à fenêtre glissante, d’attention complète, d’attention à fenêtre glissante, etc. (SWA, FA, SWA, FA, etc.). Chaque fenêtre glissante n’accède qu’à 128 jetons (y compris le jeton actuel), ce qui réduit énormément le calcul. Cependant, cela signifie aussi que la récupération et le raisonnement sur long contexte deviennent inutiles à cause de la petite fenêtre glissante. La plupart des laboratoires corrigent cela en élargissant la fenêtre glissante à 2048 ou 4096 jetons.

OpenAI s’est inspiré Puits d’attention de l’article Efficient Streaming Language Models with Attention Sinks papier qui montre qu’on peut utiliser une petite fenêtre glissante, à condition d’ajouter une attention globale sur le premier jeton ! L’article fournit une bonne illustration ci-dessous :

L’article constate que le mécanisme d’attention semble attribuer beaucoup de poids aux premiers jetons (1 à 4), et en les supprimant lors de l’opération de fenêtre glissante, ces premiers jetons « importants » disparaissent, ce qui entraîne une mauvaise récupération sur long contexte.

Si l’on trace la perplexité logarithmique (plus c’est élevé, pire c’est), et que l’on fait une inférence sur long contexte après la longueur de contexte définie du modèle préentraîné, on voit la perplexité grimper en flèche (pas bon). Cependant, la ligne rouge (qui utilise Attention Sinks) reste basse, ce qui est très bien !

L’article montre aussi que la méthode Attention Is Off By One fonctionne partiellement, sauf qu’il faut aussi ajouter quelques jetons de puits supplémentaires pour obtenir de plus faibles perplexités. L’article montre que l’ajout d’un seul jeton de puits, apprenable, fonctionne remarquablement bien ! Et c’est ce qu’OpenAI a fait pour GPT-OSS !

📐L’implémentation Flex Attention d’Unsloth

Flex Attention https://pytorch.org/blog/flexattention/ est extrêmement puissante, car elle offre au praticien 2 voies de personnalisation pour le mécanisme d’attention - un modificateur de score (f) et une fonction de masquage (M).

Le modificateur de score (f) permet de modifier les logits d’attention avant l’opération softmax, et la fonction de masquage (M) permet de sauter des opérations si nous n’en avons pas besoin (par exemple, l’attention à fenêtre glissante ne voit que les 128 derniers jetons).

L’astuce, c’est que Flex Attention fournit rapidement des kernels Triton auto-générés avec des modificateurs de score et des fonctions de masquage arbitraires !

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

Cela signifie que nous pouvons utiliser Flex Attention pour implémenter des puits d’attention ! L’implémentation d’un seul puits d’attention est fournie à la fois dans le dépôt GPT-OSS original d’OpenAI et dans l’implémentation des transformers de HuggingFace.

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

Ce qui précède montre que nous concaténons le puits tout à la fin du Q @ K.T , effectuons le softmax, puis retirons la dernière colonne, qui était le jeton de puits.

En utilisant quelques utilitaires de visualisation du dépôt Github de Flex Attention, nous pouvons visualiser cela. Supposons que la longueur de séquence soit 16, avec une fenêtre glissante de 5. À gauche se trouve la dernière colonne de puits (implémentation par défaut), et à droite, si nous déplaçons l’emplacement du puits à l’index 0 (notre implémentation).

Emplacement du puits à la fin (par défaut)

Déplacer l’emplacement du puits à l’index 0

Découverte intéressante: Les implémentations officielles de fenêtre glissante de Flex Attention considèrent la taille de la fenêtre comme le nombre des derniers jetons PLUS UN car elles incluent le jeton actuel. Les implémentations HuggingFace et GPT OSS ne voient strictement que les N derniers jetons. Autrement dit, ce qui suit provient de https://pytorch.org/blog/flexattention/ et 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 par défaut (3+1 jetons)

HuggingFace, GPT-OSS (3+0 jetons)

Nous avons aussi confirmé via l’implémentation officielle GPT-OSS d’OpenAI si l’on s’intéresse ici aux N derniers jetons ou à N+1 jetons : 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
    )

Et nous voyons que seuls les 3 derniers jetons (pas 3+1) sont pris en compte ! Cela signifie qu’au lieu d’utiliser <= SLIDING_WINDOW, utilisez < SLIDING_WINDOW (c.-à-d. utiliser inférieur à, pas égal).

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 par défaut
    window_mask = q_idx - kv_idx <  SLIDING_WINDOW # version GPT-OSS
    return causal_mask & window_mask

De plus, comme nous avons déplacé l’index du jeton de puits au premier, nous devons ajouter 1 à q_idx pour indexer correctement :

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
    """
    # Nous ajoutons (q_idx + 1) puisque la première colonne est le jeton de puits
    causal_mask = (q_idx + 1) >= kv_idx
    sink_first_column = kv_idx == 0
    return causal_mask | sink_first_column

Pour confirmer notre implémentation avec index 0, nous avons vérifié que la perte d’entraînement reste cohérente avec les exécutions standard de Hugging Face (sans Unsloth Flex Attention), comme montré dans notre graphe :

📜 Dérivation mathématique pour les puits d’attention

Il existe une autre façon de calculer les puits d’attention sans remplir K et V. Nous notons d’abord l’opération softmax, et nous voulons pour l’instant la 2e version avec puits comme un scalaire :\

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

Nous pouvons obtenir le logsumexp depuis Flex Attention via return_lse = True , puis nous faisons :

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

Et nous pouvons maintenant dériver facilement la version avec puits de l’attention. Nous constatons toutefois que ce processus présente une erreur un peu plus élevée que l’approche de remplissage par zéro, donc nous conservons par défaut notre version originale.

💾NOUVEAU : Enregistrement vers GGUF, vLLM après l’entraînement de gpt-oss

Vous pouvez maintenant affiner gpt-oss avec QLoRA et directement enregistrer, exporter ou fusionner le modèle vers llama.cpp, vLLM, ou HF - pas seulement Unsloth. Nous publierons, espérons-le, bientôt un notebook gratuit.

Auparavant, tout modèle gpt-oss affiné en QLoRA était limité à une exécution dans Unsloth. Nous avons supprimé cette limitation en introduisant la capacité de fusionner dans MXFP4 format natif en utilisant save_method="mxfp4" et déquantification à la demande de MXFP4 des modèles de base (comme gpt-oss), ce qui permet de exporter votre modèle affiné au format bf16 en utilisant save_method="merged_16bit" .

Le MXFP4 format de fusion natif offre des améliorations de performance significatives par rapport au format bf16 : il utilise jusqu’à 75 % d’espace disque en moins, réduit la consommation de VRAM de 50 %, accélère la fusion de 5 à 10× et permet une conversion beaucoup plus rapide au format GGUF .

Après avoir affiné votre modèle gpt-oss, vous pouvez le fusionner au format MXFP4 avec :

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

Si vous préférez fusionner le modèle et le pousser vers le hub Hugging Face, utilisez :

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

Pour exécuter l’inférence sur le modèle fusionné, vous pouvez utiliser vLLM et Llama.cpp, entre autres. OpenAI recommande ces paramètres d’inférence pour les deux modèles : temperature=1.0, top_p=1.0, top_k=0

✨ Enregistrement vers Llama.cpp

1. Obtenez la dernière version llama.cpp sur GitHub ici. Vous pouvez également suivre les instructions de compilation ci-dessous. Changez -DGGML_CUDA=ON en -DGGML_CUDA=OFF si vous n’avez pas de GPU ou si vous souhaitez simplement une inférence 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

2. Convertissez le MXFP4 modèle fusionné :

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

3. Exécutez l’inférence sur le modèle quantifié :

   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"

✨ Enregistrement vers SGLang

1. Construisez SGLang depuis les sources :\

   # compiler depuis les sources
   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. Lancer le serveur SGLang :\

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

3. Exécuter l’inférence :\

   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": "Quelle est la capitale de la France ?"}],
   }
   
   response = requests.post(url, json=data)
   print_highlight(response.json())

♦️Ajuster directement gpt-oss

Nous avons également ajouté la prise en charge de l’affinage direct des modèles gpt-oss en implémentant des patchs qui permettent de charger le format quantifié natif MXFP4. Cela permet de charger le modèle 'openai/gpt-oss' avec moins de 24 Go de VRAM, et de l’affiner en QLoRA. Il suffit de charger le modèle en utilisant :

model, tokenizer = FastLanguageModel.from_pretrained(
    # model_name = "unsloth/gpt-oss-20b-BF16", 
    model_name = "unsloth/gpt-oss-20b",
    dtype = dtype, # None pour détection automatique
    max_seq_length = max_seq_length, # Choisissez n’importe quelle valeur pour un long contexte !
    load_in_4bit = True,  # Quantification 4 bits pour réduire la mémoire
    full_finetuning = False, # [NOUVEAU !] Nous prenons désormais en charge le fine-tuning complet !
    # token = "hf_...", # en utilisez un si vous utilisez des modèles à accès restreint
)

ajouter une couche Peft en utilisant FastLanguageModel.get_peft_model et lancer un affinage SFT sur le modèle Peft.

🐛Corrections de bugs pour gpt-oss

Nous a récemment collaboré avec Hugging Face pour résoudre les problèmes d’inférence en utilisant les kernels d’OpenAI et en s’assurant que swiglu_limit = 7.0 est correctement appliqué lors de l’inférence MXFP4.

D’après les retours des utilisateurs, nous avons découvert que les longues sessions d’entraînement QLoRA (au-delà de 60 étapes) pouvaient provoquer une divergence de la perte puis une erreur finale. Ce problème ne se produisait que sur les appareils qui ne prennent pas en charge BF16 et retombent à la place sur F16 (par exemple, les GPU T4). Important : cela n’affectait pas l’entraînement QLoRA sur les GPU A100 ou H100, ni l’entraînement LoRA sur les GPU f16.

Après une enquête approfondie, nous avons maintenant harmonisé le comportement de la perte d’entraînement sur toutes les configurations GPU, y compris les GPU limités à F16. Si vous rencontriez auparavant des problèmes à cause de cela, nous vous recommandons d’utiliser notre nouveau notebook gpt-oss mis à jour !

Nous avons dû faire de très nombreuses expériences pour rendre la courbe de perte d’entraînement du float16 équivalente à celle des machines bfloat16 (ligne bleue). Nous avons constaté ce qui suit :

1. Le float16 pur divergera vers l’infini à l’étape 50

2. Nous avons trouvé que les projections descendantes dans le MoE présentaient de très fortes valeurs aberrantes

3. Les activations doivent être enregistrées en bfloat16 ou float32

Ci-dessous, on voit les activations de magnitude absolue pour GPT OSS 20B, et certaines connaissent de très fortes pointes - cela débordera sur les machines float16, car la plage maximale du float16 est 65504.

Nous avons corrigé cela dans Unsloth, donc tout l’entraînement float16 fonctionne désormais immédiatement !

🔢 Implémentations pour Sink Attention

L’implémentation du jeton de puits d’OpenAI est fournie ici. Nous la fournissons ci-dessous :

def sdpa(Q, K, V, S, sm_scale, sliding_window=0):
    # sliding_window == 0 signifie aucune fenêtre glissante
    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)

L’implémentation des transformers HuggingFace est fournie ici. Nous la fournissons également ci-dessous :

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)

    # Ceci ne faisait pas partie de l’implémentation originale et affecte légèrement les résultats ; cela empêche les débordements en BF16/FP16
    # lors de l’entraînement avec bsz>1, nous plafonnons les valeurs maximales.

    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]  # nous supprimons ici le jeton de puits
    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