Comment Amazon Search M5 a économisé 30 % sur les coûts de formation LLM en utilisant AWS Trainium | Services Web Amazon

Depuis des décennies, Amazon est un pionnier et innove en matière d'apprentissage automatique (ML), offrant des expériences agréables à ses clients. Dès les premiers jours, Amazon a utilisé le ML pour divers cas d'utilisation tels que les recommandations de livres, la recherche et la détection des fraudes. À l'instar du reste du secteur, les progrès du matériel accéléré ont permis aux équipes d'Amazon de poursuivre des architectures modèles utilisant les réseaux neuronaux et l'apprentissage profond (DL).

Le programme M5 d'Amazon Search possède la stratégie d'apprentissage par découverte pour Amazon et crée des modèles à grande échelle multilingues, multilocales, multi-entités, multitâches et multimodaux tels que du texte, des images et des vidéos. Le programme M5 propose des intégrations universelles et des modèles de base à grande échelle à des centaines d'équipes ML sur Amazon tout en maintenant des contrôles stricts sur l'optimisation des coûts. Pour y parvenir, l’équipe M5 évalue régulièrement de nouvelles techniques permettant de réduire les coûts.

Comme de nombreuses organisations de ML, les accélérateurs sont largement utilisés pour accélérer la formation et l'inférence DL. Lorsqu'AWS a lancé des accélérateurs spécialement conçus avec la première version de Inférence AWS en 2020, l'équipe M5 a rapidement commencé à utilisez-les pour déployer plus efficacement les charges de travail de production, ce qui permet d'économiser des coûts et de réduire la latence. L'année dernière, AWS a lancé son Formation AWS des accélérateurs qui optimisent les performances par coût pour le développement et la construction de modèles DL de nouvelle génération. Dans cet article, nous expliquons comment M5 a pu réduire le coût de formation de ses modèles de 30 % et partageons certaines des meilleures pratiques que nous avons apprises en cours de route.

Instances Trainium

Grâce aux progrès réalisés en matière d'accélérateurs spécialement conçus, Amazon propose également des accélérateurs convaincants sous la forme d'AWS Inferentia et Trainium. Comme leur nom l’indique, ces puces sont optimisées pour dépasser respectivement les besoins des charges de travail d’inférence et de formation. Pour la formation à grande échelle de modèles de fondation atteignant des milliards de paramètres, Trainium Instances Trn1 et Trn1n sont des choix idéaux en raison de leurs caractéristiques. Les instances Trn1 sont alimentées par l'état de l'art NeuronCore-v2, et disposent d'une grande quantité de calcul d'accélérateur et de mémoire. Les instances Trn1n peuvent également être choisies pour une plus grande quantité de bande passante réseau (1,600 XNUMX Go), et sont donc parfaitement adaptées à une formation performante dans un souci d'optimisation des coûts.

Pour utiliser des accélérateurs, vous avez besoin d’une couche logicielle pour les prendre en charge. Avec les puces Trn et Inf, le SDK AWS Neuron déverrouille les accélérateurs spécialement conçus par Amazon avec l'aide de PyTorch XLA. PyTorch XLA convertit le mode impatient de PyTorch en implémentation basée sur un graphique en mode paresseux. Ces graphiques sont ensuite utilisés et compilés pour être utilisés avec l'accélérateur. PyTorch Neuron (qui fait partie du SDK Neuron) permet aux utilisateurs de PyTorch d'entraîner leurs modèles sur Trainium NeuronCores avec quelques lignes de code.

Modèle et charge de travail

L'équipe M5 forme et déploie des modèles fondamentaux et des représentations universelles pour aider diverses équipes d'Amazon à apporter du plaisir à - clients. L'un de ces modèles est un modèle d'encodeur de texte suivi d'un perceptron multicouche (MLP) avec des interactions de fonctionnalités explicites ou implicites définies par l'architecture du réseau neuronal avec des centaines de millions de paramètres pouvant être entraînés. Ce modèle est formé sur des milliards de jetons et est utilisé pour générer des millions d'intégrations dans un paramètre d'inférence par lots hors ligne. Ces intégrations sont des entrées dans un service Amazon de niveau 1 destiné aux clients.

