Affinage avec contexte de 500K
Apprenez comment activer l'ajustement avec une fenêtre de contexte >500K tokens avec Unsloth.
Nous présentons de nouveaux algorithmes dans Unsloth qui repoussent les limites de l'entraînement en contexte long pour n'importe quel LLM et VLM. Les entraînements de LLMs comme gpt-oss-20b peuvent désormais atteindre des longueurs de contexte supérieures à 500K sur un seul GPU H100 80GB, contre 80K auparavant sans dégradation de la précision.
Vous pouvez atteindre >fenêtres de contexte de 750K sur un GPU B200 192GB.
Essayez le fine-tuning gpt-oss-20b à 500K de contexte sur notre notebook Colab 80GB A100.
Nous avons considérablement amélioré la façon dont Unsloth gère les schémas d'utilisation de la mémoire, la vitesse et les longueurs de contexte :
Utilisation de VRAM réduite de 60% avec Contexte 3,2× plus long via le nouveau loss cross-entropie fusionné et découpé d'Unsloth avec aucune dégradation de la vitesse ou de la précision
Décharge d'activations améliorée dans le Gradient Checkpointing
Collaboration avec Stas Bekman de Snowflake sur Tiled MLP, permettant 2× plus de contextes
Les algorithmes d'Unsloth permettent au QLoRA (4 bits) gpt-oss-20b d'atteindre 290K de contexte possible sur une H100 sans perte de précision, et plus de 500K avec Tiled MLP activé, offrant au total >des longueurs de contexte 6,4× plus longues.
📐 Refactoring du Loss d'Unsloth : Découper & Fusionner
Notre nouvelle implémentation du loss fusionné ajoute un découpage dynamique de séquence : au lieu de calculer les logits de la tête du modèle de langage et les cross-entropies sur l'ensemble de la séquence d'un coup, nous traitons des tranches gérables le long de la dimension de séquence aplatie. Cela réduit la mémoire de pointe de plusieurs Go à des tailles de tranche plus petites. Chaque tranche exécute toujours un passage avant + arrière entièrement fusionné via torch.func.grad_and_value , et conserve la précision en précision mixte en suréchantillonnant en float32 si nécessaire. Ces changements ne dégradent ni la vitesse ni la précision de l'entraînement.
L'innovation clé est que la taille de la tranche est choisie automatiquement à l'exécution en fonction de la VRAM disponible.
Si vous disposez de plus de VRAM libre, des tranches plus grandes sont utilisées pour des exécutions plus rapides
Si vous avez moins de VRAM, cela augmente le nombre de tranches pour éviter les explosions de mémoire.
Cela supprime le réglage manuel et maintient notre algorithme robuste sur les GPU anciens et nouveaux, les charges de travail et les différentes longueurs de séquence.
Grâce au réglage automatique, les contextes plus petits utiliseront plus de VRAM (moins de tranches) pour éviter des frais généraux inutiles. Pour les graphiques ci-dessus, nous ajustons le nombre de tranches de loss pour refléter des paliers de VRAM réalistes. Avec 80GB de VRAM, cela permet des contextes >3,2× plus longs.
🏁 Améliorations du Gradient Checkpointing d'Unsloth
Notre Unsloth Gradient Checkpointing algorithme, introduit en avril 2024, est rapidement devenu populaire et la norme dans l'industrie, ayant été intégré dans la plupart des packages d'entraînement aujourd'hui. Il décharge les activations vers la RAM CPU, ce qui a permis des longueurs de contexte 10× plus longues. Nos nouvelles améliorations utilisent des CUDA Streams et d'autres astuces pour ajouter au maximum 0.1% un surcoût d'entraînement sans impact sur la précision. Auparavant, il ajoutait 1 à 3% de surcoût d'entraînement.
En déchargeant les activations dès qu'elles sont produites, nous minimisons l'empreinte d'activation maximale et libérons la mémoire GPU exactement quand elle est nécessaire. Cela réduit fortement la pression mémoire lors d'entraînements en contexte long ou en grands lots, où les activations d'une seule couche de décodeur peuvent dépasser 2 Go.
Ainsi, les nouveaux algorithmes d'Unsloth et le Gradient Checkpointing contribuent à la plupart des améliorations (3,2x), permettant un fine-tuning QLoRA GPT-OSS à 290k de contexte sur une seule H100.
🔓 Tiled MLP : Déverrouiller 500K+
Avec l'aide de Stas Bekman (Snowflake), nous avons intégré Tiled MLP du papier Arctic Long Sequence Training et du billet de blog de Snowflake. TiledMLP réduit la mémoire d'activation et permet des longueurs de séquence beaucoup plus longues en mosaïquant les états cachés le long de la dimension séquence avant les fortes projections MLP.
Nous introduisons également quelques améliorations de confort d'utilisation :
Nous préservons l'état RNG lors des recomputations forward mosaïquées afin que le dropout et autres opérations stochastiques soient cohérents entre les relectures forward et backward. Cela maintient les calculs emboîtés checkpointés stables et numériquement identiques.
Notre implémentation applique automatiquement un patch à tout module nommé ou typé mlp, donc presque tous les modèles avec des modules MLP sont pris en charge nativement pour Tiled MLP.
Compromis à garder à l'esprit
TiledMLP économise de la VRAM au prix de passes avant supplémentaires. Parce qu'il vit à l'intérieur d'un bloc transformeur checkpointé et est lui-même écrit dans un style checkpoint, il devient effectivement un checkpoint imbriqué : un MLP effectue maintenant ~3 passes forward et 1 pass backward par étape. En retour, nous pouvons retirer presque toutes les activations intermédiaires du MLP de la VRAM tout en supportant des séquences extrêmement longues.
Les graphiques comparent les timelines de mémoire active pour le forward et le backward d'une seule couche de décodeur pendant une étape d'entraînement en contexte long, sans Tiled MLP (gauche) et avec (droite). Sans Tiled MLP, le pic de VRAM se produit pendant le backward du MLP ; avec Tiled MLP, il se déplace vers le calcul du loss fusionné. Nous observons environ 40% de VRAM en moins, et puisque le loss fusionné se découpe automatiquement en tranches en fonction de la VRAM disponible, le pic avec Tiled MLP serait encore plus petit sur des GPU plus petits.
Pour montrer que la cross-entropie n'est pas le nouveau goulot d'étranglement, nous fixons sa taille de tranche au lieu de la choisir dynamiquement puis doublons le nombre de tranches. Cela réduit significativement les pics de mémoire liés au loss. La mémoire maximale se produit maintenant pendant le backward dans les deux cas, et le timing global est similaire, bien que Tiled MLP ajoute un petit surcoût : un grand GEMM devient de nombreux matmuls séquentiels, plus la passe forward supplémentaire mentionnée ci-dessus.
Dans l'ensemble, le compromis en vaut la peine : sans Tiled MLP, l'entraînement en contexte long peut nécessiter environ 2× l'utilisation mémoire, tandis qu'avec Tiled MLP un seul GPU ne subit qu'environ 1,3× d'augmentation du temps par étape pour la même longueur de contexte.
Activation de Tiled MLP dans Unsloth :
Il suffit de définir unsloth_tiled_mlp = True dans from_pretrained et Tiled MLP est activé. Nous suivons la même logique que le papier Arctic et choisissons num_shards = ceil(seq_len/hidden_size). Chaque tuile opérera sur des longueurs de séquence de la même taille que la dimension cachée du modèle afin d'équilibrer le débit et les économies de mémoire.
Nous avons également discuté de la façon dont Tiled MLP effectue effectivement 3 passes forward et 1 backward, contre 2 passes forward et 1 backward avec le checkpointing de gradient normal avec Stas Bekman et DeepSpeed a fourni une mise à jour de documentation pour Tiled MLP au sein de DeepSpeed.
La prochaine fois que le fine-tuning manquera de mémoire, essayez d'activer unsloth_tiled_mlp = True. Cela devrait économiser de la VRAM tant que la longueur de contexte est supérieure à la dimension cachée du LLM.
Avec notre dernière mise à jour, il est désormais possible d'atteindre 1M de longueur de contexte avec un modèle plus petit sur un seul GPU !
Essayez le fine-tuning gpt-oss-20b à 500K de contexte sur notre notebook Colab 80GB A100.
Si vous êtes arrivé jusqu'ici, nous publions cette semaine un nouveau blog sur nos dernières améliorations de la vitesse d'entraînement, restez donc à l'écoute en rejoignant notre Reddit r/unsloth ou notre documentation.
Mis à jour
Ce contenu vous a-t-il été utile ?