La frugalité rencontre la précision : formation rentable des modèles GPT NeoX et Pythia avec AWS Trainium | Services Web Amazon

Les grands modèles de langage (ou LLM) sont devenus un sujet de conversations quotidiennes. Leur adoption rapide est évidente par le temps nécessaire pour atteindre 100 millions d'utilisateurs, qui est passé de « 4.5 ans par Facebook » à un minimum historique de seulement « 2 mois par ChatGPT ». Un transformateur génératif pré-entraîné (GPT) utilise des mises à jour causales autorégressives pour effectuer des prédictions. Une variété de tâches telles que la reconnaissance vocale, la génération de texte et la réponse aux questions se révèlent avoir des performances extraordinaires grâce à ces architectures de modèles. Plusieurs modèles récents tels que NeoX, Falcon, Flamme utilisez l’architecture GPT comme épine dorsale. La formation des LLM nécessite un temps de calcul colossal, qui coûte des millions de dollars. Dans cet article, nous résumerons la procédure de formation de GPT NeoX on Formation AWS, un accélérateur d'apprentissage automatique (ML) spécialement conçu et optimisé pour la formation en apprentissage profond. Nous expliquerons comment nous avons formé de tels modèles de manière rentable (3.2 millions de jetons/$) avec AWS Trainium sans perdre la qualité du modèle.

Vue d'ensemble de la solution

Modèles GPT NeoX et Pythia

GPT NeoX et les Pythia sont les modèles de langage causal open source d'Eleuther-AI avec environ 20 milliards de paramètres dans NeoX et 6.9 milliards dans Pythia. Les deux sont des modèles de décodeurs suivant une conception architecturale similaire à celle du Chat GPT3. Cependant, ils disposent également de plusieurs ajouts, qui sont également largement adoptés dans les modèles récents tels que Llama. En particulier, ils ont une intégration positionnelle en rotation (ROPE) avec une rotation partielle sur les dimensions de la tête. Les modèles originaux (NeoX et Pythia 6.9B) sont formés sur des outils disponibles gratuitement Jeu de données de pieux avec déduplication et utilisation du backend Megatron et Deepspeed.

Nous démontrons la pré-formation et le réglage fin de ces modèles sur les instances Trn1 basées sur AWS Trainium à l'aide Neurone NeMo bibliothèque. Pour établir la preuve de concept et la reproduction rapide, nous utiliserons un sous-ensemble de données Wikipédia plus petit tokenisé à l'aide du tokenizer GPT2 Byte-pair Encoding (BPE).

Procédure pas à pas

Téléchargez l'ensemble de données Wikipédia pré-tokenisé comme indiqué :

export DATA_DIR=~/examples_datasets/gpt2 mkdir -p ${DATA_DIR} && cd ${DATA_DIR} wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json
wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.bin . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/my-gpt2_text_document.idx . --no-sign-request
aws s3 cp s3://neuron-s3/training_datasets/gpt/wikipedia/license.txt . --no-sign-request

NeoX 20B et Pythia 6.9B utilisent tous deux ROPE avec rotation partielle, par exemple, faisant pivoter 25 % des dimensions de la tête et gardant le reste sans rotation. Pour implémenter efficacement la rotation partielle sur l'accélérateur AWS Trainium, au lieu de concaténer les dimensions rotatives et non rotatives, nous ajoutons des fréquences nulles pour les dimensions non rotatives, puis faisons pivoter l'ensemble complet des dimensions de la tête. Cette astuce simple nous a aidé à améliorer le débit (séquences traitées par seconde) sur AWS Trainium.

Étapes de formation

Pour exécuter la formation, nous utilisons Amazon Elastic Compute Cloud multi-nœuds géré par SLURM (Amazon EC2) Cluster Trn1, avec chaque nœud contenant une instance trn1.32xl. Chaque trn1.32xl dispose de 16 accélérateurs avec deux travailleurs par accélérateur. Après avoir téléchargé la dernière Neurone NeMo paquet, utilisez le néox et les la pythie scripts de pré-entraînement et de réglage fin avec des hyper-paramètres optimisés et exécutez ce qui suit pour une formation à quatre nœuds.

