ZeRO++: Potenziare l'Addestramento dei Grandi Modelli

Il deep learning è diventato uno strumento fondamentale in tanti campi, come l'analisi delle immagini, il trattamento del linguaggio e il riconoscimento vocale. Col tempo, questi modelli sono migliorati sia nella qualità che nella capacità. Uno dei motivi principali di questo miglioramento è l'aumento delle dimensioni. I modelli più grandi tendono a performare meglio perché riescono a catturare schemi più complessi. Ad esempio, le dimensioni dei modelli sono passate da 100 milioni di parametri a oltre 500 miliardi di parametri in pochi anni.

Tuttavia, addestrare questi modelli enormi presenta sfide significative. I requisiti di memoria e calcolo sono aumentati, rendendo impossibile per una singola GPU gestire il carico di lavoro. Per addestrare modelli grandi in modo efficace, è necessario unire la potenza di più GPU. Due metodi comuni per raggiungere questo obiettivo sono il Parallelismo 3D e lo Zero Redundancy Optimizer (ZeRO).

Il parallelismo 3D combina diverse tecniche per suddividere il carico di lavoro tra molte GPU. Anche se permette al sistema di utilizzare le risorse in modo efficiente, aumenta la complessità. I data scientist spesso affrontano sfide quando cercano di adattare il loro codice per il parallelismo 3D.

Al contrario, ZeRO semplifica il processo. Invece di replicare gli stati del modello su tutte le GPU, divide i parametri del modello in parti più piccole che vengono condivise tra le GPU. Questo approccio utilizza meno memoria e comporta costi di comunicazione più bassi, mantenendo comunque ottime prestazioni.

Nonostante i suoi vantaggi, ZeRO ha delle limitazioni, soprattutto in situazioni dove il numero di GPU è elevato o quando la larghezza di banda di rete disponibile è bassa. Questi fattori possono portare a colli di bottiglia nella comunicazione che rallentano l'addestramento.

#Limitazioni di ZeRO

ZeRO è popolare perché è facile da usare e si scalda bene. Tuttavia, il sovraccarico di comunicazione può limitare l'efficienza in due scenari chiave: cluster a bassa larghezza di banda e piccole dimensioni del batch.

Nei cluster a bassa larghezza di banda, anche se macchine ad alte prestazioni possono avere ottime connessioni interne, le connessioni tra diverse macchine spesso hanno una larghezza di banda inferiore. Questo porta a colli di bottiglia nella comunicazione, che riducono le prestazioni del processo di addestramento.

Anche in reti ad alta larghezza di banda, se il numero di GPU è molto grande, la dimensione del batch per ciascuna GPU può diventare troppo piccola. Questo abbassa ulteriormente il rapporto calcolo/comunicazione e crea ritardi aggiuntivi nella comunicazione.

Sebbene ZeRO abbia molti benefici, i metodi esistenti per ottimizzare la comunicazione nell'addestramento distribuito non funzionano bene con la sua architettura perché si basano sulla replicazione degli stati del modello.

Per affrontare questi problemi, introduciamo un nuovo sistema chiamato ZeRO++. Questo sistema riduce il sovraccarico di comunicazione e migliora le prestazioni durante l'addestramento. Di seguito, descriviamo le tecniche utilizzate in ZeRO++ per affrontare le sfide comunicative affrontate da ZeRO.

#Ottimizzazioni della Comunicazione in ZeRO++

ZeRO++ si concentra su tre modi principali per ridurre il sovraccarico di comunicazione:

#1. Comunicazione Pesi Bloccati-Quantizzati

Durante il passaggio in avanti, ZeRO raccoglie tutti i parametri necessari per l'addestramento tramite un'operazione di raccolta totale. Questo può creare molto traffico di comunicazione. Per ridurre questo, utilizziamo una tecnica chiamata quantizzazione basata su blocchi. Questo riduce la dimensione dei dati inviati. Invece di inviare dati in formato ad alta precisione, li convertiamo in un formato a bassa precisione. Questo dimezza il volume della comunicazione senza perdere molta accuratezza.

#2. Partizionamento Gerarchico dei Pesi

Durante il passaggio all'indietro, ZeRO deve raccogliere nuovamente i pesi, il che può creare un significativo sovraccarico di comunicazione. Per affrontare questo, manteniamo una copia completa dei pesi su ogni macchina. Questo ci consente di eseguire operazioni di raccolta all'interno di una singola macchina piuttosto che tra più macchine, rendendo il processo molto più veloce grazie a una banda più ampia.

#3. Comunicazione Gradiente Quantizzati

Quando si tratta di ridurre i costi di comunicazione durante il passaggio all'indietro, dobbiamo anche gestire i gradienti. L'uso di tecniche standard per ridurre la comunicazione può danneggiare l'accuratezza dell'addestramento. Quindi, abbiamo sviluppato un nuovo approccio che ci consente di comunicare i gradienti a bassa precisione mantenendo intatta l'accuratezza. Prima comprimiamo i gradienti e poi li ripristiniamo alla loro precisione originale prima di eseguire i calcoli necessari.

Combinando queste tre tecniche, ZeRO++ riduce significativamente il volume della comunicazione e migliora l'efficienza durante l'addestramento.

#Valutazione di ZeRO++

Abbiamo implementato ZeRO++ e condotto ampie valutazioni dimostrando tre risultati chiave:

1. Scalabilità a modelli grandi su centinaia di GPU.
2. Accelerazione costante rispetto al ZeRO originale.
3. Prestazioni migliorate anche in condizioni a bassa larghezza di banda.

