ZeRO++ : Faire avancer l'entraînement des gros modèles

L'apprentissage profond est devenu un outil majeur dans plein de domaines, comme l'analyse d'images, le traitement du langage et la reconnaissance vocale. Avec le temps, ces modèles se sont améliorés en termes de qualité et de capacité. Une des principales raisons de cette amélioration, c'est l'augmentation de la taille. Les modèles plus grands ont tendance à mieux performer parce qu'ils peuvent capturer des motifs plus complexes. Par exemple, la taille des modèles est passée de 100 millions de paramètres à plus de 500 milliards de paramètres en juste quelques années.

Cependant, entraîner ces modèles massifs pose des défis importants. Les besoins en mémoire et en calcul ont explosé, rendant impossible pour un seul GPU de gérer la charge de travail. Pour entraîner de grands modèles efficacement, on doit combiner la puissance de plusieurs GPUs. Deux méthodes courantes pour y arriver sont le Parallélisme 3D et l'Optimiseur de Redondance Zéro (ZeRO).

Le parallélisme 3D combine plusieurs techniques pour répartir les charges de travail sur de nombreux GPUs. Bien que ça permette au système d'utiliser les ressources efficacement, ça complique les choses. Les data scientists rencontrent souvent des défis pour adapter leur code au parallélisme 3D.

D'un autre côté, ZeRO simplifie le processus. Au lieu de répliquer les états du modèle sur tous les GPUs, il divise les paramètres du modèle en plus petites parties qui sont partagées entre les GPUs. Cette approche utilise moins de mémoire et réduit les coûts de communication tout en offrant de bonnes performances.

Malgré ses avantages, ZeRO a ses limites, surtout dans les situations où le nombre de GPUs est élevé ou lorsque la bande passante du réseau est faible. Ces facteurs peuvent causer des goulets d'étranglement en communication qui ralentissent l'entraînement.

#Limites de ZeRO

ZeRO est populaire parce qu'il est facile à utiliser et qu'il évolue efficacement. Cependant, les coûts de communication peuvent limiter l'efficacité dans deux scénarios clés : les clusters à faible bande passante et les petites tailles de lot.

Dans les clusters à faible bande passante, même si les machines performantes peuvent avoir d'excellentes connexions internes, les connexions entre différentes machines ont souvent une bande passante plus faible. Ça crée des goulets d'étranglement en communication, ce qui réduit les performances du processus d'entraînement.

Même dans des réseaux à haute bande passante, si le nombre de GPUs est très grand, la taille du lot pour chaque GPU peut devenir trop petite. Ça abaisse encore le ratio calcul-communication et crée des délais supplémentaires en communication.

Bien que ZeRO ait plein d'avantages, les méthodes existantes pour optimiser la communication dans l'entraînement distribué ne fonctionnent pas bien avec son architecture parce qu'elles se basent sur la réplication des états du modèle.

Pour résoudre ces problèmes, on a introduit un nouveau système appelé ZeRO++. Ce système réduit les coûts de communication et améliore les performances pendant l'entraînement. Ci-dessous, on décrit les techniques utilisées dans ZeRO++ pour surmonter les défis de communication auxquels ZeRO fait face.

#Optimisations de Communication dans ZeRO++

ZeRO++ se concentre sur trois principales manières de réduire les frais de communication :

#1. Communication de Poids Quantifiés par Bloc

Pendant le passage avant, ZeRO collecte tous les paramètres nécessaires à l'entraînement à travers une opération d'addition globale. Ça peut créer beaucoup de trafic de communication. Pour réduire ça, on utilise une technique appelée quantification basée sur des blocs. Ça rétrécit la taille des données envoyées. Au lieu d'envoyer les données au format haute précision, on les convertit en un format de précision inférieure. Ça divise le volume de communication par deux sans perdre beaucoup de précision.

#2. Partitionnement Hiérarchique des Poids

Lors du passage arrière, ZeRO doit recueillir les poids à nouveau, ce qui peut créer des frais de communication importants. Pour y remédier, on garde une copie complète des poids sur chaque machine. Ça nous permet de faire des opérations de collecte au sein d'une seule machine plutôt que d'entre plusieurs machines, rendant le processus beaucoup plus rapide grâce à une meilleure bande passante.

#3. Communication de Gradient Quantifié

Pour réduire les coûts de communication pendant le passage arrière, on doit aussi gérer les gradients. Utiliser des techniques standard pour réduire la communication peut nuire à la précision de l'entraînement. Donc, on a développé une nouvelle approche qui nous permet de communiquer les gradients en précision inférieure tout en maintenant la précision. On compresse d'abord les gradients puis on les restaure à leur précision originale avant de faire les calculs nécessaires.

En combinant ces trois techniques, ZeRO++ réduit significativement le volume de communication et améliore l'efficacité pendant l'entraînement.

#Évaluation de ZeRO++

On a implémenté ZeRO++ et réalisé des évaluations approfondies démontrant trois résultats clés :

1. Scalabilité pour de grands modèles sur des centaines de GPUs.
2. Accélération constante par rapport à l'original ZeRO.
3. Amélioration des performances même dans des paramètres à faible bande passante.

