La Tecnologia dietro BLOOM Training

La tecnologia di BLOOM Training.

Negli ultimi anni, è diventata la norma addestrare modelli di linguaggio sempre più grandi. Mentre la questione del fatto che tali modelli non vengano rilasciati per ulteriori studi viene spesso discussa, le conoscenze nascoste su come addestrare tali modelli raramente ricevono attenzione. Questo articolo mira a cambiare questa situazione, gettando un po’ di luce sulla tecnologia e l’ingegneria dietro l’addestramento di tali modelli, sia dal punto di vista hardware che software, prendendo come esempio il modello di linguaggio BLOOM con 176 miliardi di parametri.

Ma prima di tutto vorremmo ringraziare le aziende, le persone chiave e i gruppi che hanno reso possibile l’incredibile impresa di addestrare un modello con 176 miliardi di parametri da parte di un piccolo gruppo di persone dedicate.

Successivamente verranno discusse la configurazione hardware e i principali componenti tecnologici.

Ecco un breve riassunto del progetto:

Persone

Il progetto è stato ideato da Thomas Wolf (co-fondatore e CSO – Hugging Face), che ha osato competere con le grandi corporazioni non solo per addestrare uno dei modelli multilingue più grandi, ma anche per rendere il risultato finale accessibile a tutte le persone, trasformando così quello che era solo un sogno per la maggior parte delle persone in una realtà.

Questo articolo si concentra specificamente sul l’aspetto ingegneristico dell’addestramento del modello. La parte più importante della tecnologia dietro BLOOM sono state le persone e le aziende che hanno condiviso la loro esperienza e ci hanno aiutato con la scrittura del codice e l’addestramento.

Vi sono 6 gruppi principali di persone da ringraziare:

Il team BigScience di HuggingFace, che ha dedicato più di mezzo dozzina di dipendenti a tempo pieno per comprendere e gestire l’addestramento dall’inizio alla fine e ha fornito e pagato tutta l’infrastruttura oltre al calcolo di Jean Zay.
Il team Microsoft DeepSpeed, che ha sviluppato DeepSpeed e successivamente l’ha integrato con Megatron-LM, e i cui sviluppatori hanno dedicato molte settimane alle esigenze del progetto e hanno fornito molti preziosi consigli pratici ed esperienziali prima e durante l’addestramento.
Il team NVIDIA Megatron-LM, che ha sviluppato Megatron-LM e che è stato estremamente disponibile nel rispondere alle nostre numerose domande e nel fornire consigli di prima classe basati sull’esperienza.
Il team IDRIS / GENCI che gestisce il supercomputer Jean Zay, che ha donato al progetto una quantità incredibile di potenza di calcolo e un ottimo supporto di amministrazione di sistema.
Il team PyTorch che ha creato un framework estremamente potente, su cui si basava il resto del software, e che ci ha fornito un grande supporto durante la preparazione per l’addestramento, risolvendo molti bug e migliorando l’usabilità dei componenti PyTorch di cui ci siamo basati durante l’addestramento.
I volontari nel gruppo di lavoro di ingegneria di BigScience

Sarebbe molto difficile elencare tutte le persone straordinarie che hanno contribuito all’aspetto ingegneristico del progetto, quindi citerò solo alcune persone chiave al di fuori di Hugging Face che sono state la base ingegneristica di questo progetto negli ultimi 14 mesi:

Olatunji Ruwase, Deepak Narayanan, Jeff Rasley, Jared Casper, Samyam Rajbhandari e Rémi Lacroix

Siamo anche grati a tutte le aziende che hanno permesso ai loro dipendenti di contribuire a questo progetto.

Panoramica

L’architettura di BLOOM è molto simile a quella di GPT3 con alcune migliorie aggiuntive che saranno discusse in seguito in questo articolo.

Il modello è stato addestrato su Jean Zay, il supercomputer finanziato dal governo francese che è gestito da GENCI e installato presso IDRIS, il centro di calcolo nazionale del Centre National de la Recherche Scientifique (CNRS) francese. Il calcolo è stato generosamente donato al progetto da GENCI (grant 2021-A0101012475).

Durante l’addestramento sono stati utilizzati i seguenti componenti hardware:

GPU: 384 GPU NVIDIA A100 da 80 GB (48 nodi) + 32 GPU di riserva
8 GPU per nodo utilizzando NVLink 4 per le connessioni tra GPU, 4 collegamenti OmniPath
CPU: processore AMD EPYC 7543 a 32 core
Memoria CPU: 512 GB per nodo
Memoria GPU: 640 GB per nodo
Connessione tra i nodi: architettura Omni-Path (OPA) con albero grasso non bloccante
Rete di comunicazione NCCL: una subnet completamente dedicata
Rete IO disco: GPFS condiviso con altri nodi e utenti

