Entraînez des modèles gigantesques avec une mise à l'échelle quasi linéaire à l'aide du parallélisme des données partitionnées sur Amazon SageMaker

Dans la poursuite d'une précision supérieure, les modèles d'apprentissage en profondeur dans des domaines tels que le traitement du langage naturel et la vision par ordinateur ont considérablement augmenté en taille au cours des dernières années, comptant souvent des dizaines à des centaines de milliards de paramètres. La formation de ces modèles gigantesques est difficile et nécessite des stratégies de distribution complexes. Les scientifiques des données et les ingénieurs en apprentissage automatique recherchent constamment le meilleur moyen d'optimiser leur calcul de formation, mais sont aux prises avec les frais généraux de communication qui peuvent augmenter avec la taille globale du cluster.

C'est pourquoi nous avons récemment lancé sparallélisme de données renforcé on Amazon Sage Maker, une nouvelle technique d'entraînement distribué économe en mémoire dans le Bibliothèque parallèle de modèles SageMaker (SMP). Le parallélisme des données partagées est spécialement conçu pour les modèles à grande échelle et utilise Amazon en interne MiCS technologie sous le capot, un effort scientifique pour minimiser l'échelle de communication en réduisant les frais généraux de communication coûteux enracinés dans la collecte de paramètres et la synchronisation des gradients. Avec un modèle GPT-30 à paramètre 2B avec une longueur de séquence de 2048, cette nouvelle fonctionnalité a atteint 141 TFLOP, soit une accélération de 39.7 % par rapport à DeepSpeed ​​ZeRO-3. Pour un modèle 10B GPT-2 avec une longueur de séquence de 512, cette nouvelle fonctionnalité a également atteint 564 échantillons par seconde, soit une accélération de 13.9 % par rapport au FSDP (Fully Sharded Data Parallel) de PyTorch. N'oubliez pas que dans la formation de modèles gigantesques, chaque pourcentage d'accélération se traduit par des économies en dollars et une productivité gagnée dans votre équipe.

Dans cet article de blog, nous allons d'abord examiner de plus près les principaux différenciateurs du parallélisme des données partitionnées et quand l'utiliser. Ensuite, vous apprendrez à former facilement un modèle GPT-30 à paramètre 2B sur SageMaker grâce à cette nouvelle fonctionnalité. Enfin, nous comparerons les performances avec d'autres options open source, notamment en surpassant DeepSpeed ​​ZeRO jusqu'à 39.7 % sur 256 GPU.

Comment fonctionne le parallélisme des données partitionnées et quand l'utiliser

Avant d'introduire le parallélisme des données fragmentées, examinons sa famille de techniques plus large. Les approches récentes de formation distribuée pour les grands modèles sont passées à un paradigme où les paramètres du modèle, les gradients et les états de l'optimiseur sont partagés entre les nœuds parallèles aux données. Contrairement au parallélisme de pipeline, qui a la complexité innée de choisir des couches à partitionner entre les appareils, en particulier lorsque votre framework ne prend pas en charge fractionnement de modèle automatisé, ce paradigme préserve élégamment la simplicité du parallélisme des données, tout en supprimant la contrainte du parallélisme des données où un modèle doit tenir dans un seul GPU.

Dans les frameworks existants qui relèvent de ce paradigme, notamment DeepSpeed ​​ZeRO-3 et le FSDP de PyTorch en amont de FairScale, les états du modèle sont répartis entre TOUTE GPU, une stratégie qui réduit la consommation de mémoire sur chaque GPU au prix d'une surcharge de communication importante qui augmente avec la taille du cluster et entraîne donc une baisse significative de l'évolutivité à grande échelle. En revanche, le parallélisme des données partitionnées dans les états du modèle de partitions de la bibliothèque SMP dans un sensible à l'échelle manière en partitionnant chaque réplique des états du modèle uniquement dans un sous-ensemble de GPU.

Regardons de plus près le partitionnement de modèle sensible à l'échelle dans MiCS, la technologie de base derrière le parallèle de données partagées. L'intuition derrière cette conception est que le partitionnement des états d'entraînement sur l'ensemble du groupe parallèle de données peut ne pas être nécessaire pour entraîner un modèle avec des dizaines de milliards de paramètres. Par exemple, 8 GPU V100 (32 Go chacun) suffisent pour contenir la réplique des états du modèle d'un modèle à 10B paramètres qui nécessite environ 200 Go de mémoire lors de l'entraînement avec l'optimiseur Adam en précision mixte. En limitant une réplique complète des états du modèle dans le plus petits au plus sous-ensemble de GPU, nous pouvons réduire efficacement l'ampleur des frais généraux de communication par rapport à DeepSpeed ​​et PyTorch FSDP. Le parallèle de données fragmentées exploite également d'autres techniques dans MiCS telles que la communication hiérarchique et la synchronisation de gradient à 2 sauts. Pour plus d'informations, consultez Mise à l'échelle quasi linéaire de la formation de modèles gigantesques sur AWS or MiCS : Mise à l'échelle quasi-linéaire pour la formation d'un modèle gigantesque sur le cloud public.

