Principales fonctions du parallélisme des modèles distribués SageMaker - Amazon SageMaker

Principales fonctions du parallélisme des modèles distribués SageMaker

Le parallélisme des modèles distribués d'Amazon SageMaker rend le parallélisme des modèles plus accessible en fournissant une division automatisée du modèle et une planification sophistiquée de l'exécution de pipeline. Les algorithmes de division de modèle peuvent optimiser la vitesse ou la consommation de mémoire. La bibliothèque prend également en charge le partitionnement manuel. Lorsque vous utilisez la bibliothèque, l'entraînement est exécuté en pipeline sur des micro-lots afin de maximiser l'utilisation du GPU.

Vous pouvez configurer ces fonctions à l'aide de quelques lignes de code lorsque vous créez votre script d'entraînement et définissez votre SageMaker PyTorch ou Tensorflow Estimator. Utilisez les sections suivantes pour en savoir plus sur ces fonctions principales de la bibliothèque.

Note

Les bibliothèques d'entraînement distribué SageMaker sont uniquement disponibles via les conteneurs AWS Deep Learning Containers pour les cadres TensorFlow, PyTorch et HuggingFace sur la plateforme d'entraînement SageMaker. Pour utiliser les bibliothèques, vous devez utiliser le kit SDK Python SageMaker ou les API SageMaker via le kit SDK for Python (Boto3) ou la AWS Command Line Interface. La documentation contient des instructions et des exemples sur l'utilisation des bibliothèques d'entraînement distribué avec le kit SDK Python SageMaker.

Division automatisée du modèle

Lorsque vous utilisez la bibliothèque de modèles parallèles de SageMaker, vous pouvez tirer avantage de la division automatisée du modèle, également appelée partitionnement automatique du modèle. La bibliothèque utilise un algorithme de partitionnement qui équilibre la mémoire, réduit la communication entre les périphériques et optimise la performance. Vous pouvez configurer l'algorithme de partitionnement automatique de sorte à optimiser la vitesse ou la mémoire.

Vous pouvez également utiliser la division manuelle du modèle. Nous vous recommandons la division automatisée du modèle, sauf si vous connaissez très bien l'architecture du modèle et que vous savez déjà comment partitionner efficacement votre modèle.

Fonctionnement

Le partitionnement automatique intervient dès la première étape d'entraînement, lors du tout premier appel de la fonction décorée smp.step. Durant cet appel, la bibliothèque commence par créer une version du modèle sur la RAM du CPU (pour éviter les limitations de mémoire GPU), puis elle analyse le graphe du modèle et décide du partitionnement. À partir de cette décision, chaque partition de modèle est chargée sur un GPU, et ce n'est qu'alors que la première étape est exécutée. Ces étapes d'analyse et de partitionnement peuvent contribuer à allonger la première étape de l'entraînement.

Dans les deux cadres, la bibliothèque gère la communication entre les périphériques via son propre backend, qui est optimisé pour l'infrastructure AWS.

La conception de la partition automatique s'adapte aux caractéristiques du cadre, et la bibliothèque effectue le partitionnement au niveau de granularité le plus naturel dans chaque cadre. Par exemple, dans TensorFlow, chaque opération spécifique peut être affectée à un périphérique différent, alors que dans PyTorch, l'affectation se fait au niveau du module, chaque module comprenant plusieurs opérations. La section qui suit examine les spécificités de conception dans chaque cadre.

Durant la première étape d'entraînement, la bibliothèque de modèles parallèles exécute en interne une étape de traçage destinée à créer le graphe du modèle et à déterminer les formes du tenseur et des paramètres. Après cette étape de traçage, la bibliothèque crée un arbre, qui se compose des objets nn.Module imbriqués dans le modèle, ainsi que de données supplémentaires collectées à partir du traçage, comme la quantité de nn.Parameters stockés et le temps d'exécution de chaque nn.Module.