Checkpoint:

checkpoint principali
ogni checkpoint con stati di ottimizzazione fp32 e pesi bf16+fp32 occupa 2,3 TB – solo i pesi bf16 occupano 329 GB.

Set di dati:

46 lingue in 1,5 TB di testo deduplicato e pulito in modo massivo, convertito in 350 miliardi di token unici
La dimensione del vocabolario del modello è di 250.680 token
Per ulteriori dettagli, consultare il Corpus BigScience, un dataset multilingue composito da 1,6 TB

La formazione del modello BLOOM da 176 miliardi di parametri è avvenuta tra marzo e luglio 2022 e ha richiesto circa 3,5 mesi per essere completata (circa 1 milione di ore di calcolo).

Megatron-DeepSpeed

Il modello BLOOM da 176 miliardi di parametri è stato addestrato utilizzando Megatron-DeepSpeed, che è una combinazione di due tecnologie principali:

DeepSpeed è una libreria di ottimizzazione per l’apprendimento profondo che rende l’addestramento distribuito facile, efficiente ed efficace.
Megatron-LM è un framework di modelli di trasformatori di grandi dimensioni e potenti sviluppato dal team di ricerca sull’apprendimento profondo applicato di NVIDIA.

Il team di DeepSpeed ha sviluppato un’implementazione basata sulla parallelizzazione 3D combinando lo sharding ZeRO e la parallelizzazione a pipeline dalla libreria DeepSpeed con la parallelizzazione a tensori da Megatron-LM. Maggiori dettagli su ciascun componente possono essere visti nella tabella sottostante.

Si prega di notare che Megatron-DeepSpeed di BigScience è un fork del repository originale di Megatron-DeepSpeed, al quale abbiamo aggiunto varie modifiche.

Ecco una tabella che mostra quali componenti sono stati forniti da ciascun framework per addestrare BLOOM:

Si prega di notare che sia Megatron-LM che DeepSpeed hanno implementazioni di parallelismo a pipeline e di ottimizzazione BF16, ma abbiamo utilizzato quelle di DeepSpeed in quanto sono integrate con ZeRO.

Megatron-DeepSpeed implementa la parallelizzazione 3D per consentire a modelli enormi di essere addestrati in modo molto efficiente. Discutiamo brevemente i componenti 3D.

DataParallel (DP) – la stessa configurazione viene replicata più volte e ciascuna riceve un frammento dei dati. L’elaborazione avviene in parallelo e tutte le configurazioni vengono sincronizzate alla fine di ogni passaggio di addestramento.
TensorParallel (TP) – ciascun tensore viene suddiviso in più frammenti, quindi invece di avere l’intero tensore su una singola GPU, ogni frammento del tensore risiede sulla propria GPU designata. Durante l’elaborazione, ciascun frammento viene elaborato separatamente e in parallelo su diverse GPU e i risultati vengono sincronizzati alla fine del passaggio. Questo è ciò che si potrebbe definire parallelismo orizzontale, poiché la suddivisione avviene su un livello orizzontale.
PipelineParallel (PP) – il modello viene suddiviso verticalmente (a livello di layer) su più GPU, in modo che solo uno o più layer del modello siano posizionati su una singola GPU. Ogni GPU elabora in parallelo diverse fasi del processo e lavora su un piccolo insieme del batch.
Zero Redundancy Optimizer (ZeRO) – esegue anche lo sharding dei tensori in modo simile a TP, tranne che l’intero tensore viene ricostruito in tempo per una computazione in avanti o all’indietro, quindi il modello non ha bisogno di essere modificato. Supporta anche varie tecniche di spostamento per compensare la memoria limitata della GPU.

Parallelismo dei dati

La maggior parte degli utenti con solo poche GPU probabilmente conosce il DistributedDataParallel (DDP) nella documentazione di PyTorch. In questo metodo, il modello viene completamente replicato su ciascuna GPU e quindi dopo ogni iterazione tutti i modelli sincronizzano i loro stati tra loro. Questo approccio consente di accelerare l’addestramento utilizzando più risorse, ma funziona solo se il modello può adattarsi a una singola GPU.

Parallelismo dei dati ZeRO

Il parallelismo dei dati alimentato da ZeRO (ZeRO-DP) è descritto nel seguente diagramma tratto da questo post del blog

Può essere difficile capire, ma in realtà il concetto è abbastanza semplice. Si tratta dell’usuale DDP, ad eccezione del fatto che invece di replicare i parametri completi del modello, i gradienti e gli stati dell’ottimizzatore, ciascuna GPU ne memorizza solo una parte. E quindi, durante l’esecuzione, quando servono i parametri completi del layer per il layer specifico, tutte le GPU si sincronizzano per fornirsi a vicenda le parti che mancano – questo è tutto.

