Le parallélisme de tenseurs est une technique permettant de répartir le calcul d'une seule couche de réseau de neurones sur plusieurs dispositifs matériels. Dans cette approche, les matrices de poids d'une couche sont divisées en fragments, et chaque dispositif détient une partie des paramètres et effectue la partie correspondante de la multiplication matricielle. Les résultats partiels sont ensuite combinés via des opérations de communication collective, telles que l'all-reduce, pour produire la sortie de la couche. Cette méthode est une forme de parallélisme de modèle, distincte du parallélisme de données où chaque dispositif détient une copie complète du modèle et traite différents échantillons de données. Le parallélisme de tenseurs est essentiel pour l'entraînement et l'inférence de très grands modèles, tels que les grands modèles de langage, qui dépassent la capacité mémoire d'un seul accélérateur.
La motivation principale du parallélisme de tenseurs réside dans les exigences en mémoire et en calcul des modèles d'apprentissage profond modernes. Par exemple, un modèle avec des centaines de milliards de paramètres ne peut pas tenir dans la mémoire d'un seul GPU ou TPU, qui varie généralement de 16 à 80 gigaoctets. En fragmentant les matrices de poids, le parallélisme de tenseurs permet de répartir le modèle sur de nombreux dispositifs, permettant l'utilisation de modèles plus grands et des temps d'entraînement plus rapides. Il réduit également l'empreinte mémoire par dispositif et permet des tailles de lots plus grandes. Cependant, le parallélisme de tenseurs introduit un surcoût de communication, car les dispositifs doivent synchroniser les résultats partiels, ce qui peut devenir un goulot d'étranglement, surtout sur des interconnexions plus lentes.
Le parallélisme de tenseurs est couramment utilisé en conjonction avec d'autres stratégies de parallélisme, telles que le parallélisme de pipeline et le parallélisme de données, pour faire évoluer l'entraînement sur des milliers de dispositifs. C'est un composant clé de l'infrastructure d'entraînement de nombreux modèles de pointe, y compris ceux développés par OpenAI, Anthropic et Google DeepMind. La technique est également pertinente pour l'inférence, où elle permet de servir de grands modèles avec une faible latence en répartissant le calcul sur plusieurs accélérateurs.
Contexte historique
Le concept de division des opérations matricielles sur plusieurs processeurs remonte aux premières recherches en calcul parallèle. Dans les années 1980 et 1990, les chercheurs ont exploré des implémentations parallèles de réseaux de neurones, mais les outils matériels et logiciels étaient limités. La forme moderne du parallélisme de tenseurs a émergé avec l'essor de l'apprentissage profond basé sur GPU dans les années 2010. L'une des premières utilisations notables a été dans l'entraînement de grands réseaux convolutifs, où des chercheurs de l'Université de Toronto et d'autres institutions ont expérimenté le parallélisme de modèle. Cependant, c'est l'avènement des modèles basés sur le transformeur, tels que l'architecture Transformeur introduite en 2017, qui a fait du parallélisme de tenseurs une technique courante. Les couches d'attention et les réseaux feed-forward du Transformeur sont composés de grandes multiplications matricielles, qui sont des candidats naturels pour la fragmentation.
En 2019, des chercheurs de NVIDIA et d'autres organisations ont publié des descriptions détaillées du parallélisme de tenseurs pour l'entraînement de grands modèles de langage. Le framework Megatron-LM, développé par NVIDIA, a démontré un parallélisme de tenseurs efficace pour les modèles transformeurs, atteignant une mise à l'échelle quasi linéaire sur plusieurs GPU. Ce travail a jeté les bases de systèmes d'entraînement à grande échelle ultérieurs, tels que GShard de Google et DeepSpeed de Microsoft, qui intègrent le parallélisme de tenseurs comme fonctionnalité de base.
Comment fonctionne le parallélisme de tenseurs
Dans une couche typique de réseau de neurones, la passe avant calcule une multiplication matricielle : Y = XW, où X est l'activation d'entrée, W est la matrice de poids et Y est la sortie. Dans le parallélisme de tenseurs, la matrice de poids W est partitionnée le long d'une ou plusieurs dimensions. Pour une couche entièrement connectée, une approche courante consiste à diviser W en colonnes (c'est-à-dire le long de la dimension de sortie). Chaque dispositif détient un sous-ensemble des colonnes de W et calcule une sortie partielle. Les sorties partielles sont ensuite combinées à l'aide d'une opération all-reduce pour produire le Y final. C'est ce qu'on appelle le partitionnement parallèle en colonnes. Alternativement, W peut être divisé en lignes (le long de la dimension d'entrée), nécessitant une opération all-gather sur l'entrée et un all-reduce sur la sortie.
Pour les modèles transformeurs, le parallélisme de tenseurs est appliqué à la fois à l'attention multi-têtes et aux réseaux feed-forward. Dans l'attention multi-têtes, les matrices de poids de requête, de clé et de valeur sont fragmentées sur les dispositifs, et les têtes d'attention sont distribuées. La projection de sortie est ensuite combinée. Dans le réseau feed-forward, les deux couches linéaires sont généralement partitionnées : la première couche est divisée en colonnes, et la seconde couche est divisée en lignes, de sorte que les activations intermédiaires soient fragmentées et la sortie soit réduite. Cette conception minimise la communication en ne nécessitant des opérations all-reduce qu'à des points spécifiques.
Le schéma de communication est critique. Chaque opération all-reduce nécessite l'échange de données entre tous les dispositifs du groupe de parallélisme de tenseurs. Le volume de communication est proportionnel à la taille des activations, pas aux poids, ce qui est souvent plus petit. Cependant, pour de très grands modèles, la communication peut encore être significative. Pour atténuer cela, le parallélisme de tenseurs est généralement utilisé au sein d'un seul nœud avec des interconnexions à haute vitesse, telles que NVLink de NVIDIA ou Infinity Fabric de AMD, tandis que le parallélisme de données est utilisé entre les nœuds.
Comparaison avec d'autres techniques de parallélisme
Le parallélisme de tenseurs est l'une des plusieurs stratégies de parallélisme de modèle. Le parallélisme de pipeline, par exemple, divise le modèle par couches, chaque dispositif étant responsable d'un ensemble contigu de couches. Cela réduit la communication par rapport au parallélisme de tenseurs car seules les activations aux frontières des couches sont échangées. Cependant, le parallélisme de pipeline peut souffrir de temps d'inactivité en raison des bulles de pipeline. Le parallélisme de données, en revanche, réplique le modèle sur chaque dispositif et traite différents échantillons de données, nécessitant une synchronisation des gradients après chaque étape. Le parallélisme de tenseurs est complémentaire aux deux : il peut être combiné avec le parallélisme de pipeline pour réduire la mémoire par dispositif et avec le parallélisme de données pour augmenter le débit.
Une autre technique connexe est le parallélisme de séquence, qui fragmente la dimension de séquence dans les modèles transformeurs. Cela est souvent utilisé en conjonction avec le parallélisme de tenseurs pour réduire davantage l'utilisation de la mémoire. Le parallélisme d'experts, utilisé dans les modèles de mélange d'experts, fragmente les modules experts sur les dispositifs, ce qui est une forme de parallélisme de modèle mais pas strictement du parallélisme de tenseurs.
Applications et cas d'utilisation
Le parallélisme de tenseurs est largement utilisé dans l'entraînement et l'inférence de grands modèles de langage. Par exemple, le modèle GPT-3, développé par OpenAI, a été entraîné en utilisant une combinaison de parallélisme de modèle, y compris le parallélisme de tenseurs, sur des milliers de GPU NVIDIA. De même, les modèles Claude d'Anthropic et les modèles Gemini de Google DeepMind reposent sur le parallélisme de tenseurs pour atteindre des centaines de milliards de paramètres. En inférence, le parallélisme de tenseurs permet de servir des modèles comme GPT-4 avec une faible latence en répartissant le calcul sur plusieurs GPU, comme le font le serveur d'inférence Triton de NVIDIA et Text Generation Inference de Hugging Face.
Au-delà des modèles de langage, le parallélisme de tenseurs est également utilisé dans d'autres domaines, tels que la vision par ordinateur et le calcul scientifique, où de grands modèles ou de grandes tailles de lots nécessitent un calcul distribué. Par exemple, l'entraînement de grands transformeurs de vision (ViT) sur des images haute résolution peut bénéficier du parallélisme de tenseurs.
Défis et limitations
Malgré ses avantages, le parallélisme de tenseurs présente plusieurs défis. Le principal problème est le surcoût de communication. Les opérations all-reduce nécessitent des interconnexions à haute bande passante et faible latence. Sur des systèmes avec des réseaux plus lents, comme Ethernet, le parallélisme de tenseurs peut devenir inefficace. Par conséquent, il est généralement utilisé au sein d'un seul serveur ou d'un cluster étroitement couplé. Un autre défi est l'équilibrage de charge : une fragmentation inégale peut entraîner une sous-utilisation de certains dispositifs. De plus, le parallélisme de tenseurs nécessite une gestion minutieuse de la mémoire, car chaque dispositif doit stocker son fragment des poids et les activations intermédiaires. Enfin, implémenter manuellement le parallélisme de tenseurs est complexe ; cela nécessite de modifier le code du modèle et de gérer les primitives de communication. Heureusement, des frameworks comme PyTorch et TensorFlow offrent un support intégré pour le parallélisme de tenseurs, abstraisant une grande partie de la complexité.
Support logiciel et frameworks
Plusieurs frameworks et bibliothèques d'apprentissage profond prennent en charge le parallélisme de tenseurs. PyTorch offre le module torch.distributed et l'API tensor_parallel pour fragmenter les tenseurs. TensorFlow fournit tf.distribute avec un support pour le parallélisme de modèle. Des bibliothèques spécialisées telles que Megatron-LM, DeepSpeed et les Transformers de Hugging Face ont implémenté le parallélisme de tenseurs pour les modèles transformeurs. Ces outils permettent aux chercheurs et aux ingénieurs d'appliquer le parallélisme de tenseurs avec des modifications minimales du code, souvent en spécifiant simplement le nombre de dispositifs et la stratégie de fragmentation.
Directions futures
Alors que les modèles continuent de croître, le parallélisme de tenseurs restera une technique critique. Les développements futurs pourraient inclure des algorithmes de communication plus efficaces, tels que l'all-reduce hiérarchique, et une meilleure intégration avec du matériel hétérogène, y compris les GPU AMD, les accélérateurs Intel et des puces personnalisées comme AWS Trainium et les LPU de Groq. De plus, la découverte automatisée du parallélisme, où le système détermine la stratégie de fragmentation optimale, est un domaine de recherche actif. Avec la demande croissante pour l'IA à grande échelle, le parallélisme de tenseurs continuera d'évoluer pour répondre aux défis d'échelle et d'efficacité.