L'infrastructure pour le pipeline de production utilise Lot AWS comprenant stratégies de partage équitable des files d'attente, en utilisant un cluster trn1.32xlarge multi-nœuds compatible EFA comme calcul pour la formation du modèle. Sur le plan fonctionnel, le pipeline de production effectue une formation incrémentielle du modèle, une évaluation du modèle entraîné et une inférence par lots hors ligne sur le modèle entraîné, le tout en utilisant PyTorch comme bibliothèque DL sous-jacente.

Objectifs

Satisfaire nos clients est un principe primordial. Compte tenu de la nature du pipeline orienté client, il est essentiel que tous les accords de niveau de service (SLA) soient respectés sans régressions. Nous avons identifié deux critères d'acceptation critiques pour adapter notre pipeline de production de GPU existant et le faire passer à Trainium :

• Qualité du modèle – La qualité de nos modèles impacte directement l’expérience client. Nous exigeons qu'il y ait une différence de qualité de modèle inférieure à 0.1 % entre le GPU et Trainium.
• Débit de formation – Nous formons périodiquement nos modèles de manière itérative pour offrir l’expérience la plus récente à nos clients. Nous exigeons que la convergence des modèles soit atteinte dans un délai prédéfini (par exemple 1 semaine) pour respecter nos SLA de production.

Dans les sections suivantes, nous partageons notre parcours de travail à rebours à partir de ces critères et nos apprentissages pour prendre en charge les charges de travail de production à l'échelle d'Amazon.

Script de formation

Avant de commencer la formation du modèle, nous devons apporter des modifications au script de formation pour le rendre conforme à XLA. Compte tenu de la taille du modèle, nous utilisons des données distribuées parallèles (DDP) pour entraîner le modèle. DDP nous permet d'augmenter le débit de formation de modèles en augmentant le nombre de machines utilisées pour exécuter la formation de modèles, sans aucune modification du code. Nous avons suivi les instructions fournies dans le Tutoriel de formation Neuron PyTorch MLP pour ajouter des constructions spécifiques à XLA dans nos scripts de formation. Ces modifications de code sont simples à mettre en œuvre. Voici quelques enseignements techniques importants tirés de l’exercice qui ont considérablement amélioré le débit de notre modèle :

• Placement de xm.mark_step() - xm.mark_step() compile et exécute les graphiques de calcul collectés paresseusement. Invoquer mark_step trop de fois conduira à un plus grand nombre de petits graphiques, tandis que l'invoquer trop peu de fois conduira à des graphiques peu nombreux, mais de grande taille. En fonction de votre application, le débit et la mise en œuvre de votre formation de modèle varieront en fonction de votre placement de xm.mark_step(). Notre mise en œuvre en place un xm.mark_step() après une passe avant et arrière, et une après l'étape d'optimisation.
• Encapsulation du chargeur de données avec le chargeur de périphérique multitraitement XLA – Il s’agit d’une étape critique qui peut facilement être manquée. Le chargeur de périphérique multitraitement torch_xla.distributed.parallel_loader.MpDeviceLoader charge les données d'entraînement sur chaque appareil XLA avec des options de préchargement et de chevauchement du chargement des données avec les exécutions de l'appareil pour améliorer le débit. Le chargeur de périphérique appelle également xm.mark_step() et est donc capable de créer des graphiques pour le chargement des données sur l'appareil à partir de l'hôte.

Compilation pour Trainium

