La frugalità incontra la precisione: formazione economicamente vantaggiosa dei modelli GPT NeoX e Pythia con AWS Trainium | Servizi Web di Amazon

I grandi modelli linguistici (o LLM) sono diventati argomento di conversazioni quotidiane. La loro rapida adozione è evidente dalla quantità di tempo necessaria per raggiungere 100 milioni di utenti, che è passata da “4.5 anni tramite Facebook” al minimo storico di soli “2 mesi tramite ChatGPT”. Un trasformatore generativo pre-addestrato (GPT) utilizza aggiornamenti autoregressivi causali per effettuare previsioni. È stato dimostrato che una varietà di attività come il riconoscimento vocale, la generazione di testo e la risposta alle domande hanno prestazioni stupende grazie a queste architetture modello. Diversi modelli recenti come NeoX, falco, Lama utilizzare l'architettura GPT come spina dorsale. La formazione dei LLM richiede un'enorme quantità di tempo di elaborazione, che costa milioni di dollari. In questo post riassumeremo la procedura di formazione di GPT NeoX on AWSTrainium, un acceleratore di machine learning (ML) appositamente creato e ottimizzato per la formazione sul deep learning. Illustreremo come abbiamo addestrato tali modelli in modo conveniente (3.2 milioni di token/$) con AWS Trainium senza perdere la qualità del modello.

Panoramica della soluzione

Modelli GPT NeoX e Pythia

GPT NeoX ed Pythia sono i modelli di linguaggio causale open source di Eleuther-AI con circa 20 miliardi di parametri in NeoX e 6.9 ​​miliardi in Pythia. Entrambi sono modelli di decoder che seguono un design architettonico simile a Chat GPT3. Presentano però anche diverse aggiunte, ampiamente adottate anche nei modelli recenti come Llama. In particolare, hanno l'inclusione posizionale rotazionale (ROPE) con rotazione parziale attraverso le dimensioni della testa. I modelli originali (NeoX e Pythia 6.9B) sono addestrati su strumenti liberamente disponibili Insieme di dati della pila con deduplicazione e utilizzando il backend Megatron e Deepspeed.

Dimostriamo il pre-addestramento e la messa a punto di questi modelli su istanze Trn1 basate su AWS Trainium utilizzando Neurone NeMo biblioteca. Per stabilire la prova di concetto e la riproduzione rapida, utilizzeremo un sottoinsieme di set di dati di Wikipedia più piccolo tokenizzato utilizzando il tokenizzatore GPT2 Byte-pair encoding (BPE).

Soluzione

Scarica il set di dati Wikipedia pre-tokenizzato come mostrato:

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

Sia NeoX 20B che Pythia 6.9B utilizzano ROPE con rotazione parziale, ad esempio, ruotando il 25% delle dimensioni della testa e mantenendo il resto non ruotato. Per implementare in modo efficiente la rotazione parziale sull'acceleratore AWS Trainium, invece di concatenare le dimensioni rotanti e non rotanti, aggiungiamo frequenze zero per le dimensioni non rotanti e quindi ruotiamo il set completo di dimensioni della testa. Questo semplice trucco ci ha aiutato a migliorare il throughput (sequenze elaborate al secondo) su AWS Trainium.

Passaggi di formazione

Per eseguire la formazione, utilizziamo Amazon Elastic Compute Cloud multinodo gestito da SLURM (Amazon EC2) Cluster Trn1, con ciascun nodo contenente un'istanza trn1.32xl. Ogni trn1.32xl ha 16 acceleratori con due lavoratori per acceleratore. Dopo aver scaricato l'ultimo Neurone NeMo pacchetto, utilizzare il pacchetto fornito neox ed Pythia script di pre-addestramento e messa a punto con iperparametri ottimizzati ed eseguire quanto segue per un addestramento a quattro nodi.

1. Compila: precompila il modello con tre iterazioni del treno per generare e salvare i grafici:

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

2. Esegui: esegui l'addestramento caricando i grafici memorizzati nella cache dai primi passaggi

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

3. Monitorare i risultati

   tensorboard --logdir=nemo_experiments/megatron_neox

