Réduire le temps d'entraînement de Qwen de 6x sur AI Runtime

L'entraînement ou le peaufinage (fine-tuning) d'un modèle fondation de plusieurs milliards de paramètres nécessite une puissance de calcul importante, s'étalant souvent sur des centaines, voire des milliers d'heures d'entraînement sur GPU. Mais l'optimisation des performances d'entraînement pour un LLM est complexe, et les métriques de haut niveau, comme l'utilisation des GPU, peuvent ne pas refléter votre utilisation réelle. Dans cet article de blog, nous vous montrons comment franchir les couches d'abstraction et multiplier par 6 la vitesse de fine-tuning du dernier modèle Qwen 3.8 27B sur AI Runtime. Grâce à une optimisation systématique des performances, alors même que l'utilisation du GPU était déjà à 98 %, nous sommes passés de 7 heures 44 minutes à 1 heure 17 minutes pour le fine-tuning d'un jeu de données fixe.

Nous détaillerons en profondeur les techniques qui sous-entendent ces améliorations, depuis la mesure des éléments clés jusqu'aux détails d'implémentation et aux compromis de conception propres à chaque optimisation.

Si vous souhaitez commencer immédiatement, accédez directement au code d'exemple sur AI Runtime et essayez-le !

Configuration du modèle

Nous commençons par une boucle d'entraînement standard pour Qwen 3.8 27B en utilisant la bibliothèque transformers et le jeu de données d'entraînement SFT HuggingFaceH4/ultrachat_200k.

• Nous appliquons FSDP2, la toute dernière technique de sharding (fractionnement) de modèle de torch. Cela est nécessaire pour faire tenir les poids du modèle et l'état de l'optimiseur sur les GPU. Une approche purement parallèle distribuée aux données échouerait à faire tenir les poids d'un modèle de cette taille sur un seul GPU.
• Nous activons également le checkpointing d'activation, une technique qui échange de l'espace mémoire contre du calcul lors de la passe arrière (backwards pass), afin de permettre une taille de batch par GPU plus importante.
• Le dataloader regroupe des batchs 2D avec un nombre maximal de tokens par batch de 8192, avec un remplissage (padding) jusqu'à la taille maximale de l'exemple du batch.

Benchmarks d'entraînement initiaux

Notre entraînement de référence (baseline) du modèle sur l'ensemble du jeu de données prend 7 heures et 44 minutes et, pendant l'entraînement, le GPU affiche une utilisation moyenne de 98 %. Le résultat est un entraînement à 9 396 tok/s.

Cette utilisation du GPU semble élevée, mais comme mentionné précédemment, ce chiffre est trompeur. La métrique standard d'utilisation du GPU compte uniquement le pourcentage de temps pendant lequel le GPU faisait quelque chose, ce qui ne nous indique pas si le modèle utilisait bien le GPU lorsqu'il était actif. Pour comprendre cela, nous devons analyser les choses plus en détail.

L'un des moyens d'y parvenir est le mécanisme de traçage de torch ; nous pouvons ajuster notre boucle d'entraînement comme suit (tout en limitant le nombre d'étapes pour obtenir une réponse rapide) :

Pour un exemple d'étape, nous obtenons les données suivantes :

Tableau 1 : CPU - 14,4 s avec 1 714 382 appels. Le tableau contient les principales opérations par temps total.

Tableau 2 : GPU - 10,3 s avec 269 448 opérations. Le tableau contient les principales opérations par temps total.

Axes d'amélioration

À partir de ces tableaux, nous constatons quelques problèmes majeurs avec le modèle actuel :

• Le CPU prend plus de temps pour configurer chaque étape d'entraînement que le GPU n'en prend pour l'exécuter. C'est inattendu car, en général, le côté CPU de l'exécution sur GPU est asynchrone, c'est-à-dire que le CPU devrait s'exécuter bien en avance sur le GPU. La principale conséquence est le blocage du GPU, ce qui nous donne une opportunité d'amélioration.
• Il y a trop d'opérations ! Dans les tableaux ci-dessus, nous avons des centaines de milliers de petites opérations. Ici, le problème est que le CPU ne peut pas mettre en file d'attente les opérations GPU suffisamment vite pour maintenir le GPU occupé. L'idée clé est que nous devons réduire le nombre d'opérations que le CPU met en file d'attente pour le GPU. Il existe des techniques générales pour cela, dont une adaptée à notre situation :
  • Les graphes CUDA de torch enregistrent les kernels CUDA détectés et les rejouent avec un seul appel CPU. Bien que puissante, cette approche présente de nombreuses limitations et entre en conflit avec le dynamisme de FSDP2. Elle est donc écartée pour le moment.
  • Torch.compile trace le code et utilise Triton pour fusionner automatiquement les kernels lorsque c'est possible. Dans cet exemple, nous utiliserons des kernels fusionnés manuellement ; dans le cas des LLM, torch.compile ne laisse aucune marge d'amélioration supplémentaire.
  • Les kernels fusionnés manuellement utilisent des techniques de fusion spécialisées pour des opérations particulières que le compilateur de torch ne peut pas appliquer lui-même. Heureusement, il existe de nombreux kernels préconçus bien établis dans la communauté. C'est le cas par exemple de « Flash Attention », qui offre ici une légère amélioration des performances.