Grâce à nos tests, on a découvert que ZeRO++ atteint jusqu'à 2,16 fois meilleur débit avec 384 GPUs. On a aussi analysé comment chacune des trois optimisations contribue à l'amélioration générale et montré que ZeRO++ n'impacte pas négativement la convergence du modèle.

#Contributions Clés

Les principales contributions de ce travail incluent :

• Les poids quantifiés par blocs réduisent le volume de communication pour la collecte de poids de 50%.
• Le partitionnement hiérarchique des poids élimine complètement la communication inter-nœud pendant les passages arrière.
• Une nouvelle méthode de communication de gradient quantifié réduit le volume de communication pour les gradients de 75%.
• L'intégration efficace de ces techniques dans l'architecture existante de ZeRO, permettant de meilleures performances.

#Contexte et Travaux Connus

Pour mieux comprendre ZeRO++, il est essentiel de regarder comment les grands modèles sont entraînés sur plusieurs GPUs. Il y a principalement trois formes de parallélisme dans l'entraînement distribué : le parallélisme de données, le parallélisme par pipeline et le parallélisme par tenseur.

Le parallélisme de données fonctionne bien quand le modèle peut tenir dans un seul GPU. Dans cette configuration, chaque GPU a une copie complète du modèle et traite différentes données d'entrée. Quand les modèles ne peuvent pas tenir dans un GPU, le parallélisme de modèle entre en jeu. Cette méthode divise le modèle entre les GPUs, soit par couches (parallélisme par pipeline) soit à travers les dimensions du tenseur (parallélisme par tenseur).

Le parallélisme 3D combine ces méthodes mais nécessite des changements importants dans le code existant. En revanche, ZeRO offre une alternative simple en partitionnant les paramètres du modèle entre les GPUs, ce qui minimise la communication tout en maximisant l'efficacité de la mémoire.

Cependant, la méthode ZeRO standard a des demandes de communication élevées, surtout dans des environnements à grande échelle. Les solutions existantes se sont concentrées sur la réduction des coûts de communication, mais elles ne s'intègrent souvent pas bien avec la structure de ZeRO.

#Comment ZeRO++ Traite les Coûts de Communication

ZeRO++ combine différentes stratégies pour améliorer l'efficacité de la communication dans ZeRO.

#Communication de Poids Quantifiés

Pendant le passage avant, ZeRO s'appuie sur la collecte de poids, ce qui peut être lourd en communication. En appliquant la quantification basée sur des blocs, on compresse les données envoyées. Cette technique ajuste les poids à une précision inférieure, réduisant la taille tout en maintenant une précision acceptable.

#Partitionnement Hiérarchique

Pour le passage arrière, on évite la collecte inter-nœud en ayant des poids complets sur chaque machine. Cette approche nous permet de faire la collecte localement dans chaque machine, ce qui est beaucoup plus rapide que la communication entre nœuds.

#Communication de Gradient Quantifié

Pour les gradients dans le passage arrière, on a créé une nouvelle méthode qui compresse les gradients tout en maintenant leur précision. On applique d'abord la quantification pendant la phase de communication et on restaure ensuite la précision pour les calculs.

Toutes ces méthodes s'attaquent à différents points du processus de communication, rendant plus facile le partage des données entre les GPUs tout en maintenant de hautes performances.

#Évaluation des Performances de ZeRO++

On a évalué ZeRO++ en utilisant divers modèles et configurations, en se concentrant sur sa scalabilité, sa vitesse et son efficacité globale. Notre implémentation a montré de bons résultats sur différentes plateformes et tailles de modèles.

#Scalabilité

ZeRO++ a bien fonctionné sur jusqu'à 384 GPUs avec de grands modèles. Il a spécifiquement atteint un pourcentage significatif des performances maximales disponibles avec le matériel utilisé. Ça montre que les nouvelles méthodes de communication permettent d'entraîner efficacement de plus grands modèles.

#Amélioration de la Vitesse

Une des découvertes clés était que ZeRO++ a surpassé l'original ZeRO sur divers tailles et configurations de modèles. Dans certains cas, on a observé des améliorations de performances allant jusqu'à 2,4 fois par rapport à la référence. C'est particulièrement important pour les scénarios avec une bande passante limitée, où les méthodes précédentes peinaient à maintenir l'efficacité.

#Performance en Faible Bande Passante

De plus, ZeRO++ a maintenu de solides performances lors des tests dans des réseaux à faible bande passante, prouvant sa robustesse. Ça ouvre des opportunités pour entraîner de grands modèles dans des environnements qui n'ont pas forcément des configurations réseau optimales.

#Conclusion

L'introduction de ZeRO++ représente une avancée significative dans la communication collective pour entraîner efficacement de grands modèles. En se concentrant sur la réduction des frais généraux à travers des techniques de quantification et de partitionnement intelligentes, on a montré qu'il est possible de maintenir de hautes performances, même avec des ressources réseau limitées.

En résumé, ZeRO++ propose une approche prometteuse pour rendre l'entraînement de modèles à grande échelle plus accessible et efficace pour les chercheurs et les praticiens. Avec ces avancées, on envisage d'améliorer les capacités des grands modèles de langage et au-delà, favorisant des applications plus larges dans divers domaines de l'apprentissage profond.