1. Compiler : précompilez le modèle avec trois itérations d'entraînement pour générer et enregistrer les graphiques :

   sbatch --nodes 4 compile.slurm ./neoX_20B_slurm.sh

2. Exécuter : exécutez la formation en chargeant les graphiques mis en cache dès les premières étapes

   sbatch --nodes 4 run.slurm ./neoX_20B_slurm.sh

3. Surveiller les résultats

   tensorboard --logdir=nemo_experiments/megatron_neox

Les mêmes étapes doivent être suivies pour exécuter le modèle Pythia 6.9B en remplaçant neox_20B_slurm.sh by pythia_6.9B_slurm.sh.

Expériences de pré-formation et de mise au point

Nous démontrons la pré-formation des modèles GPT-NeoX et Pythia sur AWS Trainium en utilisant Neurone NeMo bibliothèque pour 10 1 itérations, et montre également le réglage fin de ces modèles pour 2 XNUMX étapes. Pour la pré-formation, nous utilisons le tokenizer GPTXNUMX BPE à l'intérieur du NeMo et suivons la même chose config tel qu'utilisé dans le modèle original. Le réglage fin sur AWS Trainium nécessite la modification de quelques paramètres (tels que facteur de division de la taille du vocabulaire), qui sont fournis dans les scripts de réglage fin pour tenir compte des différences Megatron par rapport à NeMo et des modifications GPU par rapport à AWS Trainium. Le débit de formation distribué multi-nœuds avec un nombre variable de nœuds est indiqué dans le tableau 1.

Tableau 1. Comparaison du débit moyen des modèles GPT NeoX et Pythia pour une formation jusqu'à 500 étapes avec un nombre changeant de nœuds. Le prix du trn1.32xl est basé sur le taux horaire effectif réservé sur 3 ans.

Ensuite, nous évaluons également la trajectoire de perte de la formation du modèle sur AWS Trainium et la comparons avec l'exécution correspondante sur un cluster P4d (cœurs GPU Nvidia A100). Parallèlement à la perte d'entraînement, nous comparons également des indicateurs utiles tels que la norme de gradient, qui est la norme à 2 normes des gradients du modèle calculés à chaque itération d'entraînement pour suivre la progression de l'entraînement. Les résultats de la formation sont présentés dans les figures 1 et 2 et le réglage fin de NeoX 20B dans la figure 3.

Figure 1. Perte de formation moyenne pour tous les travailleurs (à gauche) et norme de gradient (à droite) à chaque étape de formation. NeoX 20B est formé sur 4 nœuds avec un petit ensemble de données wiki sur GPU et Trainium avec les mêmes hyper-paramètres de formation (taille globale du lot = 256). Le GPU utilise BF16 et la précision mixte par défaut, tandis qu'AWS Trainium utilise le BF16 complet avec arrondi stochastique. Les trajectoires des normes de perte et de gradient correspondent pour le GPU et AWS Trainium.

Figure 2. Perte de formation moyenne pour tous les travailleurs (à gauche) et norme de gradient (à droite) à chaque étape de formation. Semblable à GPT NeoX dans la figure 1, Pythia 6.9B est formé sur 4 nœuds avec un petit ensemble de données wiki sur GPU et Trainium avec les mêmes hyper-paramètres de formation (taille globale du lot = 256). Les trajectoires des normes de perte et de gradient correspondent pour le GPU et Trainium.

Figure 3. Ajustement du modèle GPT NeoX 20B sur GPU et AWS Trainium avec perte de formation moyenne pour tous les travailleurs (à gauche) et norme de gradient (à droite). Un petit ensemble de données wiki est utilisé pour affiner la démonstration. Les trajectoires des normes de perte et de gradient correspondent pour le GPU et AWS Trainium.