Questo componente è implementato da DeepSpeed.

Parallelismo dei tensori

Nel parallelismo dei tensori (TP), ogni GPU elabora solo una porzione di un tensore e aggrega solo il tensore completo per le operazioni che richiedono l’intero tensore.

In questa sezione utilizziamo concetti e diagrammi dal paper Megatron-LM: Efficient Large-Scale Language Model Training on GPU Clusters.

Il blocco principale di qualsiasi trasformatore è un nn.Linear completamente connesso seguito da un’attivazione non lineare GeLU.

Seguendo la notazione del paper Megatron, possiamo scrivere la parte del prodotto scalare come Y = GeLU(XA), dove X e Y sono i vettori di input e output, e A è la matrice dei pesi.

Se osserviamo il calcolo in forma matriciale, è facile vedere come la moltiplicazione tra matrici possa essere suddivisa tra più GPU:

Se suddividiamo la matrice dei pesi A colonna per colonna tra N GPU e eseguiamo le moltiplicazioni tra matrici XA_1 attraverso XA_n in parallelo, otterremo N vettori di output Y_1, Y_2, ..., Y_n che possono essere alimentati in modo indipendente a GeLU: . Notare che con la suddivisione della matrice Y lungo le colonne, possiamo suddividere la seconda GEMM lungo le sue righe in modo che prenda direttamente l’output di GeLU senza alcuna comunicazione aggiuntiva.

Utilizzando questo principio, possiamo aggiornare un MLP di profondità arbitraria, sincronizzando le GPU dopo ogni sequenza riga-colonna. Gli autori del paper Megatron-LM forniscono un’illustrazione utile per questo:

In questo caso, f è un operatore di identità nel passaggio in avanti e un all reduce nel passaggio all’indietro, mentre g è un all reduce nel passaggio in avanti e un’identità nel passaggio all’indietro.

La parallelizzazione dei livelli di attenzione multi-testa è ancora più semplice, poiché sono già intrinsecamente paralleli, avendo più testate indipendenti!

Considerazioni speciali: a causa dei due all reduce per livello sia nel passaggio in avanti che in quello all’indietro, TP richiede un’interconnessione molto veloce tra i dispositivi. Pertanto, non è consigliabile utilizzare TP su più di un nodo, a meno che non si disponga di una rete molto veloce. Nel nostro caso, l’interconnessione tra i nodi era molto più lenta del PCIe. In pratica, se un nodo ha 4 GPU, il grado di TP massimo è quindi 4. Se è necessario un grado di TP di 8, è necessario utilizzare nodi che abbiano almeno 8 GPU.

Questo componente è implementato da Megatron-LM. Megatron-LM ha recentemente esteso il parallelismo dei tensori includendo il parallelismo delle sequenze che suddivide le operazioni che non possono essere suddivise come sopra, ad esempio LayerNorm, lungo la dimensione della sequenza. Il paper “Reducing Activation Recomputation in Large Transformer Models” fornisce dettagli su questa tecnica. Il parallelismo delle sequenze è stato sviluppato dopo che BLOOM è stato addestrato, quindi non è stato utilizzato nell’addestramento di BLOOM.

Parallelismo del flusso di dati

Il parallelismo del flusso di dati ingenuo (naive PP) consiste nel distribuire gruppi di livelli del modello su più GPU e semplicemente spostare i dati da una GPU all’altra come se fosse una singola grande GPU composita. Il meccanismo è relativamente semplice: passare i livelli desiderati al dispositivo desiderato utilizzando il metodo .to() e ora, ogni volta che i dati entrano e escono da quei livelli, passare i dati allo stesso dispositivo del livello e lasciare il resto inalterato.

Questo attua un parallelismo del modello verticale, perché se ricordiamo come vengono disegnati la maggior parte dei modelli, le fette verticali corrispondono ai livelli. Ad esempio, se il seguente diagramma mostra un modello di 8 livelli:

===================  ===================
|  0 | 1 | 2 | 3  |  |  4 | 5 | 6 | 7  |
===================  ===================
        GPU0                 GPU1

abbiamo appena suddiviso il modello in 2 verticalmente, posizionando i livelli 0-3 sulla GPU0 e i livelli 4-7 sulla GPU1.

Ora, mentre i dati viaggiano dal livello 0 al 1, dal 1 al 2 e dal 2 al 3, questo è simile al passaggio in avanti di un modello normale su una singola GPU. Ma quando i dati devono passare dal livello 3 al livello 4, è necessario spostarsi dalla GPU0 alla GPU1, il che introduce un overhead di comunicazione. Se le GPU partecipanti si trovano sullo stesso nodo di calcolo (ad esempio, sulla stessa macchina fisica), questa copia è piuttosto veloce, ma se le GPU si trovano su nodi di calcolo diversi (ad esempio, su macchine diverse), l’overhead di comunicazione potrebbe essere significativamente maggiore.

