All-reduce est une opération de communication collective utilisée en calcul parallèle et distribué. Elle combine les données de tous les processus participants à l'aide d'une opération spécifiée (telle que la somme, le minimum, le maximum ou la moyenne) et livre le résultat final à chaque processus. Dans le contexte de apprentissage automatique, all-reduce est le mécanisme principal pour synchroniser les gradients entre plusieurs dispositifs lors de l'entraînement distribué de réseaux de neurones et de modèles de apprentissage profond.
L'opération est définie par la norme Message Passing Interface (MPI), qui spécifie la sémantique d'all-reduce comme suit : chaque processus contribue avec un tampon de données, l'opération combine ces tampons élément par élément, et le résultat est copié vers tous les processus. Cela contraste avec une opération de réduction, qui envoie le résultat uniquement à un processus racine unique. L'opération all-reduce est essentielle pour les algorithmes qui nécessitent une vue globale des données agrégées tout en maintenant des copies locales, comme descente de gradient stochastique dans l'entraînement parallèle aux données.
Rôle dans l'entraînement distribué
Dans l'entraînement distribué parallèle aux données, chaque worker (GPU ou processeur) détient une copie du modèle et traite un sous-ensemble différent des données d'entraînement. Après avoir calculé les gradients locaux via la rétropropagation, les workers doivent moyenner leurs gradients pour mettre à jour un modèle cohérent. All-reduce accomplit cela en additionnant les gradients de tous les workers puis en divisant par le nombre de workers (si l'opération de moyenne est utilisée). Cela garantit que chaque worker dispose du gradient agrégé identique, leur permettant de mettre à jour leurs répliques locales du modèle de manière cohérente.
Sans all-reduce, les workers divergeraient, rendant le processus d'entraînement instable ou incorrect. L'opération est un goulot d'étranglement critique pour faire évoluer l'entraînement vers de nombreux dispositifs car elle nécessite une surcharge de communication significative. En conséquence, des implémentations efficaces d'all-reduce sont devenues un axe clé pour fournisseurs de cloud et fabricants de matériel.
Algorithmes et implémentations
Plusieurs algorithmes existent pour effectuer all-reduce, chacun avec des compromis différents en termes de bande passante, de latence et d'évolutivité. Les implémentations courantes incluent :
- Ring All-Reduce : Les processus sont organisés en un anneau logique. Dans la phase de réduction-dispersion, chaque processus envoie des données à son voisin, accumulant des résultats partiels. Dans la phase de collecte globale, les résultats accumulés sont circulés. Cet algorithme minimise le nombre total de messages envoyés par processus et atteint une bande passante optimale sur de nombreux systèmes, ce qui le rend populaire dans le calcul haute performance et l'entraînement de grands modèles de langage.
- All-Reduce basé sur un arbre : Utilise une topologie en arbre (par exemple, binomiale ou k-nomiale) pour combiner les données de manière hiérarchique. Il est plus efficace en latence pour de petites tailles de données mais peut avoir des coûts de bande passante plus élevés.
- Réduction/doublement récursif : Divise les données en morceaux et utilise une séquence d'échanges par paires pour combiner et redistribuer, adapté à certaines topologies de cluster.
La bibliothèque open-source Open-MPI fournit une implémentation standard d'all-reduce, tandis que des versions optimisées telles que NCCL (NVIDIA Collective Communications Library) et Gloo sont largement utilisées dans les frameworks d'apprentissage profond comme PyTorch et TensorFlow. Ces bibliothèques utilisent souvent des algorithmes basés sur un anneau par défaut pour les grands tenseurs, mais passent à des algorithmes basés sur un arbre pour les petits tenseurs afin de réduire la latence.
Accélération matérielle
Le matériel IA moderne intègre de plus en plus des moteurs de communication collective dédiés pour décharger all-reduce des cœurs de calcul principaux. Par exemple, les GPU NVIDIA disposent d'un tissu spécialisé NVLink et NVSwitch, et la bibliothèque NCCL exploite cela pour un all-reduce à haut débit. De même, AMD et Intel fournissent leurs propres bibliothèques de communication collective, telles que RCCL et OneCCL, respectivement.
AWS Trainium et d'autres puces IA personnalisées intègrent souvent des unités réseau conçues pour accélérer all-reduce directement sur l'interconnexion. Les TPU de Google Cloud utilisent une interconnexion à haute bande passante qui prend en charge un all-reduce efficace via une puce dédiée appelée Interconnect Processor (ICP). Ces optimisations matérielles sont cruciales pour passer à l'échelle de centaines ou de milliers de dispositifs, car la surcharge de communication peut sinon dominer le temps d'entraînement.
Techniques d'optimisation
Pour atténuer le coût d'all-reduce, chercheurs et ingénieurs ont développé plusieurs techniques d'optimisation :
- Compression des gradients : Des techniques comme la quantification ou la sparsification réduisent la quantité de données transférées. Par exemple, écrêtage des gradients peut être combiné avec la compression, mais des méthodes plus avancées comme la sparsification top-k nécessitent une communication supplémentaire pour les indices.
- Chevauchement avec le calcul : All-reduce peut être chevauché avec la propagation arrière en divisant les gradients en morceaux et en communiquant chaque morceau dès qu'il est prêt. Cela réduit le délai de communication visible.
- All-Reduce hiérarchique : Dans les clusters avec une topologie hiérarchique (par exemple, plusieurs serveurs chacun avec plusieurs GPU), effectuer un all-reduce local au sein d'un nœud puis un all-reduce global entre les nœuds peut réduire le trafic sur le réseau.
- Précision mixte : Accumuler les gradients en précision inférieure (par exemple, float16) avant all-reduce peut réduire de moitié le volume de communication, bien qu'il faille veiller à préserver la précision.
Ces optimisations sont essentielles pour entraîner des modèles de pointe en IA générative, qui nécessitent souvent des milliers d'accélérateurs.
Variantes et opérations connexes
All-reduce fait partie d'une famille d'opérations collectives qui inclut également la diffusion, la dispersion, la collecte et la collecte globale. Les variantes d'all-reduce incluent :
- Reduce-Scatter : Combine les données et distribue le résultat en morceaux entre les processus (chaque processus reçoit une partie du résultat total). Cela est souvent utilisé comme étape intermédiaire dans le ring all-reduce.
- All-to-All : Chaque processus envoie une donnée distincte à chaque autre processus, ce qui peut être utilisé pour des schémas de communication plus généraux mais est plus coûteux.
- Cascade All-Reduce : Une méthode pour l'all-reduce hiérarchique qui équilibre le trafic entre les nœuds, comme proposé dans certains articles de recherche.
Dans le contexte du calcul distribué, all-reduce est également utilisé dans des applications au-delà de l'entraînement, telles que l'inférence d'apprentissage automatique distribuée, les méthodes d'ensemble et les algorithmes parallèles pour le calcul scientifique.
Défis et orientations futures
À mesure que les modèles deviennent plus grands, les exigences de bande passante et de latence d'all-reduce deviennent de plus en plus difficiles. Passer à l'échelle de milliers de dispositifs nécessite une planification et un équilibrage de charge sophistiqués. Certaines approches émergentes incluent :
- All-Reduce fragmenté : Diviser le tenseur de gradient en fragments et effectuer all-reduce sur chaque fragment indépendamment tout en chevauchant la communication avec le calcul.
- All-Reduce asynchrone : Assouplir la synchronisation stricte d'all-reduce pour permettre à certains workers de continuer, bien que cela puisse entraîner des problèmes de convergence.
- Calcul dans le réseau : SHARP (Scalable Hierarchical Aggregation and Reduction Protocol) de NVIDIA et des technologies similaires déplacent les opérations de réduction dans les commutateurs réseau, réduisant considérablement le temps pour all-reduce.
La recherche se poursuit sur des algorithmes plus robustes face aux matériels hétérogènes et aux topologies réseau, en particulier dans le contexte de l'entraînement à grande échelle sur plusieurs centres de données.
Histoire et normes
Le terme « all-reduce » est originaire de la communauté du calcul parallèle. Il a été formalisé dans la norme MPI, apparue pour la première fois en 1994. Le Xerox PARC et d'autres institutions de recherche ont contribué aux premières méthodologies de calcul parallèle qui ont ensuite influencé la conception de la communication collective. Dans les années 2010, avec l'essor de l'apprentissage profond, all-reduce est devenu une primitive centrale dans les frameworks d'entraînement distribué. L'équipe de recherche de Baidu a popularisé le ring all-reduce pour TensorFlow, conduisant à son adoption généralisée dans la communauté de l'apprentissage automatique.
Aujourd'hui, all-reduce reste un sujet critique dans la recherche sur les systèmes, surtout à mesure que les tailles de modèles augmentent. Les laboratoires OpenAI et Google DeepMind, entre autres, ont publié des articles sur la mise à l'échelle de l'entraînement distribué, soulignant l'importance d'implémentations efficaces d'all-reduce. L'écosystème Hugging Face et d'autres projets open-source continuent d'améliorer les bibliothèques de communication pour prendre en charge des modèles encore plus grands.
En résumé, all-reduce est un bloc de construction fondamental pour les systèmes distribués d'intelligence artificielle. Son efficacité impacte directement le temps et le coût d'entraînement des grands modèles, ce qui en fait un domaine actif de recherche et d'innovation tant matérielle que logicielle.