Concepts de parallélisme des modèles - Amazon SageMaker AI

Concepts de parallélisme des modèles

Le parallélisme des modèles est une méthode d’entraînement distribué dans laquelle le modèle de deep learning (DL) est partitionné sur plusieurs GPU et instances. La bibliothèque SageMaker de parallélisme des modèles v2 (SMP v2) est compatible avec les API et fonctionnalités natives PyTorch. Cela vous permet d’adapter facilement votre script d’entraînement PyTorch FSDP (Fully Sharded Data Parallel) à la plateforme d’entraînement SageMaker, et de tirer parti de l’amélioration des performances apportée par SMP v2. Cette page de présentation fournit un aperçu global du parallélisme des modèles et une description de la façon dont il peut aider à surmonter les problèmes liés à l’entraînement de modèles de deep learning (DL) qui sont généralement de très grande taille. Elle fournit également des exemples de ce que propose la bibliothèque SageMaker de parallélisme des modèles pour aider à gérer les stratégies de parallélisme des modèles 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.

Afin de surmonter le problème de limites associées à l’entraînement d’un modèle sur un GPU unique, SageMaker fournit la bibliothèque de parallélisme des modèles 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 bénéficier d’un entraînement distribué optimisé à l’aide de dispositifs compatibles 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 des modèles

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 float16 (FP16) ou bfloat16 (BF16) 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 ou BF16 ~ 2 octets

  • Un gradient FP16 ou BF16 ~ 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) disponible sur un GPU individuel. 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 P4 et P5 équipées de NVIDIA A100 et H100 Tensor Core, respectivement. 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. Pour connaître les types d’instances pris en charge par SMP v2, consultez Types d’instance pris en charge.

Même avec les instances de calcul accéléré, il est évident que les 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 dispositif GPU.

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

La bibliothèque comprend différents types de fonctionnalités de parallélisme de modèle et de 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 composés de centaines de milliards de paramètres.

Parallélisme des données partitionnées

Le parallélisme de données partitionnées est une technique d'entraînement distribué qui permet d'économiser de la mémoire et qui divise l'état d'un modèle (paramètres du modèle, gradients et états de l'optimiseur) entre les GPU dans un groupe de parallélisme de données.

SMP v2 implémente le parallélisme partitionné des données via FSDP, et l’étend pour implémenter la stratégie de partitionnement hybride sensible à l’échelle décrite dans le billet de blog Near-linear scaling of gigantic-model training on AWS.

Vous pouvez appliquer le parallélisme partitionné des données à votre modèle en tant que stratégie autonome. De plus, si vous utilisez les instances de GPU les plus performantes équipées de GPU NVIDIA A100 Tensor Core, ml.p4d.24xlarge et ml.p4de.24xlarge, vous pouvez bénéficier d’une vitesse d’entraînement améliorée grâce à l’opération AllGather proposée par la bibliothèque SageMaker AI de parallélisme des données (SMDDP).

Pour explorer davantage le parallélisme partitionné des données et découvrir comment le configurer ou utiliser une combinaison du parallélisme partitionné des données et d’autres techniques telles que le parallélisme de tenseur et l’entraînement de précision mixte, consultez Parallélisme partitionné des données hybrides.

Parallélisme expert

SMP v2 s’intègre à NVIDIA Megatron pour implémenter le parallélisme expert en plus de sa prise en charge des API PyTorch FSDP natives. Vous pouvez conserver votre code d’entraînement PyTorch FSDP tel quel et appliquer le parallélisme expert SMP pour entraîner des modèles Mixture of Experts (MoE) dans SageMaker AI.

Un modèle MoE est un type de modèle de transformeur composé de plusieurs experts, chacun constitué d’un réseau neuronal, généralement un réseau à propagation avant (FFN). Un réseau de portes appelé routeur détermine quels jetons sont envoyés à quel expert. Ces experts sont spécialisés dans le traitement d’aspects spécifiques des données d’entrée, ce qui permet au modèle de s’entraîner plus rapidement, de réduire les coûts de calcul, tout en obtenant la même qualité de performance que le modèle dense équivalent. Le parallélisme expert est une technique de parallélisme qui gère la division des experts d’un modèle MoE entre différents dispositifs GPU.

Pour découvrir comment entraîner des modèles MoE avec SMP v2, consultez Parallélisme expert.

Parallélisme de tenseur

Le parallélisme de tenseur divise les couches individuelles, ou nn.Modules, entre les dispositifs pour une exécution en parallèle. La figure suivante illustre l’exemple le plus simple de la façon dont la bibliothèque SMP divise un modèle à quatre couches pour obtenir un parallélisme de tenseur bidirectionnel ("tensor_parallel_degree": 2). Dans la figure suivante, les notations pour le groupe de parallélisme des modèles, le groupe de parallélisme de tenseur et le groupe de parallélisme des données sont respectivement MP_GROUP, TP_GROUP et DP_GROUP. Les couches de chaque réplica de modèle sont divisées et distribuées dans deux GPU. La bibliothèque gère la communication entre les réplicas de modèles distribués par tenseur.

Exemple le plus simple de la façon dont la bibliothèque SMP divise un modèle à quatre couches pour obtenir un parallélisme de tenseur bidirectionnel ("tensor_parallel_degree": 2).

Pour explorer le parallélisme de tenseur et d’autres caractéristiques d’économie de mémoire pour PyTorch, et découvrir comment définir une combinaison de caractéristiques de base, consultez Parallélisme de tenseur.

Points de contrôle d’activation et déchargement

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, avec le déchargement d’activation, elle décharge les activations stockées dans la mémoire du 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 de l’activation. Pour plus d’informations sur l’utilisation de ces caractéristiques, consultez Points de contrôle d'activation et Déchargement d’activation.

Choix des techniques appropriées pour votre modèle

Pour plus d’informations sur le choix des bonnes techniques et configurations, consultez Bonnes pratiques concernant le parallélisme distribué des modèles SageMaker.