Poi i livelli 4, 5, 6 e 7 sono come avrebbe un modello normale e quando il settimo livello è completato spesso è necessario inviare i dati indietro al livello 0 dove si trovano le etichette (o viceversa inviare le etichette all’ultimo livello). Ora la perdita può essere calcolata e l’ottimizzatore può fare il suo lavoro.

Problemi:

la principale limitazione e il motivo per cui questo viene chiamato “naive” PP, è che tutte tranne una GPU sono inattive in un dato momento. Quindi se vengono utilizzate 4 GPU, è quasi identico a quadruplicare la quantità di memoria di una singola GPU, ignorando il resto dell’hardware. Inoltre c’è il costo aggiuntivo della copia dei dati tra i dispositivi. Quindi 4 schede da 6 GB saranno in grado di ospitare la stessa dimensione di 1 scheda da 24 GB utilizzando il naive PP, tranne che quest’ultima completerà l’addestramento più velocemente, poiché non ha il costo di copia dei dati. Ma, ad esempio, se si hanno schede da 40 GB e si deve adattare un modello da 45 GB si può fare con 4 schede da 40 GB (ma a malapena a causa del gradiente e degli stati dell’ottimizzatore).
le embedding condivise potrebbero dover essere copiate avanti e indietro tra le GPU.

La Parallelizzazione del Pipeline (PP) è quasi identica al naive PP descritto in precedenza, ma risolve il problema dell’inattività della GPU, suddividendo il batch in micro-batch e creando artificialmente un pipeline, che consente alle diverse GPU di partecipare contemporaneamente al processo di calcolo.

La seguente illustrazione tratta dal paper GPipe mostra il naive PP nella parte superiore e il PP nella parte inferiore:

È facile vedere dal diagramma inferiore come il PP abbia meno zone morte, dove le GPU sono inattive. Le parti inattive sono indicate come “bubble”.

Entrambe le parti del diagramma mostrano una parallelizzazione di grado 4. Cioè 4 GPU partecipano al pipeline. Quindi c’è il percorso in avanti di 4 stadi del pipeline F0, F1, F2 e F3 e poi il percorso inverso di ritorno nell’ordine B3, B2, B1 e B0.

PP introduce un nuovo iperparametro da regolare chiamato chunks. Definisce quanti chunk di dati vengono inviati in sequenza attraverso lo stesso stadio del pipeline. Ad esempio, nel diagramma inferiore, si può vedere che chunks=4. La GPU0 esegue lo stesso percorso in avanti per i chunk 0, 1, 2 e 3 (F0,0, F0,1, F0,2, F0,3) e poi attende che le altre GPU facciano il loro lavoro e solo quando il loro lavoro sta per essere completato, la GPU0 inizia nuovamente a lavorare facendo il percorso inverso per i chunk 3, 2, 1 e 0 (B0,3, B0,2, B0,1, B0,0).

Si noti che concettualmente questo è lo stesso concetto degli step di accumulo del gradiente (GAS). PyTorch utilizza chunks, mentre DeepSpeed si riferisce allo stesso iperparametro come GAS.

A causa dei chunk, PP introduce il concetto di micro-batch (MBS). DP suddivide la dimensione del batch globale in mini-batch, quindi se si ha un grado di DP di 4, una dimensione di batch globale di 1024 viene suddivisa in 4 mini-batch di 256 ciascuno (1024/4). E se il numero di chunks (o GAS) è 32, si ottiene una dimensione di micro-batch di 8 (256/32). Ogni stadio del pipeline lavora con un singolo micro-batch alla volta.

Per calcolare la dimensione del batch globale dell’impostazione DP + PP, si fa quindi: mbs*chunks*dp_degree (8*32*4=1024).

Torniamo al diagramma.

Con chunks=1 si ottiene il naive PP, che è molto inefficiente. Con un valore di chunks molto grande si ottengono dimensioni di micro-batch molto piccole che potrebbero non essere molto efficienti. Quindi bisogna sperimentare per trovare il valore che porta all’utilizzo più efficiente delle GPU.

Sebbene il diagramma mostri che c’è una bolla di tempo “morta” che non può essere parallelizzata perché l’ultimo stadio forward deve attendere che il backward completi il pipeline, lo scopo di trovare il valore migliore per chunks è quello di consentire un elevato utilizzo simultaneo delle GPU su tutte le GPU partecipanti, riducendo così la dimensione della bolla.

Questo meccanismo di programmazione è noto come tutto avanti tutto indietro. Alcune alternative sono uno avanti uno indietro e uno avanti uno indietro intercalato.

