Parallélisme de modèle

Le parallélisme de modèle est une technique de calcul distribué utilisée dans l'apprentissage automatique et le apprentissage profond pour entraîner ou inférer de grands réseaux neuronaux qui ne tiennent pas dans la mémoire d'un seul appareil, tel qu'un GPU ou un TPU. Contrairement au parallélisme de données, où chaque appareil détient une copie complète du modèle et traite différents lots de données, le parallélisme de modèle partitionne le modèle lui-même, généralement en attribuant différentes couches ou sous-ensembles de couches à différents appareils. Cette approche est essentielle pour faire évoluer des modèles comme les transformeurs basés sur les grands modèles de langage, qui peuvent avoir des milliards ou des billions de paramètres.

Vue d'ensemble

Dans le parallélisme de modèle, le réseau neuronal est divisé le long de sa profondeur (par couches) ou, dans des formes plus avancées, le long d'autres dimensions telles que la largeur ou même des opérations tensorielles individuelles. La forme la plus simple est le partitionnement par couches, où des couches consécutives sont placées sur différents appareils. Pendant la propagation avant, les activations circulent d'un appareil à l'autre ; pendant la rétropropagation, les gradients circulent dans la direction opposée. Cela crée une exécution de type pipeline, qui peut introduire des temps d'inactivité (bulles) si elle n'est pas soigneusement planifiée. Le parallélisme de modèle est souvent contrasté avec le parallélisme de données, mais les deux peuvent être combinés dans des approches hybrides, comme le parallélisme 3D utilisé pour entraîner de grands modèles.

Motivation

Les réseaux neuronaux modernes, en particulier les grands modèles de langage, ont connu une croissance exponentielle de leur taille. Par exemple, GPT-3 (2020) a 175 milliards de paramètres, nécessitant plus de 350 Go de mémoire rien que pour les poids en précision 32 bits, dépassant largement la capacité de 80 Go des GPU haut de gamme comme l'A100. Même avec des optimisations de mémoire comme la précision mixte et le checkpointing des gradients, l'entraînement sur un seul appareil est infaisable. Le parallélisme de modèle permet aux chercheurs d'entraîner des modèles qui sont des ordres de grandeur plus grands que ce qu'un seul appareil peut contenir, en distribuant le modèle sur un cluster d'appareils connectés par des interconnexions à haute vitesse.

Types de parallélisme de modèle

Parallélisme par couches (pipeline)

C'est la forme la plus simple, où le modèle est divisé en étapes séquentielles, chacune assignée à un appareil différent. Par exemple, un réseau de 100 couches pourrait être divisé en 10 étapes de 10 couches chacune, chaque étape étant sur un GPU séparé. Pendant la passe avant, les données circulent à travers l'étape 1, puis l'étape 2, et ainsi de suite. Le principal défi est l'équilibrage de charge et la réduction des bulles de pipeline. Des techniques comme GPipe et PipeDream introduisent le micro-batching et la planification pour améliorer l'utilisation.

Parallélisme tensoriel

Le parallélisme tensoriel divise les opérations individuelles (par exemple, les multiplications matricielles) sur plusieurs appareils. Par exemple, dans le mécanisme d'attention d'un transformeur, les projections de requête, de clé et de valeur peuvent être partitionnées sur les GPU, avec des résultats combinés via des opérations all-reduce. Cette approche est utilisée dans Megatron-LM et est efficace pour réduire la mémoire par appareil tout en maintenant une efficacité de calcul élevée, mais nécessite une communication à haute bande passante.

Parallélisme expert

Utilisé dans les modèles de mélange d'experts (MoE), le parallélisme expert place différents réseaux experts sur différents appareils, tandis que le routeur (réseau de gating) distribue les jetons aux experts appropriés. C'est une forme de parallélisme de modèle qui permet des nombres massifs de paramètres sans augmentation proportionnelle du calcul, car seul un sous-ensemble d'experts est activé par jeton.

Défis d'implémentation

Le parallélisme de modèle introduit plusieurs défis :

• Surcharge de communication : Les appareils doivent échanger des activations et des gradients, ce qui peut devenir un goulot d'étranglement, surtout avec des interconnexions lentes.
• Déséquilibre de charge : Des tailles de couches ou des exigences de calcul inégales peuvent entraîner l'inactivité de certains appareils pendant que d'autres travaillent.
• Bulles de pipeline : Dans le parallélisme de pipeline, les appareils peuvent attendre des données des étapes précédentes, réduisant l'utilisation.
• Fragmentation de la mémoire : Le partitionnement des modèles peut entraîner une utilisation inégale de la mémoire, nécessitant un placement et une planification soigneux.
• Tolérance aux pannes : Avec de nombreux appareils, la probabilité de défaillance augmente, nécessitant des mécanismes de checkpointing et de récupération.