Maintenant, comment savez-vous quand choisir les données partitionnées en parallèle plutôt que d'autres techniques de formation distribuées ? La règle générale est que si votre modèle a moins d'un milliard de paramètres et peut tenir dans la mémoire du GPU, bibliothèque parallèle de données SageMaker or Compilateur de formation SageMaker peut vous suffire. Si vous avez de plus grands modèles de langage ou de vision par ordinateur, nous vous suggérons de l'entraîner avec la technique de parallélisme de données fragmentées combinée à point de contrôle d'activation ainsi que déchargement d'activation dans la bibliothèque parallèle de modèles SageMaker d'abord, avant d'autres techniques telles que parallélisme tenseur ou parallélisme de pipeline.

Utilisation du parallélisme des données partagées pour former GPT-2 sur Amazon SageMaker

Apprenons maintenant à former un modèle GPT-2 avec des données partagées en parallèle, avec SMP encapsulant la complexité pour vous. Cette cahier de tuto complet vous guide tout au long du processus, depuis le traitement des données, la définition et la soumission des tâches de formation, jusqu'à la surveillance des journaux de formation. Ce qui suit est un bref aperçu mettant en évidence les étapes clés de l'utilisation de cette fonctionnalité.

1. Commencer

Le parallélisme des données partagées est disponible dans PyTorch v1.12.0+ et fonctionne avec FP16 et BF16. Le moyen le plus simple d'utiliser la bibliothèque SMP consiste à utiliser un AWS Deep Learning Container pour PyTorch prédéfini. Cependant, si vous souhaitez apporter votre propre conteneur Docker, vous pouvez vous référer à Créez votre propre conteneur Docker avec la bibliothèque parallèle de modèles distribués SageMaker. Pour commencer, suivez Modifier un script de formation PyTorch pour adapter les API des SMP dans votre script de formation. Dans cette section, nous n'appelons que quelques étapes principales avec des extraits de code du script de formation prêt à l'emploi train_gpt_simple.py. Vous pouvez suivre les commentaires dans le script et Document API pour en savoir plus sur les endroits où les API SMP sont utilisées.

Tout d'abord, importez et initialisez la bibliothèque en appelant smdistributed.modelparallel.torch.init() au début du script de formation :

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Ensuite, enveloppez le modèle à partitionner avec smdistributed.modelparallel.torch.DistributedModel et utiliser le retour DistributedModel objet pour aller de l'avant :

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Enveloppez l'optimiseur avec smdistributed.modelparallel.torch.DistributedOptimizer pour enregistrer et charger les états de l'optimiseur.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Mettez la logique avant et arrière dans une fonction pas à pas et décorez-la avec smdistributed.modelparallel.torch.step.  Tout calcul défini à l'intérieur du smp.step-decorated fonction est exécutée de manière distribuée.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Préparer le jeu de données

Nous utilisons les texte Web ouvert est le jeu de données que nous utilisons dans cet exemple. Le notebook utilise le script data_prep_512.py pour télécharger et prétraiter l'ensemble de données. Vous pouvez également vous entraîner avec d'autres ensembles de données en modifiant data_pipeline.py. Lorsque vous traitez un ensemble de données et un modèle volumineux, vous pouvez accélérer le travail d'entraînement en utilisant les données stockées dans Amazon FSx pour Lustre, qui fournit un système de fichiers hautes performances intégré nativement avec Service de stockage simple Amazon (S3). Veuillez consulter les instructions de Configurer le canal d'entrée de données pour utiliser Amazon FSx pour Lustre pour obtenir des conseils sur la configuration d'un système de fichiers FSx Lustre comme canal d'entrée de données.

3. Démarrez les travaux de formation

Cette étape suppose que vous avez déjà modifié votre script de formation et préparé l'ensemble de données comme mentionné dans les sections précédentes. À activer le parallélisme des données partagées, définissez simplement le sharded_data_parallel_degree dans l' Estimateur PyTorch. Dans ce tutoriel, nous définissons sharded_data_parallel_degree=128 ainsi que  instace_count=32 pour les nœuds p4d.24xlarge, ce qui indique que les états du modèle seront répartis sur 128 GPU parmi les 256 GPU au total. Sur la base de cette valeur sélectionnée, SMP définira alors automatiquement le degré de parallèle des données sur 2 (car 256/128=2), ce qui signifie que nous aurons deux répliques pour le parallélisme des données. Une règle générale pour choisir une valeur idéale pour sharded_data_parallel_degree consiste à ajouter un nœud supplémentaire au groupe de partage pour chaque 3 B de paramètres de modèle. Dans ce didacticiel, la taille de notre modèle est de 30 B, nous devons donc utiliser au moins 10 nœuds pour le partitionnement. Et parce que 16 nœuds (128 GPU) est la plus petite puissance de 2 au-dessus du seuil, nous avons défini sharded_data_parallel_degree=128.

