Ottimizzazione avanzata di 20B LLM con RLHF su una GPU da 24GB per consumatori

'20B LLM RLHF GPU optimization for consumers'

Siamo entusiasti di annunciare ufficialmente l’integrazione di trl con peft per rendere l’addestramento dei Grandi Modelli di Linguaggio (LLM) con il Reinforcement Learning più accessibile a tutti! In questo post, spieghiamo perché questa è un’alternativa competitiva rispetto agli approcci di addestramento esistenti.

Nota che peft è uno strumento generale che può essere applicato a molti casi d’uso di Machine Learning, ma è particolarmente interessante per RLHF in quanto questo metodo richiede molta memoria!

Se vuoi approfondire direttamente il codice, dai un’occhiata agli script di esempio direttamente nella pagina di documentazione di TRL.

Introduzione

LLM e RLHF

I LLM combinati con RLHF (Reinforcement Learning con Feedback Umano) sembrano essere l’approccio preferito per la creazione di sistemi AI molto potenti come ChatGPT.

• Jupyter X Hugging Face
• Addestra il tuo ControlNet con diffusori
• Apprendimento federato utilizzando Hugging Face e Flower

L’addestramento di un modello di linguaggio con RLHF coinvolge tipicamente i seguenti tre passaggi:

1- Fine-tuning di un LLM pre-addestrato su un dominio specifico o un corpus di istruzioni e dimostrazioni umane

2- Raccolta di un dataset annotato dall’umanità e addestramento di un modello di reward

3- Successivo fine-tuning del LLM del passaggio 1 con il modello di reward e questo dataset utilizzando RL (ad esempio PPO)

La scelta del LLM di base è molto cruciale qui. Al momento della scrittura, i LLM finetuned per le istruzioni sono i “migliori” modelli open-source che possono essere utilizzati “pronti all’uso” per molti compiti. Modelli notevoli sono: BLOOMZ, Flan-T5, Flan-UL2 e OPT-IML. Il lato negativo di questi modelli è la loro dimensione. Per ottenere un modello decente, è necessario almeno utilizzare modelli su scala di 10B+, che richiederebbero fino a 40GB di memoria GPU in precisione completa, solo per adattare il modello su un singolo dispositivo GPU senza fare alcun addestramento!

Cos’è TRL?

La libreria trl mira a semplificare e rendere più flessibile il passaggio RL in modo che chiunque possa addestrare il proprio modello di linguaggio utilizzando RL sul proprio dataset personalizzato e la propria configurazione di addestramento. Tra molte altre applicazioni, è possibile utilizzare questo algoritmo per addestrare un modello a generare recensioni positive di film, creare generazione controllata o rendere il modello meno tossico.

Utilizzando trl è possibile eseguire uno degli algoritmi Deep RL più popolari, PPO, in modo distribuito o su un singolo dispositivo! Sfruttiamo accelerate nell’ecosistema Hugging Face per rendere tutto ciò possibile, in modo che ogni utente possa scalare gli esperimenti a una scala interessante.

Il fine-tuning di un modello di linguaggio con RL segue approssimativamente il protocollo dettagliato di seguito. Questo richiede di avere 2 copie del modello originale; per evitare che il modello attivo si discosti troppo dal suo comportamento / distribuzione originale, è necessario calcolare i logits del modello di riferimento ad ogni step di ottimizzazione. Ciò aggiunge un vincolo rigido al processo di ottimizzazione, poiché è sempre necessario avere almeno due copie del modello per dispositivo GPU. Se la dimensione del modello aumenta, diventa sempre più complicato adattare l’intero setup su una singola GPU.

In trl è anche possibile utilizzare layer condivisi tra il modello di riferimento e il modello attivo per evitare copie complete. Un esempio concreto di questa funzionalità viene mostrato nell’esempio di detoxification.

Addestramento su larga scala

