Unsloth prétend accélérer 5 fois l’entraînement des LLM avec de nouveaux noyaux Triton et un auto-packing

Unsloth, la bibliothèque de fine-tuning open-source populaire parmi les praticiens de LLM, a expédié une mise à jour majeure qui, selon elle, peut offrir jusqu’à 5 fois plus de rapidité d’entraînement tout en réduisant l’utilisation de la VRAM de 30 % à 90 %. Les améliorations proviennent d’un ensemble de noyaux Triton personnalisés et d’un nouveau système de packing automatique intelligent que l’équipe dit fonctionner avec n’importe quel backend d’attention, FlashAttention 3, xFormers ou SDPA, et n’importe quelle famille de modèles.

L’ajout phare est un noyau QK RoPE fusionné. Auparavant, chaque requête et clé nécessitaient deux noyaux Triton séparés ; la nouvelle implémentation les fusionne en un seul, permettant un RoPE de longueur variable. Dans les micro-benchmarks, le noyau fusionné est 2,3 fois plus rapide sur les longueurs de contexte plus longues et 1,9 fois plus rapide sur les plus courtes. Le noyau est entièrement inplace, éliminant les opérations de clone et de transposition contiguë et réduisant davantage la pression mémoire GPU.

Unsloth a également résolu une limitation de longue date pour l’entraînement sur de longs contextes : les erreurs CUDA de dépassement de limites lors d’exécutions de 500 000 tokens causées par l’indexage int32 par défaut dans Triton et CUDA. Le correctif introduit une constante de compilation LONG_INDEXING qui permet au compilateur Triton de spécialiser les noyaux pour les contextes plus courts et plus longs sans surcharge.

« En modifiant le ratio de 50 % de séquences courtes, et en supposant que nous avons plus de séquences courtes, par exemple 20 % de séquences longues et 80 % de séquences courtes, nous obtenons un entraînement 5 fois plus rapide », a écrit l’équipe Unsloth dans la publication. « Cela signifie que l’accélération du packing dépend du nombre de lignes courtes dans votre jeu de données ; plus il y en a de courtes, plus c’est rapide. »

Le système de packing automatique (activé par défaut) regroupe automatiquement plusieurs séquences courtes en un seul tenseur unidimensionnel, conservant les métadonnées de longueur de séquence pour éviter les fuites d’attention entre les échantillons. Résultat : une amélioration rapportée du débit de 1,7 à 3 fois sur un fine-tuning de Qwen3-32B sur le jeu de données yahma/alpaca-cleaned à différentes tailles de lots, certaines configurations atteignant 5 fois ou plus. Unsloth insiste sur le fait que les pertes d’entraînement correspondent exactement à celles des exécutions sans packing.

Les résultats de benchmark publiés par l’équipe montrent que sans packing, les tokens de padding peuvent représenter près de 50 % d’un lot avec une taille de lot de 8. Avec le packing, l’utilisation de tokens valides reste proche de 100 %, quelle que soit la taille du lot. L’utilisation mémoire pour Qwen3-8B et Qwen3-32B a diminué de 60 % tandis que le débit a doublé.

La mise à niveau fonctionne pour toutes les méthodes d’entraînement, le fine-tuning complet, le pré-entraînement et le RL, et est rétrocompatible avec les anciens GPU comme Tesla T4 et RTX 2080, ainsi que les plus récents comme H100 et B200. Les utilisateurs ont seulement besoin de mettre à jour Unsloth pour bénéficier du batching sans padding par défaut ; le packing explicite peut être activé avec packing = True dans SFTConfig.

Cette publication s’appuie sur les travaux antérieurs d’Unsloth sur les noyaux RoPE de décembre 2023 et une PR communautaire de mars 2024 qui a rendu l’entraînement de bout en bout 1 à 2 % plus rapide en optimisant le noyau pour lancer des blocs pour des groupes de têtes.

Unsloth, qui s’est positionné comme une alternative légère au Trainer complet de Hugging Face pour le fine-tuning, revendique désormais le support de modèles incluant la famille Qwen3, Llama 3.1 et autres. Les notebooks de la bibliothèque sur Google Colab ont été mis à jour automatiquement.