Ensuite, la bibliothèque traverse cet arbre depuis la racine et exécute un algorithme de partitionnement qui affecte chaque nn.Module à un périphérique, ce qui équilibre la charge de calcul (mesurée par le temps d'exécution du module) et l'utilisation de la mémoire (mesurée par la taille totale des nn.Parameter stockés et les activations). Si plusieurs nn.Modules partagent le même nn.Parameter, ces modules sont alors placés sur le même périphérique afin de ne pas conserver plusieurs versions du même paramètre. Une fois la décision de partitionnement prise, les modules et les poids affectés sont chargés sur leurs périphériques.

La bibliothèque de modèles parallèles analyse les tailles des variables entraînables et la structure du graphe, et utilise en interne un algorithme de partitionnement des graphes. Cet algorithme affecte un périphérique pour chaque opération afin de réduire le volume de communication nécessaire entre les périphériques, sous réserve des deux contraintes suivantes :

  • Équilibrage du nombre de variables stockées dans chaque périphérique

  • Équilibrage du nombre d'opérations exécutées dans chaque périphérique

Si vous spécifiez speed pour optimize (dans les paramètres de modèle parallèle dans le kit SDK Python), la bibliothèque essaie d'équilibrer le nombre d'opérations et d'objets tf.Variable dans chaque périphérique. Sinon, elle essaie d'équilibrer la taille totale de tf.Variables.

Une fois la décision de partitionnement prise, la bibliothèque crée une représentation sérialisée du sous-graphe que chaque périphérique doit exécuter et l'importe sur chaque périphérique. Lors du partitionnement, la bibliothèque place les opérations qui consomment la même tf.Variable et les opérations qui font partie de la même couche Keras sur le même périphérique. Elle respecte aussi les contraintes de colocalisation imposées par TensorFlow. Cela signifie, par exemple, que si deux couches Keras partagent une tf.Variable, toutes les opérations qui font partie de ces couches sont placées sur un seul périphérique.

Comparaison de la division automatisée du modèle entre les cadres

Dans TensorFlow, l'unité fondamentale de calcul est une tf.Operation, et TensorFlow représente le modèle sous la forme d'un graphe orienté acyclique (DAG) d'opérations tf.Operation. La bibliothèque de modèles parallèles partitionne donc ce DAG pour que chaque nœud aille à un périphérique. Ce qui est intéressant ici est que les objets tf.Operation sont suffisamment riches en attributs personnalisables et qu'ils sont universels, c'est-à-dire que chaque modèle comprendra obligatoirement un graphe de ces objets.

Par contre, cette notion d'opération suffisamment riche et universelle n'a pas d'équivalent dans PyTorch. L'unité de calcul la plus proche dotée de ces caractéristiques dans PyTorch est un nn.Module, qui se situe à un niveau de granularité nettement supérieur. Voilà pourquoi la bibliothèque fait le partitionnement à ce niveau dans PyTorch.

Division manuelle du modèle

Si vous voulez spécifier manuellement le partitionnement de votre modèle entre les périphériques, vous pouvez utiliser la division manuelle du modèle à l'aide de gestionnaires de contexte de smp.partition.

Pour utiliser cette option, définissez auto_partition sur False, et définissez un default_partitiondans le kit SDK Python SageMaker. Toute opération non explicitement placée sur une partition à l'aide du gestionnaire de contexte de smp.partition est exécutée sur la default_partition. Dans ce cas, la logique de division automatisée est contournée et chaque opération est placée de la façon dont vous le spécifiez. En s'appuyant sur la structure de graphe ainsi obtenue, la bibliothèque de modèles parallèles crée automatiquement un calendrier d'exécution de pipeline.

Calendrier d'exécution de pipeline

Une caractéristique majeure de la bibliothèque de modèles parallèles distribués de SageMaker est l'exécution de pipeline, qui détermine l'ordre dans lequel les calculs sont effectués et les données sont traitées entre les périphériques pendant l'entraînement du modèle. Le Pipelining est une technique employée pour obtenir une véritable parallélisation dans le parallélisme des modèles grâce aux calculs simultanés effectués par les GPU sur différents échantillons de données, et surmonter la baisse de performance due au calcul séquentiel.

Le Pipelining repose sur la division d'un mini-lot en micro-lots, qui sont introduits individuellement dans le pipeline d'entraînement un par un et respectent un calendrier d'exécution défini par le moteur d'exécution de la bibliothèque. Un micro-lot est un sous-ensemble plus petit d'un mini-lot d'entraînement donné. Le calendrier du pipeline détermine quel micro-lot est exécuté par quel périphérique pour chaque créneau horaire.

Par exemple, en fonction du calendrier du pipeline et de la partition du modèle, le GPU i peut effectuer un calcul (vers l'avant ou vers l'arrière) sur le micro-lot b, tandis que le GPU i+1 effectue un calcul sur le micro-lot b+1, les deux GPU étant ainsi actifs en même temps. Durant une seule transmission vers l'avant et vers l'arrière, le flux d'exécution d'un seul micro-lot peut visiter le même périphérique plusieurs fois, en fonction de la décision de partitionnement. Par exemple, une opération située au début du modèle peut être placée sur le même périphérique qu'une opération située à la fin du modèle, tandis que les opérations situées entre les deux sont placées sur différents périphériques, de sorte que ce périphérique est visité deux fois.

La bibliothèque propose deux calendriers de pipeline différents, simple et entrelacé, qui peuvent être configurés à l'aide du paramètre pipeline dans le kit SDK Python SageMaker. Dans la plupart des cas, le pipeline entrelacé permet d'obtenir de meilleures performances grâce à une utilisation plus efficace des GPU.

Pipeline entrelacé

Dans un pipeline entrelacé, la priorité est donnée, dans la mesure du possible, à l'exécution vers l'arrière des micro-lots. Cela permet de libérer plus rapidement la mémoire utilisée pour les activations et donc d'utiliser la mémoire plus efficacement. Cela permet aussi de mettre à l'échelle le nombre de micro-lots à un niveau supérieur et de réduire ainsi le temps d'inactivité des GPU. À l'état d'équilibre, chaque périphérique alterne entre les transmissions vers l'avant et vers l'arrière. Cela signifie que la transmission vers l'arrière d'un micro-lot peut s'exécuter avant la fin de la transmission vers l'avant d'un autre micro-lot.

La figure précédente illustre un exemple de calendrier d'exécution du pipeline entrelacé sur 2 GPU. Sur la figure, F0 représente la transmission vers l'avant pour le micro-lot 0, et B1 la transmission vers l'arrière pour le micro-lot 1. Update représente la mise à jour des paramètres par l'optimiseur. Dans la mesure du possible, GPU0 donne toujours la priorité aux transmissions vers l'arrière (en exécutant, par exemple, B0 avant F2), ce qui permet d'effacer la mémoire utilisée pour les activations précédentes.

Pipeline simple

À contrario, un pipeline simple termine d'exécuter la transmission vers l'avant pour chaque micro-lot avant de démarrer la transmission vers l'arrière. En d'autres termes, le pipeline exécute les étapes de transmission vers l'avant et vers l'arrière en interne. La figure suivante illustre un exemple de fonctionnement, sur 2 GPU.

Exécution de pipeline dans des cadres spécifiques

Utilisez les sections suivantes pour en savoir plus sur les décisions de planification de pipeline spécifiques au cadre prises par la bibliothèque de modèles parallèles distribués de SageMaker pour Tensorflow et PyTorch.

Exécution de pipeline avec TensorFlow

L'image suivante illustre un exemple de graphe TensorFlow partitionné par la bibliothèque de modèles parallèles à l'aide d'une division automatisée du modèle. Lorsqu'un graphe est divisé, chaque sous-graphe obtenu est répliqué B fois (sauf pour les variables), B désignant le nombre de micro-lots. Sur cette figure, chaque sous-graphe est répliqué 2 fois (B=2). Une opération SMPInput est insérée à chaque entrée d'un sous-graphe, et une opération SMPOutput est insérée à chaque sortie. Ces opérations communiquent avec le backend de la bibliothèque pour transférer les tenseurs entre eux de façon bidirectionnelle.

L'image suivante illustre un exemple de 2 sous-graphes divisés avec B=2, avec ajout d'opérations de gradient. Le gradient d'une opération SMPInput est une opération SMPOutput, et vice versa. Les gradients peuvent ainsi circuler vers l'arrière pendant la rétro-propagation.

Ce GIF illustre un exemple de calendrier d'exécution de pipeline entrelacé avec B=2 micro-lots et 2 sous-graphes. Chaque périphérique exécute l'un des réplicas de sous-graphe séquentiellement afin d'améliorer l'utilisation du GPU. À mesure que B augmente, la fraction d'intervalles de temps d'inactivité tend vers zéro. Chaque fois qu'un calcul (vers l'avant ou vers l'arrière) doit être fait sur un réplica de sous-graphe spécifique, la couche de pipeline signale aux opérations SMPInput bleues correspondantes qu'il est temps de démarrer l'exécution.

Une fois que les gradients de tous les micro-lots d'un seul mini-lot sont calculés, la bibliothèque combine les gradients entre les micro-lots, qui peuvent ensuite être appliqués aux paramètres.

Exécution de pipeline avec PyTorch

En théorie, le pipelining s'exécute de façon similaire à PyTorch. Cependant, comme PyTorch n'implique pas de graphes statiques, le paradigme de pipelining utilisé par la fonction PyTorch de la bibliothèque de modèles parallèles est plus dynamique.

Comme dans TensorFlow, chaque lot est divisé en plusieurs micro-lots, qui sont exécutés individuellement sur chaque périphérique. Toutefois, le calendrier d'exécution est géré via des serveurs d'exécution lancés sur chaque périphérique. Chaque fois que le périphérique actuel a besoin de la sortie d'un sous-module placé sur un autre périphérique, une demande d'exécution est envoyée au serveur d'exécution du périphérique distant et les tenseurs d'entrée au sous-module. Le serveur exécute alors ce module avec les entrées données et renvoie la réponse au périphérique actuel.

Comme le périphérique actuel est inactif pendant l'exécution du sous-module distant, l'exécution locale du micro-lot actuel s'interrompt et le moteur d'exécution de la bibliothèque bascule l'exécution vers un autre micro-lot sur lequel le périphérique actuel peut travailler activement. La priorité donnée aux micro-lots est déterminée par le calendrier de pipeline choisi. Dans le cas d'un calendrier de pipeline entrelacé, les micro-lots qui se trouvent dans l'étape de transmission vers l'arrière du calcul sont prioritaires dans la mesure du possible.