Attraverso i nostri test, abbiamo scoperto che ZeRO++ raggiunge fino a 2.16 volte migliore throughput con 384 GPU. Abbiamo anche analizzato come ciascuna delle tre ottimizzazioni contribuisca al miglioramento complessivo e dimostrato che ZeRO++ non influisce negativamente sulla convergenza del modello.

#Contributi Chiave

I principali contributi di questo lavoro includono:

• I pesi quantizzati bloccati riducono il volume della comunicazione per la raccolta dei pesi del 50%.
• Il partizionamento gerarchico dei pesi elimina completamente la comunicazione inter-nodo durante i passaggi all'indietro.
• Un nuovo metodo di comunicazione dei gradienti quantizzati riduce il volume della comunicazione per i gradienti del 75%.
• Integrazione efficiente di queste tecniche nel framework esistente di ZeRO, consentendo prestazioni migliori.

#Contesto e Lavori Correlati

Per capire meglio ZeRO++, è essenziale guardare a come i grandi modelli vengono addestrati su più GPU. Ci sono principalmente tre forme di parallelismo nell'addestramento distribuito: parallelismo dei dati, parallelismo a pipeline e parallelismo dei tensori.

Il parallelismo dei dati funziona bene quando il modello può adattarsi a una singola GPU. In questa configurazione, ogni GPU possiede una copia completa del modello e elabora dati di input diversi. Quando i modelli non possono adattarsi a una GPU, entra in gioco il parallelismo dei modelli. Questo metodo suddivide il modello tra le GPU, sia per livelli (parallelismo a pipeline) che lungo le dimensioni del tensore (parallelismo dei tensori).

Il parallelismo 3D combina questi metodi ma richiede cambiamenti significativi al codice esistente. D'altra parte, ZeRO offre un'alternativa semplice partizionando i parametri del modello tra le GPU, il che minimizza la comunicazione massimizzando l'efficienza della memoria.

Tuttavia, il metodo standard ZeRO ha elevate richieste di comunicazione, specialmente in ambienti di grande scala. Le soluzioni esistenti si sono concentrate sulla riduzione del sovraccarico di comunicazione, ma spesso non si integrano bene con la struttura di ZeRO.

#Come ZeRO++ Affronta il Sovraccarico di Comunicazione

ZeRO++ combina diverse strategie per migliorare l'efficienza comunicativa in ZeRO.

#Comunicazione Pesi Quantizzati

Durante il passaggio in avanti, ZeRO si basa sulla raccolta dei pesi, che può essere pesante in termini di comunicazione. Applicando la quantizzazione basata su blocchi, comprimiamo i dati inviati. Questa tecnica adatta i pesi a una bassa precisione, riducendo la dimensione mantenendo un'accuratezza accettabile.

#Partizionamento Gerarchico

Per il passaggio all'indietro, evitiamo la raccolta inter-nodo mantenendo i pesi completi su ogni macchina. Questo approccio ci consente di eseguire la raccolta localmente all'interno di ogni macchina, che è notevolmente più veloce della comunicazione tra nodi.

#Comunicazione Gradiente Quantizzati

Per i gradienti nel passaggio all'indietro, abbiamo creato un nuovo metodo che comprime i gradienti mantenendo la loro accuratezza. Applichiamo prima la quantizzazione durante la fase di comunicazione e poi ripristiniamo la precisione in seguito per i calcoli.

Tutti questi metodi affrontano diversi punti all'interno del processo comunicativo, rendendo più facile condividere i dati tra GPU mantenendo alte prestazioni.

#Valutazione delle Prestazioni di ZeRO++

Abbiamo valutato ZeRO++ utilizzando vari modelli e configurazioni, concentrandoci sulla sua scalabilità, velocità e efficienza complessiva. La nostra implementazione ha mostrato risultati forti su diverse piattaforme e dimensioni di modelli.

#Scalabilità

ZeRO++ ha funzionato efficacemente su fino a 384 GPU con modelli grandi. In particolare, ha raggiunto una percentuale significativa delle prestazioni di picco disponibili con l'hardware utilizzato. Questo dimostra che i nuovi metodi di comunicazione consentono di addestrare modelli più grandi in modo efficiente.

#Miglioramento della Velocità

Una delle scoperte chiave è stata che ZeRO++ ha superato il ZeRO originale attraverso varie dimensioni di modelli e configurazioni. In alcuni casi, abbiamo osservato miglioramenti prestazionali fino a 2.4 volte rispetto alla base. Questo è particolarmente importante per scenari con larghezza di banda limitata, dove i metodi precedenti faticavano a mantenere l'efficienza.

#Prestazioni a Bassa Larghezza di Banda

Inoltre, ZeRO++ ha mantenuto prestazioni solide quando testato in reti a bassa larghezza di banda, dimostrando la sua robustezza. Questo apre opportunità per addestrare modelli grandi in ambienti che potrebbero non avere configurazioni di rete ottimali.

#Conclusione

L'introduzione di ZeRO++ rappresenta un significativo avanzamento nella comunicazione collettiva per addestrare modelli grandi in modo efficiente. Concentrandosi sulla riduzione del sovraccarico attraverso tecniche di quantizzazione e partizionamento intelligenti, abbiamo dimostrato che è possibile mantenere alte prestazioni, anche su risorse di rete limitate.

In sintesi, ZeRO++ offre un approccio promettente per rendere l'addestramento di modelli su larga scala più accessibile ed efficiente per ricercatori e professionisti. Con questi miglioramenti, prevediamo capacità migliori per grandi modelli di linguaggio e oltre, promuovendo applicazioni più ampie in vari settori del deep learning.