Dans cet article, nous avons montré une formation rentable des LLM sur le matériel d'apprentissage en profondeur AWS. Nous avons formé les modèles GPT NeoX 20B et Pythia 6.9B sur AWS Trn1 avec la bibliothèque Neuron NeMo. Le coût de débit normalisé pour 20 milliards de modèles avec AWS Trainium est d'environ 3.2 millions de jetons/$ dépensé. Parallèlement à une formation rentable sur AWS Trainium, nous obtenons une précision de modèle similaire, comme en témoignent la perte d'étape de formation et la trajectoire de la norme de gradient. Nous avons également affiné les points de contrôle disponibles pour le modèle NeoX 20B sur AWS Trainium. Pour plus d'informations sur la formation distribuée avec NeMo Megatron sur AWS Trainium, consultez Référence des neurones AWS pour NeMo Megatron. Une bonne ressource pour commencer à peaufiner le modèle Llama peut être trouvée ici, Mise au point de Llama2. Pour démarrer avec AWS Trainium géré sur Amazon Sage Maker, Voir Entraînez vos modèles ML avec AWS Trainium et Amazon SageMaker.

À propos des auteurs

Gaurav Gupta est actuellement scientifique appliqué dans les laboratoires d'IA d'Amazon Web Services (AWS). Le Dr Gupta a obtenu son doctorat à l'USC Viterbi. Ses intérêts de recherche couvrent le domaine de la modélisation de données séquentielles, l'apprentissage des équations aux dérivées partielles, la théorie de l'information pour l'apprentissage automatique, les modèles dynamiques fractionnaires et les réseaux complexes. Il travaille actuellement sur des problèmes appliqués et mathématiques sur le comportement d'entraînement des LLM, les modèles de vision avec les PDE et les modèles multimodaux de théorie de l'information. Le Dr Gupta a publié des publications dans des revues/conférences de premier plan telles que Neurips, ICLR, ICML, Nature, IEEE Control Society, ACM cyber-physical Society.

Ben Snyder est un scientifique appliqué chez AWS Deep Learning. Ses intérêts de recherche comprennent les modèles fondamentaux, l'apprentissage par renforcement et l'optimisation asynchrone. En dehors du travail, il aime faire du vélo et du camping dans l’arrière-pays.

Amith (R) Mamidala est l'ingénieur principal des applications d'apprentissage automatique chez AWS Annapurna Labs. Le Dr Mamidala a obtenu son doctorat à l'Ohio State University en calcul haute performance et en communication. Au cours de son mandat chez IBM Research, le Dr Mamidala a contribué à la classe d'ordinateurs BlueGene qui a souvent dominé le classement Top500 des supercalculateurs les plus puissants et les plus économes en énergie. Le projet a reçu la médaille nationale 2009 de la technologie et de l'innovation. Après un bref passage en tant qu'ingénieur en IA dans un fonds spéculatif financier, le Dr Mamidala a rejoint les laboratoires de l'Annapurna en se concentrant sur la formation de modèles Large Language.

Jun (Luke) Huan est l'un des scientifiques principaux d'AWS AI Labs. Le Dr Huan travaille sur l'IA et la science des données. Il a publié plus de 180 articles évalués par des pairs dans des conférences et revues de premier plan. Il a reçu le prix NSF pour le développement de carrière en début de carrière en 2009. Avant de rejoindre AWS, il a travaillé chez Baidu Research en tant que scientifique distingué et directeur du Baidu Big Data Laboratory. Il a fondé StylingAI Inc., une start-up d'IA, et a travaillé en tant que PDG et scientifique en chef de 2019 à 2021. Avant de rejoindre l'industrie, il était professeur Charles E. et Mary Jane Spahr au département EECS de l'Université du Kansas.

Shruti Koparkar est Senior Product Marketing Manager chez AWS. Elle aide les clients à explorer, évaluer et adopter l'infrastructure de calcul accéléré Amazon EC2 pour leurs besoins en machine learning.

• Contenu propulsé par le référencement et distribution de relations publiques. Soyez amplifié aujourd'hui.
• PlatoData.Network Ai générative verticale. Autonomisez-vous. Accéder ici.
• PlatoAiStream. Intelligence Web3. Connaissance Amplifiée. Accéder ici.
• PlatonESG. Carbone, Technologie propre, Énergie, Environnement, Solaire, La gestion des déchets. Accéder ici.
• PlatoHealth. Veille biotechnologique et essais cliniques. Accéder ici.
• La source: https://aws.amazon.com/blogs/machine-learning/frugality-meets-accuracy-cost-efficient-training-of-gpt-neox-and-pythia-models-with-aws-trainium/