Mixture of Experts (MoE) : Une Approche Révolutionnaire en PyTorch

Au cours de l’année écoulée, les modèles Mixture of Experts (MoE) ont connu une montée en popularité, alimentée par des modèles open source puissants tels que , , , et bien d’autres. Chez Databricks, nous avons travaillé en étroite collaboration avec l’équipe PyTorch pour mettre à l’échelle l’entraînement des modèles MoE. Dans ce billet de blog, nous parlerons de notre capacité à mettre à l’échelle jusqu’à plus de trois mille GPU en utilisant MegaBlocks, une implémentation efficace et open source de MoE en PyTorch.

Qu’est-ce qu’un MoE ?

Un modèle MoE est une architecture qui utilise plusieurs réseaux d’experts pour faire des prédictions. Un réseau de routage est utilisé pour acheminer et combiner les sorties des experts, garantissant que chaque expert est formé sur une distribution de tokens différente et spécialisée. L’architecture d’un grand modèle de langage basé sur un transformateur se compose généralement d’une couche d’incorporation qui mène à plusieurs blocs de transformateurs. Chaque bloc de transformateur contient un bloc d’attention et un réseau de neurones à propagation avant. Ces blocs de transformateurs sont empilés de telle sorte que la sortie d’un bloc de transformateur mène à l’entrée du bloc suivant. La sortie finale passe par une couche entièrement connectée et un softmax pour obtenir les probabilités du prochain token à produire.

Lors de l’utilisation d’un MoE dans les LLM, la couche de réseau de neurones à propagation avant est remplacée par une couche MoE qui se compose d’un réseau de routage et d’un certain nombre d’experts. Le réseau de routage, généralement un réseau de neurones à propagation linéaire, prend chaque token en entrée et produit un ensemble de poids qui déterminent quels tokens sont acheminés vers quels experts. Les experts eux-mêmes sont généralement implémentés sous forme de réseau de neurones à propagation avant également. Pendant l’entraînement, le réseau de routage s’adapte pour attribuer les entrées aux experts, permettant au modèle de se spécialiser et d’améliorer ses performances. Les sorties du routeur sont ensuite utilisées pour pondérer les sorties des experts afin d’obtenir la sortie finale de la couche MoE.

Comparés aux modèles denses, les MoEs offrent un entraînement plus efficace pour un budget de calcul donné. Cela est dû au fait que le réseau de routage envoie uniquement les tokens à un sous-ensemble d’experts, réduisant la charge de calcul. En conséquence, la capacité d’un modèle (son nombre total de paramètres) peut être augmentée sans augmenter proportionnellement les exigences de calcul. Lors de l’inférence, seuls certains experts sont utilisés, de sorte qu’un MoE peut effectuer une inférence plus rapide qu’un modèle dense. Cependant, l’ensemble du modèle doit être chargé en mémoire, pas seulement les experts utilisés.

La rareté des MoEs qui permet une plus grande efficacité de calcul provient du fait qu’un token particulier ne sera acheminé que vers un sous-ensemble d’experts. Le nombre d’experts et la manière dont les experts sont choisis dépendent de la mise en œuvre du réseau de routage, mais une méthode courante est le top k. Le réseau de routage prédit d’abord une valeur de probabilité pour chaque expert, puis achemine le token vers les k meilleurs experts pour obtenir la sortie. Cependant, si tous les tokens vont toujours vers le même sous-ensemble d’experts, l’entraînement devient inefficace et les autres experts finissent par être sous-entraînés. Pour pallier ce problème, une perte d’équilibrage de charge est introduite pour encourager un acheminement équitable vers tous les experts.

Le nombre d’experts et le choix des meilleurs k experts sont des facteurs importants dans la conception des MoEs. Un nombre plus élevé d’experts permet de passer à des modèles plus grands sans augmenter les coûts de calcul. Cela signifie que le modèle a une plus grande capacité d’apprentissage, cependant, au-delà d’un certain point, les gains de performance ont tendance à diminuer. Le nombre d’experts choisis doit être équilibré avec les coûts d’inférence de service du modèle, car l’ensemble du modèle doit être chargé en mémoire. De même, lors du choix du top k, un top k plus bas pendant l’entraînement entraîne des multiplications de matrices plus petites, laissant une computation libre sur la table si les coûts de communication sont suffisamment élevés. Lors de l’inférence, cependant, un top k plus élevé conduit généralement à une vitesse d’inférence plus lente.

MegaBlocks