È necessario seguire gli stessi passaggi per eseguire il modello Pythia 6.9B con sostituzione neox_20B_slurm.sh by pythia_6.9B_slurm.sh.

Esperimenti di pre-addestramento e messa a punto

Dimostriamo il pre-addestramento dei modelli GPT-NeoX e Pythia su AWS Trainium utilizzando Neurone NeMo libreria per iterazioni da 10 e mostra anche la messa a punto di questi modelli per passaggi da 1. Per il pre-addestramento, utilizziamo il tokenizer GPT2 BPE all'interno di NeMo e seguiamo lo stesso config come utilizzato nel modello originale. La messa a punto su AWS Trainium richiede la modifica di alcuni parametri (come fattore di divisione della dimensione del vocabolario), forniti negli script di ottimizzazione per tenere conto delle differenze tra Megatron e NeMo e delle modifiche tra GPU e AWS Trainium. Il throughput di addestramento distribuito su più nodi con un numero variabile di nodi è mostrato nella Tabella 1.

Tabella 1. Confronto del throughput medio dei modelli GPT NeoX e Pythia per l'addestramento fino a 500 passaggi con la modifica del numero di nodi. IL prezzo di trn1.32xl si basa sulla tariffa oraria effettiva prenotata per 3 anni.

Successivamente, valutiamo anche la traiettoria di perdita dell'addestramento del modello su AWS Trainium e la confrontiamo con l'esecuzione corrispondente su un cluster P4d (core GPU Nvidia A100). Insieme alla perdita di allenamento, confrontiamo anche indicatori utili come la norma del gradiente, che è la 2-norma dei gradienti del modello calcolata ad ogni iterazione dell'allenamento per monitorare i progressi dell'allenamento. I risultati dell'addestramento sono mostrati nella Figura 1, 2 e la messa a punto di NeoX 20B nella Figura 3.

Figura 1. Perdita di allenamento media tra tutti i lavoratori (a sinistra) e norma del gradiente (a destra) durante l'allenamento per ogni fase. NeoX 20B viene addestrato su 4 nodi con un piccolo set di dati wiki su GPU e Trainium con gli stessi iperparametri di addestramento (dimensione batch globale = 256). La GPU utilizza BF16 e la precisione mista predefinita mentre AWS Trainium utilizza BF16 completo con arrotondamento stocastico. Le traiettorie delle norme di perdita e gradiente corrispondono per GPU e AWS Trainium.

Figura 2. Perdita di allenamento media tra tutti i lavoratori (a sinistra) e norma del gradiente (a destra) durante l'allenamento per ogni fase. Similmente a GPT NeoX nella Figura 1, Pythia 6.9B viene addestrato su 4 nodi con un piccolo set di dati wiki su GPU e Trainium con gli stessi iperparametri di addestramento (dimensione batch globale=256). Le traiettorie della norma di perdita e gradiente corrispondono per GPU e Trainium.

Figura 3. Ottimizzazione del modello GPT NeoX 20B su GPU e AWS Trainium con perdita di addestramento media tra tutti i lavoratori (a sinistra) e norma del gradiente (a destra). Un piccolo set di dati wiki viene utilizzato per la dimostrazione di perfezionamento. Le traiettorie delle norme di perdita e gradiente corrispondono per GPU e AWS Trainium.

In questo post, abbiamo mostrato una formazione economicamente vantaggiosa per LLM sull'hardware di deep learning AWS. Abbiamo addestrato i modelli GPT NeoX 20B e Pythia 6.9B su AWS Trn1 con la libreria Neuron NeMo. Il throughput normalizzato dei costi per 20 miliardi di modelli con AWS Trainium è di circa 3.2 milioni di token/$ spesi. Insieme alla formazione conveniente su AWS Trainium, otteniamo una precisione del modello simile, che è evidente dalla perdita di passi di addestramento e dalla traiettoria della norma del gradiente. Abbiamo anche perfezionato i checkpoint disponibili per il modello NeoX 20B su AWS Trainium. Per ulteriori informazioni sulla formazione distribuita con NeMo Megatron su AWS Trainium, vedere Riferimento AWS Neuron per NeMo Megatron. Una buona risorsa per iniziare la messa a punto del modello Llama può essere trovata qui, Perfezionamento di Llama2. Per iniziare con AWS Trainium gestito su Amazon Sage Maker, Vedere Addestra i tuoi modelli ML con AWS Trainium e Amazon SageMaker.