L’addestramento su larga scala può essere una sfida. La prima sfida è adattare il modello e i suoi stati di ottimizzazione sui dispositivi GPU disponibili. La quantità di memoria GPU che un singolo parametro occupa dipende dalla sua “precisione” (o più specificamente dtype). Le precisioni più comuni sono float32 (32-bit), float16 e bfloat16 (16-bit). Più recentemente, sono supportate “precisioni esotiche” per addestramento e inferenza (con determinate condizioni e vincoli), come int8 (8-bit). In breve, per caricare un modello su un dispositivo GPU, ogni miliardo di parametri costa 4GB in precisione float32, 2GB in float16 e 1GB in int8. Se desideri saperne di più su questo argomento, dai un’occhiata a questo blogpost che approfondisce: https://huggingface.co/blog/hf-bitsandbytes-integration.

Se utilizzi un ottimizzatore AdamW, ogni parametro richiede 8 byte (ad esempio, se il tuo modello ha 1 miliardo di parametri, l’ottimizzatore AdamW completo del modello richiederà 8GB di memoria GPU – fonte).

Sono state adottate molte tecniche per affrontare queste sfide su larga scala. I paradigmi più familiari sono il Parallelismo di Pipeline, il Parallelismo di Tensori e il Parallelismo di Dati.

Con il parallelismo dei dati, lo stesso modello è ospitato in parallelo su diverse macchine e ogni istanza viene alimentata con un batch di dati diverso. Questa è la strategia di parallelismo più semplice, che essenzialmente replica il caso di una singola GPU ed è già supportata da trl. Con il parallelismo di Pipeline e Tensori, il modello stesso viene distribuito tra le macchine: nel parallelismo di Pipeline, il modello viene diviso per strati, mentre il parallelismo di Tensori divide le operazioni sui tensori tra le GPU (ad esempio moltiplicazioni di matrici). Con queste strategie di parallelismo del modello, è necessario suddividere i pesi del modello tra molti dispositivi, il che richiede la definizione di un protocollo di comunicazione delle attivazioni e dei gradienti tra i processi. Questo non è banale da implementare e potrebbe richiedere l’adozione di alcuni framework come Megatron-DeepSpeed o Nemo. È anche importante evidenziare altri strumenti che sono essenziali per la scalabilità dell’addestramento di LLM, come il checkpointing dell’attivazione adattiva e i kernel fusi. Ulteriori letture sui paradigmi di parallelismo possono essere trovate qui.

Pertanto, ci siamo posti la seguente domanda: fino a che punto possiamo arrivare con il solo parallelismo dei dati? Possiamo utilizzare gli strumenti esistenti per adattare processi di addestramento di dimensioni super grandi (inclusi il modello attivo, il modello di riferimento e gli stati dell’ottimizzatore) in un singolo dispositivo? La risposta sembra essere sì. Gli ingredienti principali sono: adattatori e moltiplicazione di matrici a 8 bit! Copriremo questi argomenti nelle sezioni seguenti:

Moltiplicazione di matrici a 8 bit

La moltiplicazione efficiente di matrici a 8 bit è un metodo che è stato introdotto per la prima volta nel paper LLM.int8() e mira a risolvere il problema della degradazione delle prestazioni durante la quantizzazione di modelli su larga scala. Il metodo proposto suddivide le moltiplicazioni di matrici che vengono applicate nei livelli lineari in due fasi: la parte degli stati nascosti outlier che verrà eseguita in float16 e la parte “non outlier” che viene eseguita in int8.

In poche parole, è possibile ridurre le dimensioni di un modello a precisione completa di 4 (quindi di 2 per modelli a mezza precisione) utilizzando la moltiplicazione di matrici a 8 bit.

Adattamento a rango basso e PEFT

