Présentation du parallélisme des modèles - Amazon SageMaker

Présentation du parallélisme des modèles

Le parallélisme de modèle est une méthode d'entraînement distribué dans laquelle le modèle de deep learning est partitionné sur plusieurs appareils, au sein des instances ou entre celles-ci. Cette page de présentation fournit un aperçu global du parallélisme de modèle, une description de la façon dont il peut aider à surmonter des problèmes liés à l'entraînement de modèles DL généralement de très grande taille, et des exemples de ce que propose la bibliothèque de parallélisme de modèle SageMaker pour aider à gérer les stratégies de parallélisme de modèle et la consommation de mémoire.

Qu'est-ce que le parallélisme des modèles ?

L'augmentation de la taille des modèles de deep learning (couches et paramètres) permet une meilleure précision pour des tâches complexes telles que la reconnaissance d'image et le traitement du langage naturel. Toutefois, il y a une limite à la taille maximale de modèle que vous pouvez faire tenir dans la mémoire d'un GPU individuel. Lors de l'entraînement de modèles DL, les limites de mémoire du GPU peuvent constituer un goulet d'étranglement :

  • Elles limitent la taille du modèle que vous pouvez entraîner, car l'empreinte mémoire d'un modèle évolue proportionnellement au nombre de paramètres.

  • Elles limitent la taille de lot par GPU pendant l'entraînement, ce qui réduit l'utilisation du GPU et l'efficacité de l'entraînement.

Pour que vous dépassiez les limites associées à l'entraînement d'un modèle sur un GPU unique, SageMaker propose la bibliothèque de parallélisme de modèle pour vous aider à distribuer et entraîner efficacement des modèles DL sur plusieurs nœuds de calcul. En outre, cette bibliothèque vous permet de profiter d'un entraînement distribué optimisé à l'aide d'appareils intégrant EFA, qui améliorent les performances de la communication entre les nœuds avec une faible latence, un débit élevé et le contournement du système d'exploitation.

Estimation des besoins en mémoire avant d'utiliser le parallélisme de modèle

Avant d'utiliser la bibliothèque de parallélisme de modèle SageMaker, tenez compte des éléments suivants pour vous faire une idée des besoins en mémoire associés à l'entraînement de modèles DL de grande taille.

