Perfezionamento rapido ed economico LLaMA 2 con AWS Trainium

Perfezionamento veloce ed economico con LLaMA 2 grazie ad AWS Trainium

I modelli linguistici di grandi dimensioni (LLM) hanno catturato l’immaginazione e l’attenzione di sviluppatori, scienziati, tecnologi, imprenditori e dirigenti in diversi settori. Questi modelli possono essere utilizzati per rispondere a domande, fare sintesi, traduzioni e altro ancora in applicazioni come agenti conversazionali per il supporto clienti, creazione di contenuti per il marketing e assistenti di coding.

Recentemente, Meta ha rilasciato Llama 2 sia per ricercatori che per entità commerciali, aggiungendosi alla lista di altri LLM, tra cui MosaicML MPT e Falcon. In questo post, vedremo come ottimizzare Llama 2 su AWS Trainium, un acceleratore appositamente progettato per l’addestramento di LLM, al fine di ridurre i tempi e i costi di addestramento. Esamineremo gli script di ottimizzazione forniti dal SDK AWS Neuron (utilizzando NeMo Megatron-LM), le diverse configurazioni utilizzate e i risultati del throughput che abbiamo ottenuto.

Informazioni sul modello Llama 2

Simile al modello precedente Llama 1 e ad altri modelli come GPT, Llama 2 utilizza l’architettura solo-decoder del Transformer. Viene proposto in tre dimensioni: 7 miliardi, 13 miliardi e 70 miliardi di parametri. Rispetto a Llama 1, Llama 2 raddoppia la lunghezza del contesto da 2.000 a 4.000 e utilizza l’attenzione raggruppata delle query (solo per 70B). I modelli preaddestrati di Llama 2 sono stati allenati su 2 trilioni di token e i modelli ottimizzati sono stati addestrati su oltre 1 milione di annotazioni umane.

Addestramento distribuito di Llama 2

Per adattare Llama 2 con una lunghezza sequenza di 2.000 e 4.000, abbiamo implementato lo script utilizzando NeMo Megatron per Trainium che supporta il parallelismo dei dati (DP), il parallelismo dei tensori (TP) e il parallelismo delle pipeline (PP). Per essere specifici, con l’implementazione di nuove funzionalità come unire l’embedding delle parole, l’embedding rotativo, RMSNorm e l’attivazione Swiglu, utilizzeremo lo script generico di GPT Neuron Megatron-LM per supportare lo script di addestramento di Llama 2.

• Migliora la qualità delle previsioni nei modelli di classificazione personalizzati con Amazon Comprehend
• Elaborazione intelligente dei documenti con i servizi di intelligenza artificiale AWS e Amazon Bedrock
• Perfezionare LLM su Slack Messages

La nostra procedura di addestramento ad alto livello è la seguente: per il nostro ambiente di addestramento, utilizziamo un cluster multi-istanza gestito dal sistema SLURM per l’addestramento distribuito e la pianificazione nell’ambito del framework NeMo.

Prima di tutto, scarichiamo il modello Llama 2 e i set di dati di addestramento e li preelaboriamo utilizzando il tokenizer Llama 2. Ad esempio, per utilizzare il dataset RedPajama, utilizziamo il seguente comando:

wget https://data.together.xyz/redpajama-data-1T/v1.0.0/book/book.jsonlpython nemo/scripts/nlp_language_modeling/preprocess_data_for_megatron.py

Per una guida dettagliata sul download dei modelli e sugli argomenti dello script di preelaborazione, fare riferimento a Download LlamaV2 dataset and tokenizer.

Successivamente, compiliamo il modello:

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

Dopo che il modello è stato compilato, avviamo il job di addestramento con lo script seguente, già ottimizzato con la migliore configurazione e i migliori iperparametri per Llama 2 (inclusi nel codice di esempio):

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

Infine, monitoriamo TensorBoard per tenere traccia dei progressi dell’addestramento:

tensorboard --logdir ./

Per il codice di esempio completo e gli script che abbiamo menzionato, consulta il tutorial sulla Llama 7B e il codice NeMo nell’SDK Neuron per seguire passaggi più dettagliati.

Esperimenti di fine-tuning

Abbiamo effettuato il fine-tuning del modello 7B sui dataset OSCAR (Open Super-large Crawled ALMAnaCH coRpus) e QNLI (Question-answering NLI) in un ambiente Neuron 2.12 (PyTorch). Per ogni lunghezza sequenza di 2.000 e 4.000, abbiamo ottimizzato alcune configurazioni, come batchsize e gradient_accumulation, per l’efficienza di addestramento. Come strategia di fine-tuning, abbiamo adottato il fine-tuning completo di tutti i parametri (circa 500 passi), che può essere esteso a un pre-training con passi e dataset più lunghi (ad esempio, 1T RedPajama). La parallelizzazione della sequenza può anche essere abilitata per consentire a NeMo Megatron di eseguire il fine-tuning con una lunghezza di sequenza superiore a 4.000. La tabella seguente mostra la configurazione e i risultati del throughput dell’esperimento di fine-tuning della Llama 7B. Il throughput scala quasi linearmente all’aumentare del numero di istanze fino a 4.

Libreria Distribuita	Dataset	Lunghezza Sequenza	Numero di Istanze	Parallelismo Tensoriale	Parallelismo dei Dati	Parallelismo Pipeline	Batch Globale	Throughput (seq/s)
Neuron NeMo Megatron	OSCAR	4096	1	8	4	1	256	3.7
.	.	4096	2	8	4	1	256	7.4
.	.	4096	4	8	4	1	256	14.6
.	QNLI	4096	4	8	4	1	256	14.1

L’ultimo passo è verificare l’accuratezza con il modello di base. Abbiamo implementato uno script di riferimento per gli esperimenti con GPU e abbiamo confermato che le curve di allenamento per GPU e Trainium corrispondono, come mostrato nella figura seguente. La figura illustra le curve di perdita in relazione al numero di passaggi di allenamento sul dataset QNLI. Per GPU è stato adottato il metodo a precisione mista (blu), mentre per Trainium è stato utilizzato il metodo bf16 con arrotondamento stocastico predefinito (arancione).

Conclusioni

In questo post abbiamo dimostrato che Trainium offre elevate prestazioni e un’ottimizzazione dei costi per il raffinamento di Llama 2. Per ulteriori risorse sull’utilizzo di Trainium per il pre-training distribuito e il raffinamento dei modelli di intelligenza artificiale generativi con NeMo Megatron, consulta AWS Neuron Reference per NeMo Megatron.