Traditionnellement, le cycle de développement de modèles avec GPU implique d'apporter des modifications au modèle ou au script de formation et de l'exécuter directement sur le périphérique GPU. Les accélérateurs tels que Trainium qui utilisent XLA nécessitent une étape supplémentaire avant que la formation du modèle puisse être exécutée sur l'accélérateur. Les graphiques de calcul XLA ne peuvent être exécutés qu’après avoir été compilés. Généralement, il existe deux manières d'effectuer cette compilation : Ahead of Time (AOT), où vous tracez et compilez d'abord tous les graphiques, puis les exécutez, ou Just In Time (JIT), où les graphiques sont tracés, compilés et exécutés au fur et à mesure. sont rencontrés. Le SDK Neuron fournit ces deux éléments prêts à l'emploi. En règle générale, la compilation AOT est effectuée en premier. Les graphiques sont ensuite exécutés après cette compilation. Si de nouveaux graphiques sont rencontrés, le runtime Neuron appelle une compilation JIT avant de les exécuter. Pour effectuer la compilation AOT, le SDK Neuron fournit neuron_parallel_compile, un utilitaire de compilation qui extrait les graphiques d'un essai du script de formation et effectue une compilation AOT parallèle.

Un aspect important de la compilation AOT est de garantir qu’aucun nouveau graphique de calcul n’est créé au cours de la formation. Une source de nouveaux graphiques de calcul (et donc de recompilations) réside dans les formes dynamiques des lots d'entraînement lors de l'entraînement du modèle. Nous avons constaté que l'utilisation de formes statiques et de lots de taille fixe élimine les compilations de temps de formation et améliore considérablement le débit de formation sans aucun effet sur la précision du modèle. En imposant de telles contraintes sur la formation, nous avons observé que seulement 4 à 5 étapes de formation du modèle, une étape de validation du modèle et un point de contrôle du modèle une fois sont nécessaires pour tracer tous les graphiques lors de la compilation AOT. Il est important de noter que le SDK Neuron évolue constamment et qu’à l’avenir, il prendra également en charge les formes dynamiques.

De plus, les graphiques compilés sont stockés dans le Cache persistant des neurones sur disque ou dans un Service de stockage simple Amazon (Amazon S3). Ceci est particulièrement utile pour les charges de travail de production où l’architecture du modèle et la configuration de la formation ne changent pas. Par conséquent, la surcharge de compilation n’est encourue qu’une seule fois. Utiliser le cache est aussi simple que de définir un indicateur d'environnement :

export NEURON_COMPILE_CACHE_URL="s3://BUCKET/KEY"

Le compilateur Neuron fournit également trois options d'optimisation au niveau du compilateur (O1, O2, O3) pour équilibrer le temps de compilation et le débit d’exécution du modèle. O1 permet des optimisations de base sur le graphe de calcul et minimise le temps de compilation, O3 offre un débit d'exécution de modèle amélioré au prix d'un temps de compilation plus élevé, et O2 (option par défaut) constitue un équilibre entre les deux. Pour notre cas d'utilisation, nous avons utilisé l'optimisation O1 et observé une réduction de 86 % du temps de compilation sans modification des métriques de précision du modèle, tout en observant une réduction d'environ 5 à 7 % du débit par rapport à l'optimisation par défaut (O2). Selon le cas d'utilisation, vous pouvez choisir différents niveaux d'optimisation.

Pour résumer, nous avons utilisé les indicateurs suivants pour la compilation :

NEURON_CC_FLAGS="--target trn1 --auto-cast all --auto-cast-type bf16 --model-type transformer --optlevel O1"

Compatibilité des points de contrôle

Une fois la compilation terminée avec succès, nous pouvons procéder à la formation de nos modèles sur Trainium. Comme mentionné précédemment, nous entraînons progressivement nos modèles, ce qui signifie que nous chargeons un point de contrôle de modèle précédemment formé et poursuivons l'entraînement avec de nouvelles données. PyTorch et PyTorch XLA permettent une transition transparente entre les accélérateurs grâce à l'interopérabilité des points de contrôle. La flexibilité de passer du GPU à Trainium nous a permis de charger de manière transparente le modèle de GPU précédent et de nous entraîner sur les machines Trainium. Cela était essentiel pour garantir que nous puissions initialiser notre modèle avec le meilleur modèle précédemment formé sans aucun temps d'arrêt de production ni perte de précision du modèle.