Anche se sia Megatron-LM che DeepSpeed hanno la propria implementazione del protocollo PP, Megatron-DeepSpeed utilizza l’implementazione di DeepSpeed poiché è integrata con altri aspetti di DeepSpeed.

Un altro problema importante qui è la dimensione della matrice di embedding delle parole. Normalmente una matrice di embedding delle parole consuma meno memoria rispetto al blocco del transformer, ma nel nostro caso con un vocabolario enorme di 250k, il layer di embedding richiedeva 7.2GB di pesi in bf16 e il blocco del transformer solo 4.9GB. Pertanto, abbiamo dovuto istruire Megatron-Deepspeed a considerare il layer di embedding come un blocco del transformer. Quindi avevamo una pipeline di 72 layer, di cui 2 erano dedicati all’embedding (il primo e l’ultimo). Questo ha permesso di bilanciare il consumo di memoria della GPU. Se non l’avessimo fatto, avremmo avuto il primo e l’ultimo stadio che consumano la maggior parte della memoria della GPU e il 95% delle GPU avrebbe utilizzato molto meno memoria e quindi l’addestramento sarebbe stato inefficiente.

DP+PP

Il seguente diagramma del tutorial sulla pipeline di DeepSpeed mostra come combinare DP con PP.

Qui è importante vedere come il rank DP 0 non vede la GPU2 e il rank DP 1 non vede la GPU3. Per DP ci sono solo le GPU 0 e 1, dove vengono forniti i dati come se ci fossero solo 2 GPU. La GPU0 “segretamente” scarica parte del suo carico sulla GPU2 usando PP. E la GPU1 fa lo stesso arruolando l’aiuto della GPU3.

Poiché ogni dimensione richiede almeno 2 GPU, qui avresti bisogno di almeno 4 GPU.

DP+PP+TP

Per ottenere un addestramento ancora più efficiente, PP viene combinato con TP e DP, che viene chiamato parallelismo 3D. Questo può essere visto nel seguente diagramma.

Questo diagramma proviene da un post sul blog 3D parallelism: Scaling to trillion-parameter models, che è anche una buona lettura.

Poiché ogni dimensione richiede almeno 2 GPU, qui avresti bisogno di almeno 8 GPU per il parallelismo 3D completo.

ZeRO DP+PP+TP

Una delle principali caratteristiche di DeepSpeed è ZeRO, che è un’estensione super-scalabile di DP. È già stato discusso in ZeRO Data Parallelism. Normalmente è una funzionalità autonoma che non richiede PP o TP. Ma può essere combinato con PP e TP.

Quando ZeRO-DP è combinato con PP (e opzionalmente TP), di solito abilita solo ZeRO stage 1, che divide solo gli stati dell’ottimizzatore. ZeRO stage 2 divide anche i gradienti e lo stage 3 divide anche i pesi del modello.

Anche se teoricamente è possibile utilizzare ZeRO stage 2 con Pipeline Parallelism, avrà un impatto negativo sulle prestazioni. Sarebbe necessario un ulteriore raggruppamento collettivo per ogni micro-batch per aggregare i gradienti prima della divisione, il che aggiunge un overhead di comunicazione potenzialmente significativo. Per natura del Parallelismo della Pipeline, vengono utilizzati micro-batch piccoli e l’attenzione è incentrata nel cercare di bilanciare l’intensità aritmetica (dimensione del micro-batch) con la riduzione della Pipeline (numero di micro-batch). Pertanto, questi costi di comunicazione avranno un impatto negativo.

Inoltre, ci sono già meno layer rispetto al normale a causa di PP e quindi i risparmi di memoria non saranno enormi. PP riduce già la dimensione del gradiente di 1/PP, quindi i risparmi di divisione del gradiente su quello sono meno significativi rispetto a DP puro.

ZeRO stage 3 può essere utilizzato anche per addestrare modelli a questa scala, tuttavia richiede più comunicazione rispetto all’implementazione parallela 3D di DeepSpeed. Dopo una valutazione attenta nel nostro ambiente, che è avvenuta un anno fa, abbiamo scoperto che il parallelismo 3D di Megatron-DeepSpeed ha fornito le migliori prestazioni. Da allora le prestazioni di ZeRO stage 3 sono migliorate notevolmente e se dovessimo valutarle oggi forse avremmo scelto lo stage 3 al suo posto.

BF16Optimizer

Allenare modelli LLM enormi in FP16 è un no-no.

Lo abbiamo dimostrato a noi stessi dedicando diversi mesi all’addestramento di un modello 104B che, come si può vedere da tensorboard, è stato un completo fallimento. Abbiamo imparato molte cose mentre lottavamo contro l’lm-loss sempre divergente:

e abbiamo ricevuto lo stesso consiglio anche dalle squadre di Megatron-LM e DeepSpeed dopo aver addestrato il modello 530B. La recente versione di OPT-175B riportava anche che hanno avuto molte difficoltà nell’addestramento in FP16.