Comparaison avec le parallélisme de données

Dans le parallélisme de données, chaque appareil détient une copie complète du modèle et traite un mini-lot de données différent. Les gradients sont moyennés sur les appareils après chaque étape. C'est simple et évolue bien pour les modèles qui tiennent dans la mémoire d'un seul appareil. Le parallélisme de modèle, en revanche, est nécessaire lorsque le modèle lui-même est trop grand. Cependant, le parallélisme de modèle a souvent des coûts de communication plus élevés et une efficacité plus faible en raison des dépendances séquentielles. En pratique, l'entraînement à grande échelle utilise les deux : le parallélisme de données entre les nœuds et le parallélisme de modèle au sein d'un nœud.

Applications

Le parallélisme de modèle est essentiel pour entraîner de grands modèles de langage comme GPT-4, PaLM et LLaMA, ainsi que pour les servir en production. Des entreprises comme OpenAI, Google (via Google Cloud) et Anthropic s'appuient sur le parallélisme de modèle pour entraîner et déployer leurs modèles. Il est également utilisé dans d'autres domaines, comme la vision par ordinateur (par exemple, les CNN 3D) et le calcul scientifique, où les modèles sont trop grands pour un seul accélérateur.

Support logiciel

Plusieurs frameworks offrent un support intégré pour le parallélisme de modèle :

• PyTorch : Propose torch.distributed avec le parallélisme de pipeline (par exemple, torch.distributed.pipeline.sync.Pipe) et le parallélisme tensoriel via des bibliothèques tensor_parallel.
• TensorFlow : Fournit tf.distribute avec des stratégies comme TPUStrategy et MultiWorkerMirroredStrategy qui peuvent intégrer le parallélisme de modèle.
• Megatron-LM : Une bibliothèque spécialisée de NVIDIA pour le parallélisme tensoriel et de pipeline.
• DeepSpeed : De Microsoft, offre ZeRO (Zero Redundancy Optimizer) qui est une forme de parallélisme de modèle qui partitionne les états de l'optimiseur, les gradients et les paramètres.
• JAX : Avec jax.sharding, permet un placement explicite des appareils et une exécution parallèle.

Considérations matérielles

Le parallélisme de modèle repose fortement sur la communication inter-appareils à haute vitesse. Au sein d'un seul nœud, NVLink et PCIe fournissent des connexions rapides ; entre les nœuds, InfiniBand ou Ethernet haute vitesse sont utilisés. Le matériel spécialisé comme les GPU NVIDIA et les TPU Google incluent souvent des interconnexions dédiées (par exemple, NVSwitch, liens TPU) pour faciliter le parallélisme de modèle. Des entreprises comme Cerebras et SambaNova ont conçu des systèmes avec une grande mémoire sur puce pour réduire le besoin de parallélisme de modèle, mais pour une échelle extrême, les approches distribuées restent nécessaires.

Avancées récentes

Des recherches récentes se sont concentrées sur l'amélioration de l'efficacité du parallélisme de modèle. Des techniques comme le parallélisme de séquence (diviser la dimension de séquence) et le parallélisme de contexte (pour les longues séquences) ont été développées. Des algorithmes de partitionnement automatique de modèles, comme ceux d'Alpa et FlexFlow, optimisent le placement des opérations sur les appareils. De plus, le parallélisme hybride combinant le parallélisme de données, tensoriel et de pipeline est désormais standard dans l'entraînement à grande échelle, comme en témoigne le modèle Megatron-Turing NLG (530B paramètres) entraîné par NVIDIA et Microsoft.

Limitations et orientations futures

Malgré sa puissance, le parallélisme de modèle a des limitations. La surcharge de communication peut dominer pour les petits modèles ou les réseaux lents. La complexité du partitionnement et de la planification augmente avec la taille du modèle et l'hétérogénéité du cluster. Les orientations futures incluent un auto-parallélisme plus intelligent, une meilleure compression de la communication et de nouvelles architectures intrinsèquement plus parallélisables. Avec l'essor de l'IA générative et des modèles toujours plus grands, le parallélisme de modèle restera une pierre angulaire de l'apprentissage profond évolutif.

Voir aussi

• Parallélisme de données
• Parallélisme de pipeline
• Parallélisme tensoriel
• Mélange d'experts
• Entraînement distribué
• Grand modèle de langage
• Cluster GPU

Références

• Narayanan, D., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training.
• Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.
• Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.
• Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.