Nel 2021, un articolo chiamato LoRA: Low-Rank Adaption of Large Language Models ha dimostrato che il fine-tuning di grandi modelli di linguaggio può essere eseguito congelando i pesi preaddestrati e creando versioni a rango basso delle matrici di attenzione dei livelli di query e valore. Queste matrici a rango basso hanno molti meno parametri rispetto al modello originale, consentendo il fine-tuning con una quantità molto inferiore di memoria GPU. Gli autori dimostrano che il fine-tuning degli adattatori a rango basso ha ottenuto risultati comparabili al fine-tuning del modello preaddestrato completo.

Questa tecnica consente il fine-tuning di LLM utilizzando una frazione dei requisiti di memoria. Tuttavia, ci sono anche alcuni svantaggi. Il passaggio in avanti e all’indietro è approssimativamente due volte più lento, a causa delle moltiplicazioni di matrici aggiuntive nei livelli degli adattatori.

Cos’è PEFT?

Parameter-Efficient Fine-Tuning (PEFT) è una libreria di Hugging Face, creata per supportare la creazione e il fine-tuning di livelli di adattatori su LLM. peft è integrato in modo trasparente con 🤗 Accelerate per modelli su larga scala che sfruttano DeepSpeed e Big Model Inference.

La libreria supporta molti modelli all’avanguardia e ha un ampio set di esempi, tra cui:

• Modellazione del linguaggio causale
• Generazione condizionale
• Classificazione delle immagini
• Addestramento a 8 bit int8
• Adattamento a rango basso di modelli Dreambooth
• Segmentazione semantica
• Classificazione di sequenze
• Classificazione di token

La libreria è ancora in fase di sviluppo estensivo e attivo, con molte nuove funzionalità che verranno annunciate nei prossimi mesi.

Fine-tuning di modelli a 20B di parametri con adattatori a rango basso

Ora che le premesse sono state spiegate, passiamo attraverso l’intera pipeline passo dopo passo e spieghiamo con le figure come è possibile eseguire il fine-tuning di un LLM a 20B di parametri con RL utilizzando gli strumenti sopra menzionati su una singola GPU da 24GB!

Passo 1: Carica il tuo modello attivo con precisione a 8 bit

Una riduzione di memoria “tutto incluso” di un LLM utilizzando transformers consiste nel caricare il modello con precisione a 8 bit utilizzando il metodo descritto in LLM.int8. Ciò può essere eseguito semplicemente aggiungendo il flag load_in_8bit=True quando si chiama il metodo from_pretrained (puoi leggere di più a riguardo qui).

Come indicato nella sezione precedente, un “hack” per calcolare la quantità di memoria GPU necessaria per caricare il tuo modello è pensare in termini di “miliardi di parametri”. Poiché un byte corrisponde a 8 bit, hai bisogno di 4 GB per miliardo di parametri per un modello a piena precisione (32 bit = 4 byte), 2 GB per miliardo di parametri per un modello a mezza precisione e 1 GB per miliardo di parametri per un modello int8.

Quindi, in primo luogo, carichiamo semplicemente il modello attivo a 8 bit. Vediamo cosa dobbiamo fare per il secondo passaggio!

Passo 2: Aggiungi adattatori addestrabili aggiuntivi utilizzando peft

Il secondo passaggio consiste nel caricare gli adattatori all’interno del modello e renderli addestrabili. Questo consente una drastica riduzione del numero di pesi addestrabili necessari per il modello attivo. Questo passaggio sfrutta la libreria peft e può essere eseguito con poche righe di codice. Nota che una volta che gli adattatori sono addestrati, puoi facilmente caricarli nell’Hub per usarli in seguito.

Passo 3: Utilizza lo stesso modello per ottenere i logit di riferimento e attivi

Poiché gli adattatori possono essere disattivati, possiamo utilizzare lo stesso modello per ottenere i logit di riferimento e attivi per PPO, senza dover creare due copie dello stesso modello! Questo sfrutta una funzionalità nella libreria peft, che è il gestore di contesto disable_adapters.

Panoramica degli script di addestramento:

Ora descriveremo come abbiamo addestrato un modello gpt-neox con 20 miliardi di parametri utilizzando transformers, peft e trl. L’obiettivo finale di questo esempio era ottimizzare un LLM per generare recensioni positive di film in un ambiente con vincoli di memoria. Passaggi simili potrebbero essere applicati ad altre attività, come i modelli di dialogo.

In generale, c’erano tre passaggi chiave e script di addestramento:

1. Script – Ottimizzazione di un adattatore di basso rango su un modello a 8 bit congelato per la generazione di testo sul dataset imdb.
2. Script – Unione dei livelli degli adattatori nei pesi del modello di base e archiviazione di questi nell’hub.
3. Script – Ottimizzazione del sentimento di un adattatore di basso rango per creare recensioni positive.

Abbiamo testato questi passaggi su una GPU NVIDIA 4090 da 24 GB. Sebbene sia possibile eseguire l’intero processo di addestramento su una GPU da 24 GB, gli addestramenti completi sono stati eseguiti su un singolo A100 nel cluster di ricerca di 🤗.

Il primo passaggio nel processo di addestramento consisteva nell’ottimizzazione del modello preaddestrato. Tipicamente, questo richiederebbe diverse GPU A100 da 80 GB di fascia alta, quindi abbiamo scelto di addestrare un adattatore di basso rango. Abbiamo trattato questo come un ambiente di modellazione del linguaggio causale e ci siamo addestrati per un’epoca di esempi dal dataset imdb, che contiene recensioni di film e etichette che indicano se sono di sentimento positivo o negativo.

Per prendere il modello adattato e eseguire ulteriori ottimizzazioni con RL, abbiamo prima dovuto combinare i pesi adattati, ciò è stato ottenuto caricando il modello preaddestrato e l’adattatore in virgola mobile a 16 bit e sintetizzando le matrici di pesi (con l’applicazione della scala appropriata).

Infine, abbiamo potuto ottimizzare un altro adattatore di basso rango, sopra il modello imdb-finetuned congelato. Abbiamo utilizzato un classificatore di sentimenti imdb per fornire le ricompense per l’algoritmo RL.

Il rapporto completo di Weights and Biases per questo esperimento è disponibile qui, se desideri controllare ulteriori grafici e generazioni di testo.

Conclusioni

Abbiamo implementato una nuova funzionalità in trl che consente agli utenti di ottimizzare grandi modelli di linguaggio utilizzando RLHF a un costo ragionevole sfruttando le librerie peft e bitsandbytes. Abbiamo dimostrato che è possibile ottimizzare gpt-neo-x (40 GB in bfloat16!) su una GPU da consumatore da 24 GB, e ci aspettiamo che questa integrazione venga ampiamente utilizzata dalla comunità per ottimizzare modelli più grandi utilizzando RLHF e condividere grandi artefatti.

Abbiamo individuato alcune direzioni interessanti per i prossimi passi per spingere i limiti di questa integrazione.

• Come sarà adattabile questa soluzione nell’ambito del multi-GPU? Esploreremo principalmente come questa integrazione si adatterà rispetto al numero di GPU, se sarà possibile applicare il Parallelismo dei Dati immediatamente o se richiederà l’adozione di nuove funzionalità in una delle librerie coinvolte.
• Quali strumenti possiamo sfruttare per aumentare la velocità di addestramento? Abbiamo osservato che il principale svantaggio di questa integrazione è la velocità complessiva di addestramento. In futuro saremmo interessati a esplorare le possibili direzioni per rendere l’addestramento molto più veloce.

Riferimenti

• paradigmi di parallelismo: https://huggingface.co/docs/transformers/v4.17.0/en/parallelism
• integrazione a 8-bit in transformers: https://huggingface.co/blog/hf-bitsandbytes-integration
• articolo LLM.int8: https://arxiv.org/abs/2208.07339
• Spiegazione del checkpointing dei gradienti: https://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel-extended-features-pytorch-activation-checkpointing.html