Quindi, a gennaio, sapendo che avremmo addestrato su A100 che supporta il formato BF16, Olatunji Ruwase ha sviluppato un BF16Optimizer che abbiamo usato per addestrare BLOOM.

Se non sei familiare con questo formato dati, dai un’occhiata alla disposizione dei bit. La chiave del formato BF16 è che ha la stessa esponente di FP32 e quindi non soffre di overflow come FP16! Con FP16, che ha un intervallo numerico massimo di 64k, puoi moltiplicare solo numeri piccoli. Ad esempio, puoi fare 250*250=62500, ma se provassi a fare 255*255=65025 avresti un overflow, che è ciò che causa i principali problemi durante l’addestramento. Questo significa che i tuoi pesi devono rimanere piccoli. Una tecnica chiamata loss scaling può aiutare con questo problema, ma l’intervallo limitato di FP16 è comunque un problema quando i modelli diventano molto grandi.

BF16 non ha tali problemi, puoi facilmente fare 10_000*10_000=100_000_000 e non c’è problema.

Ovviamente, poiché BF16 e FP16 hanno la stessa dimensione di 2 byte, non si ottiene un pranzo gratis e si paga con una precisione molto scarsa quando si usa BF16. Tuttavia, se ricordi che l’addestramento utilizzando la discesa del gradiente stocastico e le sue variazioni è una sorta di camminata inciampante, quindi se non hai subito la direzione perfetta non è un problema, ti correggerai nei passaggi successivi.

Indipendentemente dall’uso di BF16 o FP16, c’è anche una copia dei pesi che è sempre in FP32: è questo che viene aggiornato dall’ottimizzatore. Quindi i formati a 16 bit vengono utilizzati solo per i calcoli, l’ottimizzatore aggiorna i pesi FP32 con piena precisione e poi li converte nel formato a 16 bit per la successiva iterazione.

Tutti i componenti di PyTorch sono stati aggiornati per assicurarsi che eseguano qualsiasi accumulo in FP32, quindi non ci sono perdite in quel punto.

Un problema cruciale è l’accumulo del gradiente, ed è una delle principali caratteristiche del parallelismo a pipeline poiché i gradienti da ogni microbatch di elaborazione vengono accumulati. È fondamentale implementare l’accumulo del gradiente in FP32 per mantenere l’addestramento preciso, ed è quello che fa BF16Optimizer.

Oltre ad altre migliorie, riteniamo che l’uso dell’addestramento a precisione mista BF16 abbia trasformato un potenziale incubo in un processo relativamente fluido, come si può osservare dal seguente grafico dell’lm-loss:

Kernel CUDA fusi

La GPU esegue due cose. Può copiare i dati da/per la memoria e eseguire calcoli su quei dati. Mentre la GPU sta copiando, le unità di calcolo della GPU sono inattive. Se vogliamo utilizzare efficientemente la GPU, vogliamo ridurre al minimo il tempo di inattività.

Un kernel è un insieme di istruzioni che implementa una specifica operazione di PyTorch. Ad esempio, quando chiami torch.add, passa attraverso un dispatcher di PyTorch che guarda il tensore/i di input e varie altre cose e decide quale codice deve eseguire, quindi lo esegue. Un kernel CUDA è un’implementazione specifica che utilizza la libreria API CUDA e può essere eseguita solo su GPU NVIDIA.

Ora, quando si istruisce la GPU a calcolare c = torch.add(a, b); e = torch.max([c,d]), un approccio ingenuo, e ciò che PyTorch farà a meno di istruzioni diverse, è lanciare due kernel separati, uno per eseguire l’addizione di a e b e un altro per trovare il valore massimo tra c e d. In questo caso, la GPU recupera dalla memoria a e b, esegue l’addizione e quindi copia il risultato nella memoria. Poi recupera c e d e esegue l’operazione di max e copia nuovamente il risultato nella memoria.

Se unissimo queste due operazioni, cioè mettendole in un unico “kernel fuso”, e lanciassimo solo quel kernel, non copieremmo il risultato intermedio c nella memoria, ma lo lasciamo nei registri della GPU e abbiamo solo bisogno di recuperare d per completare l’ultima computazione. Questo risparmia molta overhead e impedisce l’idling della GPU, rendendo l’intera operazione molto più efficiente.

I kernel fusi sono proprio così. Sostituiscono principalmente molteplici calcoli discreti e spostamenti dei dati da/per la memoria in calcoli fusi che hanno pochissimi spostamenti di memoria. Inoltre, alcuni kernel fusi riscrivono la matematica in modo che certi gruppi di calcoli possano essere eseguiti più velocemente.