MegaBlocks est une implémentation efficace de MoE qui utilise une multiplication de matrices clairsemées pour calculer les sorties des experts en parallèle malgré une affectation inégale des tokens. MegaBlocks implémente un MoE sans suppression de tokens tout en utilisant des noyaux GPU qui maintiennent un entraînement efficace. Avant MegaBlocks, les formulations de routage dynamique imposaient un compromis entre la qualité du modèle et l’efficacité matérielle. Auparavant, les utilisateurs devaient soit supprimer des tokens du calcul, soit gaspiller du calcul et de la mémoire sur le rembourrage. Les experts peuvent recevoir un nombre variable de tokens et le calcul de l’expert peut être effectué efficacement en utilisant une multiplication de matrices clairsemées par blocs. Nous avons travaillé pour permettre la mise à l’échelle de l’entraînement MoE sur des milliers de GPU.

Parallélisme des Experts

À mesure que les modèles augmentent en taille et ne rentrent plus sur un seul GPU, nous avons besoin de formes plus avancées de parallélisme. Le parallélisme des experts est une forme de parallélisme de modèle où nous plaçons différents experts sur différents GPU pour de meilleures performances. Au lieu que les poids des experts soient communiqués à tous les GPU, les tokens sont envoyés à l’appareil contenant l’expert. En déplaçant les données au lieu des poids, nous pouvons agréger les données sur plusieurs machines pour un seul expert. Le routeur détermine quels tokens de la séquence d’entrée doivent être envoyés à quels experts. Cela se fait généralement en calculant un score de routage pour chaque paire token-expert, puis en acheminant chaque token vers les experts les mieux classés. Une fois les affectations de token à expert déterminées, une étape de communication all-to-all est effectuée pour envoyer les tokens aux appareils hébergeant les experts pertinents. Cela implique que chaque appareil envoie les tokens assignés aux experts sur d’autres appareils, tout en recevant les tokens assignés à ses experts locaux.

L’avantage clé du parallélisme des experts est le traitement de quelques multiplications de matrices plus grandes au lieu de plusieurs petites multiplications de matrices. Comme chaque GPU n’a qu’un sous-ensemble d’experts, il ne doit effectuer de calcul que pour ces experts. De même, en agrégeant les tokens sur plusieurs GPU, la taille de chaque matrice est proportionnellement plus grande. Comme les GPU sont optimisés pour les calculs parallèles à grande échelle, les opérations plus importantes peuvent mieux exploiter leurs capacités, ce qui conduit à une meilleure utilisation et efficacité. Une explication plus approfondie des avantages des multiplications de matrices plus grandes peut être trouvée . Une fois le calcul terminé, une autre étape de communication all-to-all est effectuée pour renvoyer les sorties des experts à leurs appareils d’origine.

Nous exploitons , une abstraction de bas niveau pour décrire comment les tenseurs sont fragmentés et répliqués, pour mettre en œuvre efficacement le parallélisme des experts. Nous plaçons d’abord manuellement les experts sur différents GPU, en général en fragmentant sur un nœud pour nous assurer que nous pouvons exploiter NVLink pour une communication GPU rapide lorsque nous acheminons les tokens. Nous pouvons ensuite construire un sur cette disposition, ce qui nous permet de décrire de manière succincte le parallélisme sur l’ensemble du cluster. Nous pouvons utiliser cette grille d’appareils pour facilement vérifier ou réorganiser les experts lorsque nous avons besoin de formes de parallélisme alternatives.

Mise à l’échelle de ZeRO-3 avec PyTorch FSDP

En conjonction avec le parallélisme des experts, nous utilisons le parallélisme des données pour toutes les autres couches, où chaque GPU stocke une copie du modèle et de l’optimiseur et traite un morceau de données différent. Après que chaque GPU ait terminé une passe avant et arrière, les gradients sont accumulés sur les GPU pour une mise à jour globale du modèle.

ZeRO-3 est une forme de parallélisme des données où les poids et les optimiseurs sont fragmentés sur chaque GPU au lieu d’être répliqués. Chaque GPU ne stocke désormais qu’un sous-ensemble du modèle complet, réduisant considérablement la pression sur la mémoire. Lorsqu’une partie du modèle est nécessaire pour le calcul, elle est rassemblée sur tous les GPU, et après que le calcul est terminé, les poids rassemblés sont jetés. Nous utilisons l’implémentation de ZeRO-3 de PyTorch, appelée .

À mesure que nous passons à des milliers de GPU, le coût de la communication entre les appareils augmente, ralentissant l’entraînement. La communication augmente en raison de la nécessité de synchroniser et de partager les paramètres du modèle, les gradients et les états de l’optimiseur sur tous les GPU, ce qui implique des opérations de collecte et de réduction. Pour atténuer ce problème tout en conservant les avantages de FSDP, nous utilisons le parallélisme de données sharded hybride (HSDP) pour fragmenter le modèle et l’optimiseur sur un certain nombre de GPU et répliquer cela plusieurs fois pour utiliser pleinement le cluster. Avec HSDP, une opération de réduction supplémentaire est nécessaire dans la passe arrière pour synchroniser les gradients entre les répliques. Cette approche nous permet d’équilibrer l’efficacité mémoire et le coût de communication lors de l’entraînement distribué à grande échelle. Pour utiliser HSDP, nous pouvons étendre notre grille d’appareils précédente du parallélisme des experts et laisser PyTorch faire le gros du travail pour fragmenter et rassembler lorsque cela est nécessaire.