De plus, nous avons un problème de remplissage (padding) excessif dans le batch : le tenseur d'entrée doit être rectangulaire et chaque ligne est complétée jusqu'au nombre maximal de tokens pour ce batch. Nous allons montrer une astuce adaptée aux LLM permettant de le minimiser.

Dans les sections suivantes, nous examinerons plus en détail chaque sujet technique, en couvrant sa définition, les meilleures pratiques recommandées et les principaux compromis à prendre en compte lors de son application.

Kernels fusionnés

Vous remarquerez peut-être que certains kernels apparaissent très souvent ! En particulier, pour Qwen, le kernel de multiplication matricielle par batch est exécuté plus de 30 000 fois en une seule étape.

Dans ce cas, nous allons d'abord explorer les kernels fusionnés. Un kernel fusionné prend plusieurs kernels GPU plus petits et les combine en une seule unité. Cela permet d'améliorer les performances en évitant les transferts coûteux entre la mémoire à haute bande passante du GPU et sa mémoire intégrée (on-die) à bande passante encore plus élevée. Cela permet également une planification plus efficace des cœurs GPU.

Avec autant de noyaux supplémentaires, même si torch.compile peut fonctionner, l'exécution prendra beaucoup de temps. Réduire le nombre brut d'opérations avec des noyaux fusionnés manuellement constitue donc une bonne première étape. De plus, des fusions bien connues telles que Flash Attention ne sont pas le genre de choses que torch.compile trouverait de lui-même.

La bibliothèque transformers utilise par défaut la fonction SDPA de torch pour les couches d'attention standard. Cependant, il s'avère que beaucoup de ces petites opérations proviennent des couches non standard de réseaux gated delta. Heureusement, pour ces couches, il existe deux noyaux, causal-conv1d et flash-linear-attention, qui ont été écrits pour fusionner plusieurs de ces opérations par couche. Comment le découvrir ? L'implémentation de la bibliothèque transformers de Qwen nous indique déjà quels noyaux elle prend en charge :

• flash-linear-attention
• causal-conv1d

Il existe deux méthodes principales pour obtenir ces noyaux :

• Installer les noyaux vous-même. Cela peut aller de :
  • Un simple pip install (par ex. pour flash-linear-attention) depuis PyPI
  • Un clone git et la compilation du dépôt concerné
  • Trouver une autre source pour un package précompilé au cas où la compilation prendrait trop de temps
• Utiliser le Kernel Hub de Hugging Face pour récupérer des noyaux précompilés.

Nous recommandons généralement la première option, en particulier avec notre prise en charge des conteneurs Docker : construisez votre conteneur avec tous vos noyaux installés au moment du build afin de réduire le temps de démarrage de votre tâche.

Bien que Kernel Hub soit pratique, comme les noyaux dépendent souvent explicitement des versions de torch et de CUDA installées, un noyau compatible avec votre configuration spécifique peut ne pas être disponible. Le modèle transformers utilisera alors silencieusement un noyau non optimisé offrant de faibles performances.

Après avoir installé les packages Python flash-linear-attention et causal-conv1d, nous avons amélioré notre temps global d'un facteur 3,7x, le réduisant à 2 heures et 5 minutes. Dans les tableaux ci-dessous, vous pouvez voir les nouveaux noyaux (en surbrillance) remplacer les fonctions élément par élément des tableaux précédents. De plus, le nombre total d'opérations exécutées a été considérablement réduit, d'un facteur 15x :

Tableau 3 : CPU - 3,9 s avec 118 316 appels. Le tableau présente les principales opérations par temps total.

Tableau 4 : GPU - 3,9 s avec 18 060 appels. Le tableau présente les principales opérations par temps total.

Avertissement : pendant l'entraînement, vous aurez besoin d'un noyau pour les passes avant (forward) et arrière (backward), mais certains noyaux sont uniquement réservés à la passe avant, car ils n'ont été implémentés que pour accélérer l'inférence.

À retenir