Per addestrare BLOOM in modo veloce ed efficiente, è stato necessario utilizzare diversi kernel CUDA fusi personalizzati forniti da Megatron-LM. In particolare, c’è un kernel ottimizzato per eseguire LayerNorm così come kernel per fondere varie combinazioni di operazioni di scaling, mascheramento e softmax. L’aggiunta di un termine di bias è anche fusa con l’operazione GeLU utilizzando la funzionalità JIT di PyTorch. Queste operazioni sono tutte legate alla memoria, quindi è importante fonderle per massimizzare la quantità di calcoli effettuati una volta che un valore è stato recuperato dalla memoria. Quindi, ad esempio, l’aggiunta del termine di bias durante l’operazione di GeLU legata alla memoria non aggiunge tempo aggiuntivo. Tutti questi kernel sono disponibili nel repository Megatron-LM.

Insiemi di dati

Un’altra caratteristica importante di Megatron-LM è il caricatore di dati efficiente. Durante l’avvio dell’addestramento iniziale, ogni set di dati viene suddiviso in campioni della lunghezza di sequenza richiesta (2048 per BLOOM) e viene creato un indice per numerare ogni campione. In base ai parametri di addestramento, viene calcolato il numero di epoche per un set di dati e viene creato un ordine per tante epoche e quindi mescolato. Ad esempio, se un set di dati ha 10 campioni e dovrebbe essere attraversato due volte, il sistema dispone prima gli indici dei campioni in ordine [0, ..., 9, 0, ..., 9] e quindi mescola quell’ordine per creare l’ordine globale finale per il set di dati. Si noti che questo significa che l’addestramento non attraverserà semplicemente l’intero set di dati e poi si ripeterà, è possibile vedere lo stesso campione due volte prima di vedere un altro campione del tutto, ma alla fine dell’addestramento il modello avrà visto ogni campione due volte. Questo aiuta a garantire una curva di addestramento uniforme durante l’intero processo di addestramento. Questi indici, inclusi gli offset nel set di dati di base di ogni campione, vengono salvati in un file per evitare di ricalcolarli ogni volta che viene avviato un processo di addestramento. Diversi di questi set di dati possono quindi essere mescolati con pesi variabili nel set di dati finale visualizzato dal processo di addestramento.

LayerNorm di incorporamento

Mentre stavamo cercando di impedire a 104B di divergere, abbiamo scoperto che l’aggiunta di un ulteriore LayerNorm subito dopo l’incorporamento della prima parola rendeva l’addestramento molto più stabile.

Questa intuizione è emersa sperimentando con bitsandbytes che contiene un StableEmbedding che è un’incorporamento normale con layernorm e utilizza una inizializzazione uniforme xavier.

Codifica posizionale

Abbiamo anche sostituito l’incorporamento posizionale usuale con una codifica AliBi – basata sul paper: Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation, che consente di fare previsioni per sequenze di input più lunghe rispetto a quelle su cui il modello è stato addestrato. Quindi, anche se addestriamo sequenze con lunghezza 2048, il modello può gestire anche sequenze molto più lunghe durante l’inferenza.

Difficoltà di addestramento

Con l’architettura, l’hardware e il software in posizione, siamo stati in grado di iniziare l’addestramento all’inizio di marzo 2022. Tuttavia, non è stato tutto facile da lì. In questa sezione discutiamo alcune delle principali difficoltà che abbiamo incontrato.

C’erano molti problemi da risolvere prima dell’inizio dell’addestramento. In particolare, abbiamo riscontrato diversi problemi che si manifestavano solo una volta che abbiamo iniziato ad addestrare su 48 nodi e non si verificano su piccola scala. Ad esempio, era necessario impostare CUDA_LAUNCH_BLOCKING=1 per evitare che il framework si blocchi, e abbiamo dovuto suddividere i gruppi di ottimizzatori in gruppi più piccoli, altrimenti il framework si bloccava di nuovo. Puoi leggere maggiori dettagli su questi problemi nel capitolo dei precedenti di addestramento.

Il tipo principale di problema riscontrato durante l’addestramento sono stati i guasti hardware. Essendo questo un nuovo cluster con circa 400 GPU, in media avevamo 1-2 guasti di GPU a settimana. Salvavamo un checkpoint ogni 3 ore (100 iterazioni), quindi in media avremmo perso 1,5 ore di addestramento a causa di un crash hardware. Gli amministratori di Jean Zay avrebbero quindi sostituito le GPU difettose e riportato il nodo in funzione. Nel frattempo avevamo nodi di backup da utilizzare al loro posto.

Ci siamo imbattuti in una varietà di altri problemi che hanno causato tempi di inattività di 5-10 ore diverse volte, alcuni legati a un bug di deadlock in PyTorch, altri dovuti all’esaurimento dello spazio su disco. Se sei curioso di conoscere i dettagli specifici, consulta le cronache di addestramento .

