Entraînement gpt-oss à long contexte

Nous sommes ravis de présenter la prise en charge d'Unsloth Flex Attention pour l'entraînement OpenAI gpt-oss qui permet >8× des longueurs de contexte plus longues, >50% de réduction de l'utilisation de la VRAM et >1,5× d'entraînement 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 rend possible l'entraînement avec une longueur de contexte de 60K sur un GPU H100 80GB VRAM pour BF16 LoRA. De plus :

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

• Nous avons corrigé l'entraînement gpt-oss les pertes allant vers l'infini sur les GPU float16 (comme les T4 Colab)

• Nous corrigé l'implémentation gpt-oss problèmes sans rapport avec Unsloth, notamment en veillant à ce que swiglu_limit = 7.0 soit correctement appliqué pendant l'inférence MXFP4 dans transformers

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

Avec la prise en charge Flex Attention d'Unsloth, un seul H100 80GB 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 longueur de contexte importante, plus vous bénéficierez des gains d'Unsloth Flex Attention :

En comparaison, toutes les autres implémentations non-Unsloth plafonnent à 9K de longueur de contexte sur un GPU 80GB, et ne peuvent atteindre que 15K de contexte avec FA3. Mais, FA3 est inadapté pour l'entraînement gpt-oss car il ne prend pas en charge la passe arrière pour les attention sinks. Donc si vous utilisiez auparavant FA3 pour l'entraînement gpt-oss, nous vous recommandons de ne pas l'utiliser pour l'instant. Ainsi, la longueur de contexte maximale que vous pouvez obtenir sans Unsloth sur 80GB 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 avec la longueur de contexte, atteignant jusqu'à 2× plus rapide. Parce que Flex Attention s'adapte au contexte, les séquences plus longues produisent 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 Flex Attention, qui a directement inspiré le développement de l'implémentation Flex Attention d'Unsloth.

🕶️ Attention Sinks

Le modèle GPT OSS d'OpenAI utilise un schéma alterné d'attention en fenêtre glissante, attention complète, attention en fenêtre glissante, etc. (SWA, FA, SWA, FA, etc.). Chaque fenêtre glissante n'attend que 128 tokens (y compris le token courant), donc le calcul est considérablement réduit. Cependant, cela signifie aussi que la récupération et le raisonnement sur de longs contextes deviennent inutiles en raison de la petite fenêtre glissante. La plupart des laboratoires corrigent cela en étendant la fenêtre glissante à 2048 ou 4096 tokens.

OpenAI s'est inspiré de Attention Sinks l'article Efficient Streaming Language Models with Attention Sinks Arctic Long Sequence Training qui montre que vous pouvez utiliser une petite fenêtre glissante, à condition d'ajouter une attention globale sur le premier token ! L'article fournit une bonne illustration ci-dessous :

L'article constate que le mécanisme d'attention semble attribuer beaucoup de poids aux premiers tokens (1 à 4), et en les supprimant pendant l'opération de fenêtre glissante, ces premiers tokens "importants" disparaissent, entraînant de mauvaises performances de récupération sur de longs contextes.

Si nous traçons la perplexité logarithmique (plus c'est élevé, pire c'est), et effectuons une inférence sur de longs contextes au-delà de la longueur de contexte définie du modèle préentraîné, nous voyons la perplexité augmenter brusquement (pas bon). Cependant la courbe rouge (utilise Attention Sinks) reste basse, ce qui est très bien !

L'article montre également que la méthode Attention Is Off By One fonctionne partiellement, sauf qu'il faut aussi ajouter quelques tokens sink supplémentaires pour obtenir des perplexités plus faibles. L'article montre qu'ajouter un seul token sink qui est apprenable donne des résultats remarquables ! Et c'est ce qu'OpenAI a fait pour GPT-OSS !

📐L'implémentation Flex Attention d'Unsloth

et notre fonctionnalité Standby dans 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 un fonction de masquage (M).

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

L'astuce est que Flex Attention fournit des kernels Triton auto-générés rapides 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 attention sinks ! L'implémentation d'un seul attention sink est fournie à la fois dans le dépôt original GPT-OSS d'OpenAI et 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 sink à la toute fin du Q @ K.T , effectuons le softmax, et supprimons la dernière colonne qui était le token sink.

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

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

Déplacer l'emplacement du sink à l'index 0

Constat intéressant : Les implémentations officielles de Flex Attention pour la fenêtre glissante considèrent la taille de la fenêtre comme le nombre des derniers tokens PLUS UN car elle inclut le token courant. Les implémentations HuggingFace et GPT OSS voient strictement seulement les N derniers tokens. Par ex. 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 tokens)

HuggingFace, GPT-OSS (3+0 tokens)