Étant donné que le modèle GPU a été enregistré à l'aide des utilitaires de sauvegarde de modèle PyTorch standard, nous avons pu utiliser l'utilitaire de chargement de point de contrôle PyTorch pour charger le modèle GPU sur les appareils Trainium.

Par exemple, sur GPU/CPU, vous pouvez enregistrer le modèle avec le code suivant :

torch.save(model.state_dict(), PATH)

Ensuite, vous chargez à nouveau le modèle sur Trainium :

import torch_xla.core.xla_model as xm
xla_device = xm.xla_device()
model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(xla_device)

De même, vous pouvez sauvegarder le modèle sur Trainium avec le code suivant :

import torch_xla.core.xla_model as xm
# automatically moves the data to CPU for the master device
xm.save(model.state_dict(), PATH) 

Et chargez à nouveau le modèle sur GPU/CPU :

model = MyModel(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # can be any device

En fait, comme nous utilisons DDP pour la formation du modèle, le chargement du modèle est indépendant du nombre de machines utilisées pour former le point de contrôle précédent. Cela nous permet de faire évoluer horizontalement la flotte Trn1 sans modification du code ni effet négatif sur la formation du modèle. Ces points de contrôle basés sur PyTorch peuvent être directement utilisés ou même scriptés au flambeau pour des cas d'utilisation d'inférence sur AWS Inferentia2 ou d'autres accélérateurs.

Stabilité opérationnelle

On ne soulignera jamais assez que l’exécution de charges de travail en production nécessite le respect de plusieurs SLA. Pour notre cas d'utilisation, outre les SLA sur la qualité du modèle et le débit de formation, il est impératif que le pipeline de production soit opérationnellement stable, ce qui signifie un minimum de temps d'arrêt et de perturbations pendant la formation, l'évaluation et l'inférence du modèle.

Comme pour le pipeline existant basé sur GPU, nous avons ajouté de nombreux mécanismes pour rendre le pipeline opérationnellement stable. Avant de commencer la formation du modèle, nous effectuons plusieurs tests d’intégrité pour évaluer la santé des machines. Ces tests comprennent généralement des opérations tensorielles simples pour vérifier la santé des dispositifs accélérateurs. Nous avons observé que pour la formation distribuée, il est important d’exécuter des tests pour vérifier également la communication collective entre les instances. Nous avons utilisé le Suite de tests NCCOM du SDK Neuron pour y parvenir, en exécutant diverses opérations telles que la collecte totale, la réduction totale et la réduction de la dispersion.

Même après avoir suivi les suggestions que nous avons mentionnées, nous avons observé que les problèmes transitoires sont inévitables dans tout pipeline, quel que soit l’accélérateur sous-jacent. Pour renforcer la résilience de tout pipeline de formation, nous vous recommandons d'intégrer des mécanismes de nouvelle tentative pour résoudre ces problèmes potentiels. Nous utilisons Nouvelles tentatives automatisées AWS Batch pour réessayer les tâches qui rencontrent un échec transitoire pendant la formation du modèle. Ces redémarrages peuvent être coûteux si une panne survient vers la fin de la formation. Pour contrer ce problème, nous avons adapté nos scripts de formation pour charger un point de contrôle de modèle préalablement formé et poursuivre la formation à partir de ce point. Grâce à cette fonctionnalité, nous sommes en mesure de redémarrer de manière agressive les tâches de formation ayant échoué avec un minimum de surcharge.

Grâce à ces mécanismes de résilience en place, nous avons pu atteindre des taux de réussite de 98.5 % pour nos charges de travail sur Trn1, comparables aux taux de réussite de nos pipelines GPU existants.

Résultats

Pour valider l'exactitude de nos modèles, nous avons initialisé deux modèles à partir du même point de contrôle GPU et formé l'un sur Trainium et l'autre sur un GPU comparable. Les deux modèles ont été formés avec les mêmes hyperparamètres de formation. L'ensemble de données utilisé pour le calcul des métriques est un ensemble de données d'exclusion, et nous évaluons la précision du modèle sur cet ensemble de données toutes les N étapes globales. L'axe X est l'étape globale et l'axe Y est la précision du modèle. Nous avons observé une différence de moins de 0.1 % dans la précision du modèle à chaque point du graphique suivant.

De plus, pour évaluer la rentabilité de la formation du modèle, nous préférons comparer le temps nécessaire pour atteindre la convergence du modèle. Nous pensons que cela fournit une vision plus pratique des économies de coûts par rapport à des mesures telles que le coût par jeton, le FLOPS/dollar obtenu et d'autres facteurs. Compte tenu du temps de formation de trn1.32xl et comparable Cloud de calcul élastique Amazon (Amazon EC2), nous avons observé que Trainium offre un coût jusqu'à 30 % moins cher pour modéliser la convergence.

Conclusion

De nombreux facteurs doivent être pris en compte lors de l’évaluation des différents accélérateurs pour vos charges de travail DL. Certains des plus importants sont la qualité du modèle, le débit, le coût et la disponibilité. Il est primordial de garantir que la qualité et le débit de votre modèle ne sont pas sacrifiés en fonction de l'accélérateur que vous choisissez.

Grâce à notre partenariat et à notre collaboration avec l'équipe Annapurna Neuron, l'équipe Amazon Search M5 a pu économiser jusqu'à 30 % en passant à Trainium. L'équipe est capable d'utiliser Trainium et d'atteindre une qualité de modèle et une parité de débit avec des accélérateurs comparables sur le marché. L'interopérabilité des points de contrôle et les modifications minimes du code avec la prise en charge de XLA ont permis à M5 de choisir entre plusieurs accélérateurs pour leurs charges de travail. Cela a permis à l'équipe M5 de tirer parti de la grande puissance de calcul de Trainium et de créer des solutions indépendantes des accélérateurs pour ravir les clients d'Amazon.com. D'un point de vue opérationnel, Trainium s'est avéré capable de prendre en charge les services de niveau 1 à l'échelle d'Amazon. L'équipe M5 continue de déplacer davantage de charges de travail vers Trainium afin de fournir les meilleurs modèles pour Amazon aux coûts les plus bas.

En résumé, l'équipe M5 a été en mesure d'effectuer une formation ML rentable et de niveau production en ajoutant Trainium à la flotte d'accélérateurs. Nous vous encourageons à jeter un œil à Trainium et à d'autres appareils Neuron comme AWS Inferentia pour profiter des avantages du silicium Amazon spécialement conçu pour les charges de travail ML. Démarrez facilement avec l'un des nombreux tutoriels présentant différents modèles, comme Lama 2, disponible sur Trainium.

À propos des auteurs

Patni abhinandais est ingénieur logiciel senior chez Amazon Search. Il se concentre sur la construction de systèmes et d'outils pour la formation évolutive en apprentissage profond distribué et l'inférence en temps réel.

James Park est architecte de solutions chez Amazon Web Services. Il travaille avec Amazon.com pour concevoir, créer et déployer des solutions technologiques sur AWS, et s'intéresse particulièrement à l'IA et à l'apprentissage automatique. Dans ses temps libres, il aime rechercher de nouvelles cultures, de nouvelles expériences et se tenir au courant des dernières tendances technologiques. Vous pouvez le retrouver sur LinkedIn.

Jerry Mannil est ingénieur logiciel chez Amazon Search. Il travaille à l'amélioration de l'efficacité, de la robustesse et de l'évolutivité de l'infrastructure de formation distribuée.

Ken Su est ingénieur logiciel chez Amazon Search. Il travaille à l'amélioration de l'efficacité de la formation et au flux de travail de formation distribué évolutif. En dehors du travail, il aime la randonnée et le tennis.

RJ est ingénieur au sein d'Amazon. Il construit et optimise des systèmes pour les systèmes distribués à des fins de formation et travaille à l'optimisation de l'adoption de systèmes afin de réduire la latence pour l'inférence ML. En dehors de son travail, il explore l’utilisation de l’IA générative pour créer des recettes alimentaires.

• 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/how-amazon-search-m5-saved-30-for-llm-training-cost-by-using-aws-trainium/