Abbiamo pianificato tutti questi tempi di inattività quando abbiamo deciso sulla fattibilità di addestrare questo modello – abbiamo scelto le dimensioni del modello in base a quella fattibilità e alla quantità di dati che volevamo che il modello consumasse. Con tutti i tempi di inattività siamo riusciti a completare l’addestramento nel tempo stimato. Come accennato in precedenza, ci sono volute circa 1 milione di ore di calcolo per completare l’operazione.

Un altro problema è stato che SLURM non è stato progettato per essere utilizzato da un team di persone. Un lavoro SLURM è di proprietà di un singolo utente e se non sono presenti, gli altri membri del gruppo non possono fare nulla sul lavoro in esecuzione. Abbiamo sviluppato una soluzione alternativa di interruzione che consentiva agli altri utenti del gruppo di interrompere il processo corrente senza richiedere la presenza dell’utente che ha avviato il processo. Ciò ha funzionato bene nel 90% dei casi. Se i progettisti di SLURM leggono questo – per favore aggiungete un concetto di gruppi Unix, in modo che un lavoro SLURM possa essere di proprietà di un gruppo.

Poiché l’addestramento avveniva 24/7, avevamo bisogno che qualcuno fosse disponibile – ma poiché avevamo persone sia in Europa che sulla costa occidentale del Canada, non c’era bisogno che qualcuno portasse un cercapersone, ci sovrapponiamo semplicemente bene. Naturalmente, qualcuno doveva monitorare l’addestramento anche nei fine settimana. Abbiamo automatizzato la maggior parte delle cose, inclusa la ripresa dai crash hardware, ma a volte era necessario anche un intervento umano.

Conclusioni

La parte più difficile e intensa dell’addestramento è stata i 2 mesi che hanno preceduto l’inizio dell’addestramento. Siamo stati sotto molta pressione per iniziare l’addestramento il prima possibile, poiché l’allocazione delle risorse era limitata nel tempo e non avevamo accesso ad A100 fino all’ultimo momento. Quindi è stato un periodo molto difficile, considerando che il BF16Optimizer è stato scritto all’ultimo momento e abbiamo dovuto risolvere vari bug. E come spiegato nella sezione precedente, abbiamo scoperto nuovi problemi che si sono manifestati solo una volta che abbiamo iniziato l’addestramento su 48 nodi e non si verificheranno su piccola scala.

Ma una volta risolti questi problemi, l’addestramento stesso è stato sorprendentemente fluido e senza problemi importanti. La maggior parte del tempo avevamo una persona che monitorava l’addestramento e solo alcune volte diverse persone sono state coinvolte per risolvere i problemi. Abbiamo ricevuto un ottimo supporto dall’amministrazione di Jean Zay che ha prontamente affrontato la maggior parte delle esigenze emerse durante l’addestramento.

Nel complesso è stata un’esperienza super-intensa ma molto gratificante.

L’addestramento di modelli linguistici di grandi dimensioni è ancora una sfida, ma speriamo che costruendo e condividendo questa tecnologia in modo aperto, gli altri possano costruire sulla base della nostra esperienza.

Risorse

Link importanti

documento principale di addestramento
tensorboard
script di addestramento slurm
cronache di addestramento

Articoli e documenti

Non avremmo potuto spiegare tutto in dettaglio in questo articolo, quindi se la tecnologia presentata qui ha suscitato la tua curiosità e desideri saperne di più, ecco i documenti da leggere:

Megatron-LM:

Efficient Large-Scale Language Model Training on GPU Clusters .
Reducing Activation Recomputation in Large Transformer Models

DeepSpeed:

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
ZeRO-Offload: Democratizing Billion-Scale Model Training
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
DeepSpeed: Extreme-scale model training for everyone

Joint Megatron-LM and Deepspeeed:

Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model .

ALiBi:

Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation
What Language Model to Train if You Have One Million GPU Hours? – là troverai gli esperimenti che ci hanno portato a scegliere ALiBi.

BitsNBytes:

8-bit Optimizers via Block-wise Quantization (in the context of Embedding LayerNorm but the rest of the paper and the technology is amazing – the only reason were weren’t using the 8-bit optimizer is because we were already saving the optimizer memory with DeepSpeed-ZeRO).

Crediti del blog

Un enorme ringraziamento alle seguenti gentili persone che hanno posto buone domande e hanno contribuito a migliorare la leggibilità dell’articolo (elencate in ordine alfabetico): Britney Muller, Douwe Kiela, Jared Casper, Jeff Rasley, Julien Launay, Leandro von Werra, Omar Sanseviero, Stefan Schweter e Thomas Wang.

I principali grafici sono stati creati da Chunte Lee.