Nous avons également confirmé via l'implémentation officielle GPT-OSS d'OpenAI si nous assistons aux N derniers ou N+1 tokens ici : 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 seulement les 3 derniers tokens (et non 3+1) sont pris en compte ! Cela signifie qu'au lieu d'utiliser <= SLIDING_WINDOW, utilisez < SLIDING_WINDOW (c.-à-d. utiliser moins que, et non é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, puisque nous avons déplacé l'index du token sink en 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 token sink
    causal_mask = (q_idx + 1) >= kv_idx
    sink_first_column = kv_idx == 0
    return causal_mask | sink_first_column

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

📜 Dérivation mathématique pour les attention sinks

Il existe une autre façon de calculer les attention sinks sans padding de K et V. Nous notons d'abord que l'opération softmax fait, et nous voulons la 2ᵉ version avec sinks pour l'instant 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 , et donc 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 facilement dériver la version sink de l'attention. Nous constatons toutefois que ce processus présente une erreur quelque peu plus élevée que l'approche du padding à zéro, donc nous restons par défaut sur notre version originale.

💾NOUVEAU : Sauvegarde en GGUF, vLLM après entraînement gpt-oss

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

Auparavant, tout modèle gpt-oss affiné avec QLoRA était limité à être exécuté dans Unsloth. Nous avons supprimé cette limitation en introduisant la possibilité de fusionner dans le MXFP4 format natif en utilisant save_method="mxfp4" et déquantification à la demande de MXFP4 modèles de base (comme gpt-oss) le rendant possible de exporter votre modèle affiné au format bf16 en utilisant save_method="merged_16bit" .

Le MXFP4 le 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 en GGUF format.

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

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

Si vous préférez fusionner le modèle et le pousser sur 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 le dernier llama.cpp sur GitHub ici. Vous pouvez suivre les instructions de compilation ci-dessous également. Changez -DGGML_CUDA=ON en -DGGML_CUDA=OFF si vous n'avez pas de GPU ou si vous voulez 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.cp

2. Convertir 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écuter 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 "Le sens de la vie et de l'univers est"

✨ Enregistrement vers SGLang

1. Construire SGLang depuis la source :\

   # construire depuis la source
   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())

♦️Affinage direct de gpt-oss

Nous avons également ajouté la prise en charge de l'affinage direct des modèles gpt-oss en implémentant des correctifs qui permettent de charger le format quantifié natif MXFP4. Cela rend possible de charger le modèle 'openai/gpt-oss' avec moins de 24GB de VRAM, et de l'affiner avec QLoRA. Chargez simplement 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 contexte long !
    load_in_4bit = True,  # quantification 4 bits pour réduire la mémoire
    full_finetuning = False, # [NOUVEAU!] Nous avons maintenant le fine-tuning complet !
    # token = "hf_...", # utilisez-en un si vous utilisez des modèles à accès restreint
)

ajoutez une couche Peft en utilisant FastLanguageModel.get_peft_model et exécutez l'affinage SFT sur le modèle Peft.

🐛Corrections de bugs pour gpt-oss

Nous a récemment collaboré avec Hugging Face pour résoudre des problèmes d'inférence en utilisant les kernels d'OpenAI et en veillant à ce que swiglu_limit = 7.0 soit correctement appliqué pendant l'inférence MXFP4.

D'après les retours des utilisateurs, nous avons découvert que des sessions d'entraînement QLoRA prolongées (au-delà de 60 étapes) pouvaient provoquer la perte à diverger et finalement générer une erreur. Ce problème ne se produisait que sur des appareils qui ne prennent pas en charge BF16 et qui basculent plutôt en F16 (par ex., les GPU T4). Il est important de noter que cela n'a pas impacté 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 aligné le comportement de la perte d'entraînement sur tous les configurations GPU, y compris les GPU limités à F16. Si vous rencontrez précédemment des problèmes à cause de cela, nous vous recommandons d'utiliser notre nouveau notebook gpt-oss mis à jour !

Nous avons dû réaliser de nombreuses expériences pour amener la courbe de perte d'entraînement du float16 à être équivalente à celle des machines bfloat16 (ligne bleue). Nous avons constaté ce qui suit :

1. Le float16 pur ira à l'infini à l'étape 50

2. Nous avons trouvé que les projections vers le bas dans le MoE avaient de très grands outliers

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

Ci-dessous sont montrées les magnitudes absolues des activations pour GPT OSS 20B, et certaines présentent de fortes pointes - cela débordera sur des machines float16 puisque la plage maximale du float16 est 65504.

Nous avons corrigé cela dans Unsloth, donc tout l'entraînement en float16 fonctionne immédiatement !

🔢 Implémentations pour Sink Attention

L'implémentation du token sink d'OpenAI est fourni ici. Nous la fournissons ci-dessous :

def sdpa(Q, K, V, S, sm_scale, sliding_window=0):
    # sliding_window == 0 signifie pas de 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 fourni 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 n'était pas dans 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 limitons 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 le sink ici
    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