Informazioni sugli autori

Gaurav Gupta è attualmente uno scienziato applicato presso i laboratori AI di Amazon Web Services (AWS). Il dottor Gupta ha completato il suo dottorato di ricerca presso l'USC Viterbi. I suoi interessi di ricerca abbracciano il dominio della modellazione di dati sequenziali, l'apprendimento delle equazioni differenziali parziali, la teoria dell'informazione per l'apprendimento automatico, i modelli dinamici frazionari e le reti complesse. Attualmente sta lavorando su problemi applicati e matematici sul comportamento formativo degli LLM, modelli di visione con PDE, modelli multimodali di teoria dell'informazione. Il Dr. Gupta ha pubblicazioni nelle migliori riviste/conferenze come Neurips, ICLR, ICML, Nature, IEEE Control Society, ACM cyber-physical society.

Ben Snyder è uno scienziato applicato con AWS Deep Learning. I suoi interessi di ricerca includono modelli fondazionali, apprendimento per rinforzo e ottimizzazione asincrona. Al di fuori del lavoro, gli piace andare in bicicletta e fare campeggio in backcountry.

Amith (a destra) Mamidala è l'ingegnere senior delle applicazioni di machine learning presso AWS Annapurna Labs. Il Dr. Mamidala ha completato il suo dottorato di ricerca presso la Ohio State University in comunicazione e calcolo ad alte prestazioni. Durante il suo mandato alla ricerca IBM, il dottor Mamidala ha contribuito alla classe di computer BlueGene che spesso guidava la classifica Top500 dei supercomputer più potenti ed efficienti dal punto di vista energetico. Il progetto è stato insignito nel 2009 della Medaglia Nazionale per la Tecnologia e l'Innovazione. Dopo un breve periodo come ingegnere di intelligenza artificiale presso un hedge fund finanziario, il dottor Mamidala si è unito ai laboratori dell'Annapurna concentrandosi sulla formazione del modello Large Language.

Jun (Luca) Huan è uno scienziato principale presso AWS AI Labs. Il dottor Huan lavora su intelligenza artificiale e scienza dei dati. Ha pubblicato più di 180 articoli sottoposti a peer review in importanti conferenze e riviste. Ha ricevuto il premio NSF Faculty Early Career Development Award nel 2009. Prima di unirsi ad AWS, ha lavorato presso Baidu Research come illustre scienziato e capo del Baidu Big Data Laboratory. Ha fondato StylingAI Inc., una start-up basata sull'intelligenza artificiale, e ha lavorato come CEO e Chief Scientist nel 2019-2021. Prima di entrare nell'industria, è stato professore Charles E. e Mary Jane Spahr presso il dipartimento EECS dell'Università del Kansas.

Shruti Koparkar è Senior Product Marketing Manager presso AWS. Aiuta i clienti a esplorare, valutare e adottare l'infrastruttura di elaborazione accelerata Amazon EC2 per le loro esigenze di machine learning.

• Distribuzione di contenuti basati su SEO e PR. Ricevi amplificazione oggi.
• PlatoData.Network Generativo verticale Ai. Potenzia te stesso. Accedi qui.
• PlatoAiStream. Intelligenza Web3. Conoscenza amplificata. Accedi qui.
• PlatoneESG. Carbonio, Tecnologia pulita, Energia, Ambiente, Solare, Gestione dei rifiuti. Accedi qui.
• Platone Salute. Intelligence sulle biotecnologie e sulle sperimentazioni cliniche. Accedi qui.
• Fonte: https://aws.amazon.com/blogs/machine-learning/frugality-meets-accuracy-cost-efficient-training-of-gpt-neox-and-pythia-models-with-aws-trainium/