Pour les points de contrôle, nous fournissons également un ensemble d'utilitaires de points de contrôle dans sharded_data_parallel_checkpoint.py , y compris un utilitaire pour reconstruire l'intégralité state_dict pour les cas d'utilisation avancés. Enfin, nous pouvons lancer une tâche de formation distribuée en appelant fit() sur l'Estimator.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Surveillez les tâches de formation

Vous pouvez accéder aux journaux d'entraînement et suivre l'utilisation du GPU et de la mémoire sur Amazon Cloud Watch. Assurez-vous de regarder les journaux de "algo-1" car il s'agit du nœud principal dont le flux de sortie contient les journaux de tâches d'entraînement de toutes les instances.

Analyse comparative des performances

Nous avons comparé le parallélisme des données fragmentées dans la bibliothèque SMP sur 16 et 32 ​​nœuds p4d.24xlarge pour une longueur de séquence de 512 et 2048, respectivement. Le modèle GPT30 à paramètre 2B est configuré pour utiliser une largeur cachée de 7168, 48 couches et 64 têtes. Vous pouvez adopter exactement la même configuration où la longueur de la séquence est de 2048 en définissant model_config = "gpt2-30b" dans le cahier du tutoriel. Avec ce paramètre, SMP a atteint 73.52 échantillons par seconde, soit une accélération de 39.7 % par rapport à DeepSpeed ​​ZeRO-3. Si votre taille de jeton est de 500 milliards, cette accélération signifie près de 367 heures d'économies sur les nœuds p4d.24xlarge, soit l'équivalent de plus de 12,000 XNUMX $ de budget économisé par formation ! Le tableau suivant résume nos résultats de référence.

1/ Pour chaque configuration de modèle, nous avons testé différentes fonctionnalités, étapes et configurations dans DeepSpeed ​​ZeRO et avons choisi celle qui offre le meilleur débit comme ligne de base DeepSpeed. Le benchmark a été exécuté sur Cloud de calcul élastique Amazon (Amazon EC2). 2/ Ces résultats s'appuient sur des collectifs de communication améliorés optimisés pour AWS qui seront bientôt disponibles. 3/ Le temps de formation est projeté à partir de la vitesse basée sur le nombre de jetons traités.

En résumé, nous avons observé un débit systématiquement plus élevé avec le parallélisme des données partagées dans SMP par rapport à DeepSpeed ​​sur une gamme de modèles et de configurations. Cette nouvelle fonctionnalité a également démontré une meilleure efficacité de la mémoire par rapport à DeepSpeed, permettant à SMP de s'adapter à une taille de lot plus grande et de réduire le niveau d'accumulation de gradient requis pour s'adapter à une taille de lot globale particulière.

Conclusion

Dans cet article, nous avons présenté une nouvelle technique de formation distribuée - le parallélisme des données partagées - et comment elle accélère la formation de modèles gigantesques avec une mise à l'échelle quasi linéaire sur Amazon SageMaker. Nous avons également expliqué comment former un modèle GPT-2 avec la nouvelle technique suivant ce exemple complet. Vous pouvez suivre le Dépôt GitHub d'exemples Amazon SageMaker pour suivre tous les exemples parallèles de modèles SageMaker ou assister à notre prochain ateliers de formation distribués. Pour en savoir plus sur le parallélisme des données partagées, veuillez consulter le Documentation.

À propos des auteurs

Emilie Webber a rejoint AWS juste après le lancement de SageMaker et essaie d'en parler au monde depuis ! En dehors de la création de nouvelles expériences ML pour les clients, Emily aime méditer et étudier le bouddhisme tibétain.

Can Karakus est un scientifique appliqué senior chez AWS, optimisant l'apprentissage en profondeur distribué à grande échelle sur AWS. Ses intérêts de recherche portent sur l'apprentissage en profondeur, l'optimisation distribuée, les systèmes distribués et la théorie de l'information. En dehors du travail, il aime faire du vélo, voyager, lire et apprendre.

Rahul Huilgol est ingénieur logiciel senior chez AWS. Il travaille sur des systèmes d'apprentissage en profondeur distribués, afin de rendre facile et performant l'entraînement de grands modèles d'apprentissage en profondeur dans le cloud. Dans ses temps libres, il aime la photographie, le vélo et le jardinage.

Suhit Kodgule est un ingénieur en développement logiciel au sein du groupe d'intelligence artificielle d'AWS travaillant sur des cadres d'apprentissage en profondeur. Dans ses temps libres, il aime faire de la randonnée, voyager et cuisiner.

Erin Ho est chef de produit pour AWS Deep Learning. Elle travaille sur des produits qui permettent aux clients de former plus facilement des modèles d'apprentissage en profondeur sur AWS. Pour s'amuser en dehors du travail, elle aime faire de la randonnée et du ski.