Zuinigheid ontmoet nauwkeurigheid: kostenefficiënte training van GPT NeoX- en Pythia-modellen met AWS Trainium | Amazon-webservices

Grote taalmodellen (of LLM's) zijn een onderwerp van dagelijkse gesprekken geworden. Hun snelle acceptatie blijkt uit de hoeveelheid tijd die nodig is om 100 miljoen gebruikers te bereiken, die is gestegen van “4.5 jaar via Facebook” naar een dieptepunt van slechts “2 maanden via ChatGPT.” Een generatieve, vooraf getrainde transformator (GPT) maakt gebruik van causale autoregressieve updates om voorspellingen te doen. Deze modelarchitecturen hebben aangetoond dat een verscheidenheid aan taken, zoals spraakherkenning, het genereren van tekst en het beantwoorden van vragen, verbazingwekkende prestaties leveren. Diverse recente modellen zoals NeoX, valk, Lama gebruik de GPT-architectuur als ruggengraat. Het trainen van LLM's vereist een enorme hoeveelheid rekentijd, wat miljoenen dollars kost. In dit bericht vatten we de trainingsprocedure van GPT samen NeoX on AWS Trainium, een speciaal gebouwde machine learning (ML)-versneller die is geoptimaliseerd voor deep learning-training. We zullen schetsen hoe we dergelijke modellen kosteneffectief (3.2 miljoen tokens/$) hebben getraind met AWS Trainium zonder enige modelkwaliteit te verliezen.

Overzicht oplossingen

GPT NeoX- en Pythia-modellen

GPT NeoX en Pythia zijn de open-source causale taalmodellen van Eleuther-AI met ongeveer 20 miljard parameters in NeoX en 6.9 miljard in Pythia. Beide zijn decodermodellen die een soortgelijk architectonisch ontwerp volgen als Chat GPT3. Ze hebben echter ook verschillende toevoegingen, die ook breed worden overgenomen in de recente modellen zoals Llama. In het bijzonder hebben ze een roterende positionele inbedding (ROPE) met gedeeltelijke rotatie over de hoofdafmetingen. De originele modellen (NeoX en Pythia 6.9B) zijn getraind op open beschikbaar Stapel dataset met deduplicatie en gebruik van Megatron en Deepspeed backend.

We demonstreren de pre-training en verfijning van deze modellen op AWS Trainium-gebaseerde Trn1-instanties met behulp van Neuron NeMo bibliotheek. Om het proof-of-concept en de snelle reproductie vast te stellen, gebruiken we een kleinere Wikipedia-dataset-subset die is getokeniseerd met behulp van GPT2 Byte-pair encoding (BPE) tokenizer.

walkthrough

Download de vooraf getokeniseerde Wikipedia-dataset zoals weergegeven:

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

Zowel NeoX 20B als Pythia 6.9B gebruiken ROPE met gedeeltelijke rotatie, waarbij bijvoorbeeld 25% van de afmetingen van het hoofd wordt geroteerd en de rest ongeroteerd blijft. Om de gedeeltelijke rotatie op de AWS Trainium-versneller efficiënt te implementeren, in plaats van de roterende en niet-roterende dimensies aaneen te schakelen, voegen we nulfrequenties toe voor niet-roterende dimensies en roteren we vervolgens de volledige set kopafmetingen. Deze eenvoudige truc heeft ons geholpen de doorvoer (reeksen verwerkt per seconde) op AWS Trainium te verbeteren.

Training stappen

Om de training uit te voeren, gebruiken we SLURM beheerde multi-node Amazon Elastic Compute Cloud (Amazon EC2) Trn1-cluster, waarbij elk knooppunt een trn1.32xl-instantie bevat. Elk trn1.32xl heeft 16 versnellers met twee werkers per versneller. Na het downloaden van de nieuwste Neuron NeMo pakket, gebruik het meegeleverde neox en Pythia scripts vooraf trainen en afstemmen met geoptimaliseerde hyperparameters en het volgende uitvoeren voor een training met vier knooppunten.

1. Compileren: compileer het model vooraf met drie treiniteraties om de grafieken te genereren en op te slaan:

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

2. Uitvoeren: Voer de training uit door de in de cache opgeslagen grafieken uit de eerste stappen te laden

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

3. Resultaten volgen

   tensorboard --logdir=nemo_experiments/megatron_neox

Dezelfde stappen moeten worden gevolgd voor het uitvoeren van het Pythia 6.9B-model met vervanging neox_20B_slurm.sh by pythia_6.9B_slurm.sh.

Experimenten voorafgaand aan de training en verfijning

We demonstreren de voortraining van GPT-NeoX- en Pythia-modellen op AWS Trainium met behulp van Neuron NeMo bibliotheek voor 10k-iteraties, en laat ook de fijnafstemming van deze modellen zien voor stappen van 1k. Voor pre-training gebruiken we de GPT2 BPE-tokenizer in de NeMo en volgen deze config zoals gebruikt in het originele model. Voor het afstemmen van AWS Trainium zijn enkele wijzigingen in de parameters nodig (zoals factor voor de verdeling van de woordenschatgrootte), die zijn opgenomen in de verfijningsscripts om rekening te houden met de verschillen tussen Megatron en NeMo en tussen GPU en AWS Trainium. De gedistribueerde trainingsdoorvoer over meerdere knooppunten met een variërend aantal knooppunten wordt weergegeven in Tabel 1.

Tafel 1. Vergelijking van de gemiddelde doorvoer van GPT NeoX- en Pythia-modellen voor het trainen van maximaal 500 stappen met een veranderend aantal knooppunten. De prijs van trn1.32xl is gebaseerd op het 3 jaar gereserveerde effectieve uurtarief.

Vervolgens evalueren we ook het verliestraject van de modeltraining op AWS Trainium en vergelijken we dit met de overeenkomstige run op een P4d (Nvidia A100 GPU-cores) cluster. Naast het trainingsverlies vergelijken we ook nuttige indicatoren, zoals de gradiëntnorm, wat de 2-norm is van de modelgradiënten die bij elke trainingsiteratie worden berekend om de voortgang van de training te volgen. De trainingsresultaten worden getoond in Figuur 1, 2 en de verfijning van NeoX 20B in Figuur 3.

Figuur 1. Het trainingsverlies gemiddeld voor alle werknemers (links) en de gradiëntnorm (rechts) bij het trainen van elke stap. NeoX 20B wordt getraind op 4 knooppunten met een kleine wiki-dataset op GPU en Trainium met dezelfde trainingshyperparameters (globale batchgrootte = 256). GPU gebruikt BF16 en standaard gemengde precisie, terwijl AWS Trainium volledige BF16 gebruikt met stochastische afronding. De verlies- en gradiëntnormtrajecten komen overeen voor GPU en AWS Trainium.

Figuur 2. Het trainingsverlies gemiddeld voor alle werknemers (links) en de gradiëntnorm (rechts) bij het trainen van elke stap. Vergelijkbaar met GPT NeoX in Figuur 1, wordt Pythia 6.9B getraind op 4 knooppunten met een kleine wiki-dataset op GPU en Trainium met dezelfde trainingshyperparameters (globale batchgrootte = 256). De verlies- en gradiëntnormtrajecten komen overeen voor GPU en Trainium.

Figuur 3. Verfijning van het GPT NeoX 20B-model op GPU en AWS Trainium met het gemiddelde trainingsverlies voor alle werknemers (links) en de gradiëntnorm (rechts). Een kleine wiki-dataset wordt gebruikt voor het verfijnen van demonstraties. De verlies- en gradiëntnormtrajecten komen overeen voor GPU en AWS Trainium.

In dit bericht lieten we een kostenefficiënte training van LLM's zien op AWS deep learning-hardware. We hebben GPT NeoX 20B- en Pythia 6.9B-modellen getraind op AWS Trn1 met de Neuron NeMo-bibliotheek. De kostengenormaliseerde doorvoer voor 20 miljard modellen met AWS Trainium bedraagt ​​ongeveer 3.2 miljoen tokens/uitgegeven $. Samen met de kostenefficiënte training op AWS Trainium verkrijgen we een vergelijkbare modelnauwkeurigheid, wat blijkt uit het verlies van trainingsstappen en het verloop van de gradiëntnorm. We hebben ook de beschikbare controlepunten voor het NeoX 20B-model op AWS Trainium verfijnd. Voor aanvullende informatie over de gedistribueerde training met NeMo Megatron op AWS Trainium, zie AWS-neuronreferentie voor NeMo Megatron. Een goede bron om te beginnen met het verfijnen van het Llama-model kunt u hier vinden: Llama2 fijnafstemming. Om aan de slag te gaan met managed AWS Trainium op Amazon Sage Maker, Zie Train uw ML-modellen met AWS Trainium en Amazon SageMaker.

Over de auteurs

Gaurav Gupta is momenteel een toegepast wetenschapper bij Amazon Web Services (AWS) AI-labs. Dr. Gupta promoveerde aan USC Viterbi. Zijn onderzoeksinteresses bestrijken het domein van sequentiële datamodellering, het leren van partiële differentiaalvergelijkingen, informatietheorie voor machinaal leren, fractionele dynamische modellen en complexe netwerken. Momenteel werkt hij aan toegepaste en wiskundige problemen met betrekking tot het trainingsgedrag van LLM's, visiemodellen met PDE's, informatietheoretische multimodaliteitsmodellen. Dr. Gupta heeft publicaties in toptijdschriften/conferenties zoals Neurips, ICLR, ICML, Nature, IEEE Control Society, ACM cyber-physical society.

Ben Snijder is een toegepast wetenschapper bij AWS Deep Learning. Zijn onderzoeksinteresses omvatten fundamentele modellen, versterkend leren en asynchrone optimalisatie. Buiten zijn werk houdt hij van fietsen en kamperen in het binnenland.

Amith (R) Mamidala is de senior machine learning-applicatie-engineering bij AWS Annapurna Labs. Dr. Mamidala voltooide zijn doctoraat aan de Ohio State University in high performance computing en communicatie. Tijdens zijn ambtstermijn bij IBM Research heeft Dr. Mamidala bijgedragen aan de BlueGene-klasse van computers, die vaak bovenaan de Top 500-ranglijst stond van de krachtigste en energiezuinigste supercomputers. Het project werd in 2009 bekroond met de Nationale Medaille voor Technologie en Innovatie. Na een korte periode als AI-ingenieur bij een financieel hedgefonds, trad Dr. Mamidala toe tot de Annapurna-laboratoria, waar hij zich richtte op training van grote taalmodellen.

Jun (Lucas) Huan is hoofdwetenschapper bij AWS AI Labs. Dr. Huan werkt aan AI en Data Science. Hij heeft meer dan 180 peer-reviewed artikelen gepubliceerd in toonaangevende conferenties en tijdschriften. In 2009 ontving hij de NSF Faculty Early Career Development Award. Voordat hij bij AWS kwam, werkte hij bij Baidu Research als een vooraanstaand wetenschapper en hoofd van het Baidu Big Data Laboratory. Hij richtte StylingAI Inc. op, een AI-start-up, en werkte in 2019-2021 als CEO en Chief Scientist. Voordat hij bij de industrie kwam, was hij Charles E. en Mary Jane Spahr Professor op de EECS-afdeling van de Universiteit van Kansas.

Shruti Koparkar is Senior Product Marketing Manager bij AWS. Ze helpt klanten bij het verkennen, evalueren en toepassen van de Amazon EC2-versnelde computerinfrastructuur voor hun behoeften op het gebied van machine learning.

• Door SEO aangedreven content en PR-distributie. Word vandaag nog versterkt.
• PlatoData.Network Verticale generatieve AI. Versterk jezelf. Toegang hier.
• PlatoAiStream. Web3-intelligentie. Kennis versterkt. Toegang hier.
• PlatoESG. carbon, CleanTech, Energie, Milieu, Zonne, Afvalbeheer. Toegang hier.
• Plato Gezondheid. Intelligentie op het gebied van biotech en klinische proeven. Toegang hier.
• Bron: https://aws.amazon.com/blogs/machine-learning/frugality-meets-accuracy-cost-efficient-training-of-gpt-neox-and-pythia-models-with-aws-trainium/