Le parallélisme de pipeline est une stratégie d'entraînement distribué pour les réseaux de neurones profonds qui partitionne les couches d'un modèle en étapes séquentielles, chacune étant assignée à un dispositif différent (comme des GPU ou des TPU). Il combine le parallélisme de modèle avec le parallélisme de données en traitant des micro-lots à travers les étapes de manière logicielle-pipeline, visant à augmenter le débit d'entraînement et à réduire le temps d'inactivité (bulles) que le parallélisme naïf par couches entraînerait.
Contexte et motivation
L'entraînement de grands réseaux de neurones, en particulier les grands modèles de langage et d'autres modèles de apprentissage profond, nécessite d'énormes ressources de calcul. À mesure que la taille des modèles atteint des milliards ou des billions de paramètres, l'entraînement sur un seul dispositif devient irréalisable en raison des limites de mémoire et de calcul. Le parallélisme de données réplique le modèle entier sur chaque dispositif et divise les données, mais il échoue lorsque le modèle dépasse la mémoire d'un seul dispositif. Le parallélisme de modèle divise le modèle lui-même entre les dispositifs, mais le placement naïf par couches entraîne une sous-utilisation sévère : un seul dispositif calcule à la fois pendant que les autres attendent, ce qui donne un faible débit et un temps d'inactivité élevé.
Le parallélisme de pipeline répond à cela en divisant le modèle en étapes (généralement des groupes contigus de couches) et en traitant plusieurs micro-lots simultanément. L'idée clé est de chevaucher le calcul entre les étapes : pendant que l'étape 1 calcule sur le micro-lot 2, l'étape 2 peut calculer sur le micro-lot 1, et ainsi de suite. Cela ressemble à une chaîne de montage, où chaque étape traite continuellement les données entrantes et transmet les résultats à l'étape suivante.
Concepts fondamentaux
L'unité de travail fondamentale dans le parallélisme de pipeline est le micro-lot. Un grand lot d'entraînement est divisé en micro-lots plus petits qui traversent les étapes du pipeline séquentiellement. Chaque étape effectue des passes avant et arrière sur ses couches assignées. Le pipeline est "rempli" en alimentant les micro-lots un après l'autre, et après une période de préchauffage initiale, toutes les étapes peuvent être occupées simultanément, atteignant un débit élevé.
L'efficacité d'un pipeline est souvent mesurée par le ratio de bulles, qui est la fraction du temps où les étapes sont inactives. Pour un pipeline avec p étapes et m micro-lots, le débit idéal est atteint lorsque m est beaucoup plus grand que p, minimisant le surcoût relatif des bulles. Cependant, augmenter m augmente également l'utilisation de la mémoire pour les activations, créant un compromis.
Schémas principaux
Plusieurs schémas de parallélisme de pipeline ont été proposés, différant dans la façon dont ils planifient les passes avant et arrière et comment ils gèrent les mises à jour de gradient.
GPipe
Google a introduit GPipe en 2019. Il partitionne le modèle en étapes et utilise un schéma simple : chaque micro-lot effectue des passes avant à travers toutes les étapes, puis des passes arrière à travers toutes les étapes dans l'ordre inverse. GPipe utilise des mises à jour de gradient synchrones, ce qui signifie que toutes les étapes attendent que tous les micro-lots soient terminés avant de mettre à jour les poids. Cela garantit des gradients cohérents mais introduit une bulle qui croît avec le nombre d'étapes. GPipe nécessite également soit le recalcul des activations, soit leur stockage, échangeant du calcul contre de la mémoire.
PipeDream
PipeDream, de Microsoft Research (2019), utilise un schéma asynchrone où chaque étape traite les micro-lots de manière round-robin, effectuant les passes avant et arrière dès que possible. Cela réduit la bulle mais introduit un retard de poids : différentes étapes peuvent utiliser différentes versions des poids, ce qui peut nuire à la convergence. PipeDream utilise également une technique appelée "weight stashing" pour maintenir plusieurs versions des poids pour différents micro-lots, augmentant le surcoût mémoire.
PipeDream-2BW et variantes
PipeDream-2BW (2020) améliore PipeDream en utilisant deux tampons de poids (un pour la passe avant, un pour la passe arrière) pour réduire la mémoire et améliorer la convergence. Des variantes comme V-Pipe et PipeMare explorent différentes stratégies de synchronisation.
Megatron-LM et planification entrelacée
NVIDIA's Megatron-LM (2019) a combiné le parallélisme de tenseurs (divisant les couches individuelles) avec le parallélisme de pipeline. Plus tard, Megatron-2 (2020) a introduit un schéma entrelacé qui divise le modèle en plus d'étapes que de dispositifs, permettant à chaque dispositif de gérer plusieurs étapes. Cela réduit la taille de la bulle en augmentant le nombre de micro-lots pouvant être en vol, au prix de plus de communication.
Chimera et autres schémas récents
Chimera (2021) utilise un pipeline bidirectionnel pour réduire davantage les bulles. D'autres schémas comme PTD-P (2021) et ZeroBubble (2023) explorent des planifications plus sophistiquées pour éliminer ou réduire les bulles de pipeline.
Considérations mathématiques et pratiques
Le parallélisme de pipeline est souvent combiné avec d'autres stratégies de parallélisme. Le parallélisme de tenseurs divise les couches individuelles entre les dispositifs, tandis que le parallélisme de pipeline divise le modèle verticalement. Le parallélisme de données réplique le pipeline entre des groupes de dispositifs. Cette approche hybride est standard dans l'entraînement de grands modèles, comme on le voit dans des systèmes comme Megatron-Turing NLG et OpenAI's GPT-4.
Le choix de la profondeur du pipeline (nombre d'étapes) et de la taille des micro-lots affecte les performances. Des pipelines plus profonds réduisent la mémoire par étape mais augmentent la communication et le surcoût des bulles. Des micro-lots plus grands améliorent l'utilisation mais consomment plus de mémoire. La configuration optimale dépend de la taille du modèle, de la mémoire du dispositif et de la bande passante d'interconnexion.
Applications et impact
Le parallélisme de pipeline a été essentiel pour entraîner certains des plus grands modèles d'IA. Par exemple, GPT-3 (175B paramètres) a été entraîné en utilisant une combinaison de parallélisme de modèle et de pipeline. Les modèles Claude d'Anthropic et PaLM de Google reposent également sur des techniques similaires. L'approche est prise en charge par les principaux frameworks, y compris PyTorch (via le module torch.distributed.pipeline), TensorFlow (via tf.distribute) et JAX (avec pjit).
Défis et limitations
Le parallélisme de pipeline introduit plusieurs défis. Le surcoût de communication découle de l'envoi d'activations et de gradients entre les étapes, ce qui peut devenir un goulot d'étranglement sur des interconnexions lentes. Le déséquilibre de charge se produit lorsque les étapes ont des temps de calcul inégaux, réduisant l'efficacité. La pression mémoire due au stockage des activations pour les passes arrière peut être atténuée par le recalcul des activations, mais au prix d'un coût de calcul. De plus, les schémas asynchrones peuvent souffrir de problèmes de convergence dus à des poids obsolètes, tandis que les schémas synchrones entraînent des bulles.
Directions futures
La recherche continue sur l'amélioration du parallélisme de pipeline, en se concentrant sur la réduction des bulles, la minimisation de la mémoire et l'adaptation au matériel hétérogène. Des techniques comme le partitionnement automatique des étapes, l'équilibrage de charge dynamique et l'intégration avec les modèles mixture-of-experts sont des domaines actifs. L'essor de modèles extrêmement grands et le besoin d'un entraînement efficace sur des clusters cloud maintiendront probablement le parallélisme de pipeline comme un composant clé des systèmes d'entraînement distribué.