Pour une tâche d'entraînement utilisant les optimiseurs AMP (FP16) et Adam, la mémoire GPU requise par paramètre est d'environ 20 octets, que nous pouvons décomposer comme suit :

  • Un paramètre FP16 ~ 2 octets

  • Un gradient FP16 ~ 2 octets

  • Un état d'optimiseur FP32 ~ 8 octets en fonction des optimiseurs Adam

  • Une copie de paramètre FP32 ~ 4 octets (requis pour l'opération optimizer apply (OA))

  • Une copie de gradient FP32 ~ 4 octets (requis pour l'opération OA)

Même un modèle DL relativement petit, avec 10 milliards de paramètres, peut nécessiter au moins 200 Go de mémoire, ce qui dépasse nettement la mémoire GPU standard (par exemple, NVIDIA A100 avec 40/80 Go de mémoire et V100 avec 16/32 Go) disponible sur un GPU individuel. Notez qu'en plus des besoins en mémoire pour les états de modèle et d'optimiseur, il existe d'autres consommateurs de mémoire tels que les activations générées dans la transmission vers l'avant. La mémoire requise peut être largement supérieure à 200 Go.

Pour un entraînement distribué, nous vous recommandons d'utiliser des instances Amazon EC2 P3 et P4 dotées de GPU Tensor Core NVIDIA A100 ou V100. Pour plus de détails sur les spécifications telles que les cœurs de processeur, la RAM, le volume de stockage attaché et la bande passante du réseau, consultez la section Calcul accéléré de la page Types d'instances Amazon EC2.

Même avec les instances de calcul accéléré, il est évident que des modèles avec environ 10 milliards de paramètres, tels que Megatron-LM et T5, et des modèles encore plus grands avec des centaines de milliards de paramètres, tels que GPT-3, ne peuvent pas faire tenir les réplicas de modèles dans chaque périphérique GPU.

Utilisation par la bibliothèque des techniques d'économie de mémoire et de parallélisme de modèle

La bibliothèque comprend deux types de fonctions de parallélisme de modèle, le parallélisme de pipeline et le parallélisme de tenseur, ainsi que d'autres fonctionnalités d'économie de mémoire, telles que le partitionnement de l'état de l'optimiseur, les points de contrôle d'activation et le déchargement d'activation. Toutes ces techniques peuvent être combinées pour entraîner efficacement des modèles de grande taille, tels que GPT-2, composés de centaines de milliards de paramètres.

Parallélisme de pipeline

Le parallélisme de pipeline partitionne l'ensemble de couches ou d'opérations sur l'ensemble de périphériques, laissant chaque opération intacte. Lorsque vous spécifiez une valeur pour le nombre de partitions de modèle (pipeline_parallel_degree), le nombre total de GPU (processes_per_host) doit être divisible par le nombre de partitions de modèle. Pour configurer cela correctement, vous devez spécifier les bonnes valeurs pour les paramètres pipeline_parallel_degree et processes_per_host. Le calcul simple est le suivant :

(pipeline_parallel_degree) x (data_parallel_degree) = processes_per_host

La bibliothèque se charge de calculer le nombre de réplicas du modèle (également appelé data_parallel_degree) en fonction des deux paramètres d'entrée que vous fournissez.

Par exemple, si vous définissez "pipeline_parallel_degree": 2 et "processes_per_host": 8 pour utiliser une instance ML avec huit workers GPU tels que "pipeline_parallel_degree": 2, la bibliothèque configure automatiquement le modèle distribué sur l'ensemble des GPU et un parallélisme de données à quatre voies. L'image suivante illustre la distribution d'un modèle sur les huit GPU avec un parallélisme des données à quatre voies et un parallélisme de pipeline bidirectionnel. Chaque réplica de modèle pour lequel nous le définissons comme un groupe de parallélisme de pipeline et lui appliquons l'étiquette PP_GROUP est partitionné sur deux GPU. Chaque partition du modèle est affectée à quatre GPU, et les quatre réplicas de partition se trouvent dans un groupe de parallélisme de données et sont étiquetés DP_GROUP. Sans parallélisme de tenseur, le groupe de parallélisme de pipeline est essentiellement le groupe de parallélisme de modèle.

Pour explorer le parallélisme de pipeline, consultez Principales fonctions de la bibliothèque de parallélisme de modèles SageMaker.

Pour commencer à exécuter votre modèle à l'aide du parallélisme de pipeline, consultez Exécuter une tâche d'entraînement distribué SageMaker avec la bibliothèque de parallélisme de modèle SageMaker.

Parallélisme de tenseur (disponible pour PyTorch)

Le parallélisme de tenseur divise les couches individuelles, ou nn.Modules, entre les appareils, pour qu'elles soient exécutées en parallèle. La figure suivante illustre l'exemple le plus simple de la façon dont la bibliothèque divise un modèle à quatre couches pour obtenir un parallélisme de tenseur bidirectionnel ("tensor_parallel_degree": 2). Les couches de chaque réplica de modèle sont divisées et distribuées dans deux GPU. Dans cet exemple, la configuration du parallélisme de modèle inclut également "pipeline_parallel_degree": 1 et "ddp": True (utilise le package PyTorch DistributedDataParallel en arrière-plan), de sorte que le degré de parallélisme des données devient huit. La bibliothèque gère la communication entre les réplicas de modèles distribués par tenseur.

L'utilité de cette fonctionnalité réside dans le fait que vous pouvez sélectionner des couches spécifiques ou un sous-ensemble de couches pour appliquer le parallélisme de tenseur. Pour explorer le parallélisme de tenseur et d'autres fonctionnalités d'économie de mémoire pour PyTorch, et pour apprendre à définir une combinaison de parallélisme de pipeline et de tenseur, consultez Fonctions étendues de la bibliothèque de parallélisme de modèles SageMaker pour PyTorch.

Partitionnement de l'état de l'optimiseur (disponible pour PyTorch)

Pour comprendre comment la bibliothèque effectue le partitionnement de l'état de l'optimiseur, envisagez un exemple de modèle simple à quatre couches. L'idée clé de l'optimisation du partitionnement de l'état est que vous n'avez pas besoin de répliquer l'état de votre optimiseur dans tous vos GPU. Au lieu de cela, un seul réplica de l'état de l'optimiseur est partitionné entre les rangs parallèles de données, sans redondance entre les appareils. Par exemple, le GPU 0 conserve l'état de l'optimiseur pour la couche 1, le GPU 1 suivant contient l'état de l'optimiseur pour la couche 2 (L2), etc. La figure animée suivante montre une propagation vers l'arrière avec la technique de partitionnement de l'état de l'optimiseur. À la fin de la propagation vers l'arrière, il y a le temps de calcul et de réseau pour l'opération optimizer apply (OA) pour mettre à jour les états de l'optimiseur et l'opération all-gather (AG) pour mettre à jour les paramètres du modèle pour la prochaine itération. Surtout, l'opération reduce peut chevaucher le calcul sur GPU 0, ce qui entraîne une propagation vers l'arrière plus rapide et plus efficace en termes de mémoire. Dans l'implémentation actuelle, les opérations AG et OA ne chevauchent pas compute. Cela peut entraîner un calcul étendu pendant l'opération AG, pouvant donner lieu à un compromis.

Pour plus d'informations sur l'utilisation de cette fonctionnalité, consultez Partitionnement de l'état de l'optimiseur.

Déchargement et points de contrôle d'activation (disponibles pour PyTorch)

Pour enregistrer la mémoire GPU, la bibliothèque prend en charge les points de contrôle d'activation afin d'éviter de stocker des activations internes dans la mémoire GPU pour les modules spécifiés par l'utilisateur pendant la transmission vers l'avant. La bibliothèque recalcule ces activations pendant la transmission vers l'arrière. En outre, la fonctionnalité de déchargement d'activation décharge les activations stockées dans la mémoire CPU et les récupère dans le GPU pendant la transmission vers l'arrière, afin de réduire encore l'empreinte mémoire d'activation. Pour plus d'informations sur l'utilisation de ces fonctionnalités, consultez Points de contrôle d'activation et Déchargement de l'activation.

Choix des techniques appropriées pour votre modèle

Pour plus d'informations sur le choix des techniques et configurations appropriées, consultez Bonnes pratiques concernant le parallélisme des modèles distribués SageMaker et Conseils et pièges relatifs à la configuration.