• Pensez à utiliser des noyaux bien établis et hautement optimisés qui sont facilement disponibles pour les modèles courants. Tirer parti de ces noyaux peut apporter des améliorations significatives de performances avec relativement peu d'efforts. Dans ce cas précis, nous avons constaté une amélioration de 3,7x de la vitesse d'entraînement, réduisant notre temps d'entraînement de 7 heures 44 minutes à 2 heures 5 minutes.

Mise en lots, tokenisation et padding

Notre modèle Qwen est désormais assez rapide, mais nous pouvons aller encore plus vite. Par défaut, notre configuration de mise en lots utilise des tenseurs rectangulaires 2D de 8 192 tokens, car c'est la plus grande taille que nous puissions utiliser sans manquer de mémoire. L'inconvénient de cette approche est que nous avons tendance à gaspiller une quantité considérable d'espace dans le tenseur avec du padding.

Si Qwen n'utilisait que des mécanismes d'attention standard, la fonctionnalité NestedTensor de torch est maintenant bien prise en charge et permettrait de réduire ce coût en supprimant complètement le padding. Cependant, flash-linear-attention ne prend en charge qu'une approche plus traditionnelle de tenseur compacté 1D :

Mais nous ne pouvons pas faire cela de manière arbitraire : transmettre simplement le tenseur 1D au modèle peut être acceptable pour le préentraînement, mais pour le réglage fin (finetuning), nous devons nous assurer que les tokens d'un exemple ne prêtent attention qu'aux tokens du même exemple. Pour résoudre ce problème, nous devons transmettre les bons position_ids d'exemple au modèle. À son tour, le modèle les transmettra aux noyaux d'attention fusionnés, qui appliqueront le masque d'attention correct. Si vous consultez le code d'exemple ici, vous verrez comment nous calculons les entrées nécessaires pour que cela fonctionne.

Grâce à ce changement en apparence simple, notre temps d'entraînement passe de 2 heures 5 minutes à 1 heure 17 minutes.

Il y a deux raisons à cela :

1. La raison principale est que ce compactage augmente le nombre de séquences par lot et réduit le nombre total de lots d'environ 36 %. Notez que cela dépend de la distribution et de la longueur des exemples dans votre jeu de données.
2. Une raison secondaire est que les noyaux flash linear attention doivent être compilés à la volée (just-in-time) pour chaque forme d'entrée, et la forme était auparavant dynamique, ce qui entraînait quelques minutes supplémentaires de compilation.

Compromis : ce type de modification rend le pipeline d'entrée plus complexe et augmente le risque d'introduire des erreurs qui affectent la qualité du modèle. Mais dans ce cas, l'économie de 36 % en vaut la peine. Il existe d'autres techniques pour réduire le padding, comme conserver le mécanisme de mise en lots d'origine mais trier toutes les entrées par longueur pour minimiser le padding, mais cet ordre peut introduire une corrélation inattendue entre les échantillons adjacents et biaiser le modèle.

À retenir : Lors de l'entraînement de LLM sur un jeu de données contenant des entrées plus courtes que votre nombre maximal de tokens entraînabiles (par exemple 8 192 dans ce cas), le compactage efficace des données en lots permet de gagner un temps considérable. Pour ce jeu de données, nous avons constaté une amélioration de 1,6x de la vitesse d'entraînement, ce qui a réduit notre temps d'entraînement de 2 heures 5 minutes à 1 heure 17 minutes.

Conclusion

Dans ce blog, nous vous avons montré comment multiplier par 6 l'efficacité de votre entraînement. Cela s'est résumé à : (1) mesurer d'abord ce qui compte, pas seulement l'utilisation des GPU, mais aussi les étapes où votre modèle passe du temps lors de l'entraînement, (2) optimiser et fusionner les noyaux à l'aide de bibliothèques connues comme Flash Linear Attention, et (3) adapter la manière dont les lots de données sont transmis au modèle.

Pour faire cela à grande échelle, vous avez besoin d'une plateforme qui propose par défaut une trajectoire performante. Dans AIR, bon nombre de ces améliorations de performances sont intégrées. Nous incluons les dernières versions de CUDA et de flash attention dans nos environnements de base, suivons automatiquement pour vous des métriques clés telles que l'utilisation des GPU, la température et leur état de santé, et veillons à ce que les versions de torch (ainsi que les profilers) soient à jour. Notre agent interne, Genie Code, est également réglé pour rechercher ce type d'opportunités afin d'optimiser davantage votre entraînement. De plus, la lecture de données depuis Databricks est optimisée grâce à notre UCVolumeDataset.

N'hésitez pas à essayer AIR pour votre prochaine charge de travail d'entraînement ! Prêt à tenter l'expérience ? Consultez le cookbook AI Runtime pour découvrir comment exécuter vous-même ce modèle.

(Cet article de blog a été traduit à l'aide d'outils basés sur l'intelligence artificielle) Article original