Avec PyTorch, nous pouvons combiner efficacement ces deux types de parallélisme, en exploitant l’API de plus haut niveau de FSDP tout en utilisant l’abstraction de plus bas niveau lorsque nous voulons implémenter quelque chose de personnalisé comme le parallélisme des experts. Nous disposons maintenant d’une grille d’appareils 3D avec une dimension de fragment de parallélisme des experts, une dimension de fragment ZeRO-3 et une dimension de réplication pour le parallélisme des données pur. Ensemble, ces techniques offrent une mise à l’échelle presque linéaire sur de très grands clusters, nous permettant d’atteindre des chiffres MFU de plus de 40 %.

Checkpointing Élastique avec Torch Distributed

La tolérance aux pannes est cruciale pour garantir que les LLM peuvent être entraînés de manière fiable sur des périodes prolongées, en particulier dans des environnements distribués où les défaillances de nœuds sont courantes. Pour éviter de perdre de la progression lorsque les travaux rencontrent inévitablement des défaillances, nous mettons en place un système de checkpointing de l’état du modèle, qui comprend les paramètres, les états de l’optimiseur et d’autres métadonnées nécessaires. Lorsqu’une défaillance se produit, le système peut reprendre à partir du dernier état enregistré plutôt que de recommencer. Pour garantir la robustesse aux défaillances, nous devons effectuer fréquemment des checkpoints et enregistrer et charger les checkpoints de la manière la plus performante possible pour minimiser les temps d’arrêt. De plus, si trop de GPU échouent, la taille de notre cluster peut changer. En conséquence, nous avons besoin de la capacité de reprendre de manière élastique sur un nombre différent de GPU.

PyTorch prend en charge le checkpointing élastique grâce à son framework de formation distribuée, qui comprend des utilitaires pour sauvegarder et charger des checkpoints dans différentes configurations de cluster. PyTorch Distributed Checkpoint garantit que l’état du modèle peut être sauvegardé et restauré avec précision sur tous les nœuds du cluster d’entraînement en parallèle, indépendamment de tout changement dans la composition du cluster en raison de défaillances ou d’ajouts de nœuds.

De plus, lors de l’entraînement de modèles très volumineux, la taille des checkpoints peut être très grande, ce qui entraîne des temps de téléchargement et de téléchargement de checkpoint très lents. PyTorch Distributed Checkpoint prend en charge les checkpoints fragmentés, ce qui permet à chaque GPU de sauvegarder et de charger uniquement sa partie du modèle. En combinant le checkpointing fragmenté avec l’entraînement élastique, chaque GPU lit le fichier de métadonnées pour déterminer quels fragments télécharger lors de la reprise. Le fichier de métadonnées contient des informations sur les parties de chaque tenseur stockées dans chaque fragment. Le GPU peut ensuite télécharger les fragments pour sa partie du modèle et charger cette partie du checkpoint.

En parallélisant le checkpointing sur les GPU, nous pouvons répartir la charge du réseau, améliorant la robustesse et la vitesse. Lors de l’entraînement d’un modèle avec plus de 3000 GPU, la bande passante du réseau devient rapidement un goulot d’étranglement. Nous tirons parti de la réplication dans HSDP pour d’abord télécharger les checkpoints sur une réplique, puis envoyer les fragments nécessaires aux autres répliques. Avec notre intégration dans , nous pouvons télécharger de manière fiable les checkpoints vers le stockage cloud aussi fréquemment que toutes les 30 minutes et reprendre automatiquement à partir du dernier checkpoint en cas de défaillance d’un nœud en moins de 5 minutes.

Conclusion

Nous sommes très enthousiastes de voir comment PyTorch permet de former des LLM de pointe avec de grandes performances. Dans notre billet, nous avons montré comment nous avons mis en œuvre un entraînement efficace de MoE grâce à Pytorch Distributed et MegaBlocks sur Foundry. De plus, le checkpointing élastique de Pytorch nous a permis de reprendre rapidement l’entraînement sur un nombre différent de GPU en cas de défaillance de nœud. L’utilisation de Pytorch HSDP nous a permis de mettre à l’échelle efficacement l’entraînement ainsi que d’améliorer les temps de reprise du checkpointing. Nous sommes impatients de continuer à construire une communauté open source solide et dynamique pour aider à apporter de grands modèles d’IA à tout le monde. Rejoignez-nous pour construire de grands modèles sur et .

Source de l’article