Ottieni risultati migliori e addestra più velocemente con ZeRO tramite DeepSpeed e FairScale

Migliora risultati e addestramento più veloce con ZeRO, DeepSpeed e FairScale.

Un post ospite di Stas Bekman, collega di Hugging Face

Dato che i modelli di Machine Learning recenti stanno crescendo molto più velocemente della quantità di memoria GPU aggiunta alle schede appena rilasciate, molti utenti non sono in grado di allenare o anche solo caricare alcuni di quei modelli enormi sul proprio hardware. Mentre c’è uno sforzo in corso per ridurre la dimensione di alcuni di quei modelli enormi in modo da renderli più gestibili, tale sforzo non sta producendo modelli abbastanza piccoli abbastanza velocemente.

Nell’autunno del 2019, Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase e Yuxiong He hanno pubblicato un articolo: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, che contiene una moltitudine di nuove idee ingeniose su come fare in modo che l’hardware faccia molto di più di quanto si pensasse fosse possibile. Poco tempo dopo è stato rilasciato DeepSpeed, che ha dato al mondo l’implementazione open source della maggior parte delle idee in quel documento (alcune idee sono ancora in lavorazione) e contemporaneamente un team di Facebook ha rilasciato FairScale, che ha anche implementato alcune delle idee principali del documento ZeRO.

Se utilizzi il Trainer di Hugging Face, a partire dalla versione 4.2.0 di transformers hai il supporto sperimentale per le funzionalità ZeRO di DeepSpeed e FairScale. I nuovi argomenti della riga di comando --sharded_ddp e --deepspeed del Trainer forniscono rispettivamente l’integrazione di FairScale e DeepSpeed. Qui trovi la documentazione completa .

In questo post del blog verrà descritto come puoi beneficiare di ZeRO indipendentemente dal fatto che tu possieda un singolo GPU o un intero stack di essi.

• Modelli TensorFlow più veloci in Hugging Face Transformers
• Hugging Face su PyTorch / XLA TPUs
• Recupero potenziato dalla generazione con Huggingface Transformers e Ray

Facciamo un piccolo esperimento di addestramento con un compito di traduzione, utilizzando un modello t5-large e lo script finetune_trainer.py che puoi trovare in examples/seq2seq nel repository GitHub di transformers.

Abbiamo 2x 24GB (Titan RTX) GPU con cui fare i test.

Questi sono solo benchmark di prova di concetto, quindi sicuramente le cose possono essere migliorate ulteriormente, quindi effettueremo il benchmark su un piccolo campione di 2000 elementi per l’addestramento e 500 elementi per la valutazione per effettuare i confronti. La valutazione di default effettua una ricerca a fascio di dimensioni 4, quindi è più lenta dell’addestramento con lo stesso numero di campioni, ecco perché sono stati utilizzati 4 volte meno elementi di valutazione in questi test.

Ecco i principali argomenti della riga di comando del nostro punto di riferimento:

export BS=16
python -m torch.distributed.launch --nproc_per_node=2 ./finetune_trainer.py \
--model_name_or_path t5-large --n_train 2000 --n_val 500 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro [...]

Stiamo utilizzando solo il Parallelismo dei Dati Distribuito (DDP) e nient’altro per aumentare le prestazioni del punto di riferimento. Sono riuscito a utilizzare una dimensione di batch (BS) di 16 prima di incorrere in un errore di Out of Memory (OOM).

Nota che per semplicità e per renderlo più facile da capire, ho mostrato solo gli argomenti della riga di comando importanti per questa dimostrazione. Troverai la riga di comando completa in questo post .

Successivamente, eseguiremo nuovamente il benchmark ogni volta aggiungendo uno dei seguenti:

1. --fp16
2. --sharded_ddp (fairscale)
3. --sharded_ddp --fp16 (fairscale)
4. --deepspeed senza l’uso della CPU
5. --deepspeed con l’uso della CPU

Poiché l’ottimizzazione chiave qui è che ogni tecnica utilizza in modo più efficiente la RAM della GPU, cercheremo di aumentare continuamente la dimensione del batch e ci aspettiamo che l’addestramento e la valutazione vengano completati più velocemente (mantenendo costanti o addirittura migliorando alcune metriche, ma non ci concentreremo su queste qui).

Ricorda che le fasi di addestramento e valutazione sono molto diverse tra loro, perché durante l’addestramento vengono modificati i pesi del modello, vengono calcolati i gradienti e vengono memorizzati gli stati dell’ottimizzatore. Durante la valutazione, nulla di tutto ciò accade, ma in questo particolare compito di traduzione il modello cercherà di cercare la migliore ipotesi, quindi in realtà deve effettuare più prove prima di essere soddisfatto. Ecco perché non è veloce, soprattutto quando un modello è grande.

Guardiamo i risultati di queste sei esecuzioni di test:

È facile vedere che sia FairScale che DeepSpeed offrono grandi miglioramenti rispetto alla base, sia nel tempo totale di addestramento e valutazione, ma anche nella dimensione del batch. DeepSpeed implementa più magia al momento della scrittura e sembra essere il vincitore a breve termine, ma FairScale è più facile da implementare. Per DeepSpeed è necessario scrivere un semplice file di configurazione e modificare il launcher della riga di comando, mentre per FairScale è sufficiente aggiungere l’argomento della riga di comando --sharded_ddp, quindi potresti volerlo provare prima in quanto è la soluzione più semplice.

Seguendo la regola dell’80:20, ho dedicato solo alcune ore a questi benchmark e non ho cercato di sfruttare ogni MB e secondo raffinando gli argomenti della riga di comando e la configurazione, poiché è abbastanza ovvio dalla semplice tabella cosa vorresti provare successivamente. Quando ti troverai di fronte a un progetto reale che verrà eseguito per ore e forse giorni, dedica sicuramente più tempo per assicurarti di utilizzare gli iperparametri più ottimali per completare il tuo lavoro più velocemente e a un costo minimo.

Se desideri sperimentare questo benchmark da solo o desideri conoscere maggiori dettagli sull’hardware e il software utilizzati per eseguirlo, ti preghiamo di fare riferimento a questo post .

Mentre FairScale ci offre un incremento solo con più GPU, DeepSpeed ha un vantaggio anche per coloro che dispongono di una singola GPU.

Proviamo l’impossibile: addestriamo t5-3b su una scheda RTX-3090 da 24 GB.

Prima cerchiamo di ottimizzare l’enorme t5-3b utilizzando la normale configurazione a singola GPU:

export BS=1
CUDA_VISIBLE_DEVICES=0 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 [...]

Niente da fare, anche con BS=1 otteniamo:

RuntimeError: Memoria CUDA esaurita. Tentativo di allocare 64,00 MiB (GPU 0; capacità totale di 23,70 GiB;
21,37 GiB già allocati; 45,69 MiB disponibili; 22,05 GiB riservati in totale da PyTorch)

Nota, come già detto in precedenza, mostro solo le parti importanti e gli argomenti completi della riga di comando possono essere trovati qui .

Ora aggiorna il tuo transformers alla versione 4.2.0 o superiore, quindi installa DeepSpeed:

pip install deepspeed

e riproviamo, questa volta aggiungendo DeepSpeed alla riga di comando:

export BS=20
CUDA_VISIBLE_DEVICES=0 deepspeed --num_gpus=1 ./finetune_trainer.py \
--model_name_or_path t5-3b --n_train 60 --n_val 10 \
--per_device_eval_batch_size $BS --per_device_train_batch_size $BS \
--task translation_en_to_ro --fp16 --deepspeed ds_config_1gpu.json [...]

et voilà! Otteniamo una dimensione di batch di 20 addestrata perfettamente. Probabilmente potrei spingerla ancora più avanti. Il programma è fallito con OOM a BS=30 .

Ecco i risultati rilevanti:

2021-01-12 19:06:31 | INFO | __main__ |   train_n_objs = 60
2021-01-12 19:06:31 | INFO | __main__ |   train_runtime = 8,8511
2021-01-12 19:06:35 | INFO | __main__ |   val_n_objs = 10
2021-01-12 19:06:35 | INFO | __main__ |   val_runtime = 3,5329

Non possiamo confrontare questi con la base, poiché la base non si avvia nemmeno e fallisce immediatamente con OOM.

Semplicemente incredibile!

Ho utilizzato solo un piccolo campione poiché ero principalmente interessato a poter addestrare e valutare con questo enorme modello che normalmente non si adatta a una GPU da 24 GB.

Se desideri sperimentare questo benchmark da solo o desideri conoscere maggiori dettagli sull’hardware e il software utilizzati per eseguirlo, ti preghiamo di fare riferimento a questo post .

Dato che transformers ha solo integrato queste soluzioni favolose e non ne è stato parte dell’invenzione, condividerò le risorse in cui puoi scoprire tutti i dettagli da solo. Ma ecco alcune informazioni rapide che potrebbero aiutarti a capire come ZeRO gestisce queste imprese straordinarie.

La caratteristica chiave di ZeRO è l’aggiunta di archiviazione dati distribuita al concetto abbastanza familiare di addestramento parallelo dei dati.

Il calcolo su ogni GPU è esattamente lo stesso dell’addestramento parallelo dei dati, ma i parametri, i gradienti e gli stati dell’ottimizzatore vengono archiviati in modo distribuito/partizionato su tutte le GPU e recuperati solo quando necessario.

Il seguente diagramma, tratto da questo articolo del blog, illustra come funziona:

L’approccio ingegnoso di ZeRO consiste nel partizionare i parametri, i gradienti e gli stati dell’ottimizzatore in modo equo su tutte le GPU e dare a ciascuna GPU una singola partizione (chiamata anche frammento). Ciò porta a nessuna sovrapposizione nell’archiviazione dei dati tra le GPU. Durante l’esecuzione ogni GPU costruisce i dati di ciascun layer al volo chiedendo alle GPU partecipanti di inviare le informazioni che mancano.

Quest’idea potrebbe essere difficile da comprendere e troverai il mio tentativo di spiegazione qui .

Alla data di scrittura, FairScale e DeepSpeed eseguono solo la partizione (sharding) degli stati dell’ottimizzatore e dei gradienti. Si prevede che la partizione dei parametri del modello verrà presto implementata in DeepSpeed e FairScale.

L’altra caratteristica potente è ZeRO-Offload ( documento ). Questa funzionalità sposta parte dell’elaborazione e delle esigenze di memoria sulla CPU dell’host, consentendo così di utilizzare più memoria sulla GPU. Hai visto il suo impatto drammatico nel successo dell’esecuzione di t5-3b su una GPU da 24GB.

Un altro problema di cui molte persone si lamentano nei forum di pytorch è la frammentazione della memoria della GPU. Spesso si riceve un errore OOM che potrebbe apparire così:

RuntimeError: CUDA out of memory. Tried to allocate 1.48 GiB (GPU 0; 23.65 GiB total capacity;
16.22 GiB already allocated; 111.12 MiB free; 22.52 GiB reserved in total by PyTorch)

Il programma cerca di allocare ~1.5GB e la GPU ha ancora circa 6-7GB di memoria non utilizzata, ma segnala di avere solo ~100MB di memoria libera contigua e fallisce con l’errore OOM. Questo accade quando vengono allocati e deallocati pezzi di dimensioni diverse e nel tempo vengono creati buchi che portano alla frammentazione della memoria, in cui c’è molta memoria non utilizzata ma nessun frammento contiguo della dimensione desiderata. Nell’esempio sopra il programma potrebbe probabilmente allocare 100MB di memoria contigua, ma chiaramente non può ottenerne 1.5GB in un unico frammento.

DeepSpeed affronta questo problema gestendo la memoria della GPU da solo e garantendo che le allocazioni di memoria a lungo termine non si mescolino con quelle a breve termine e quindi ci sia molto meno frammentazione. Sebbene l’articolo non fornisse dettagli, il codice sorgente è disponibile, quindi è possibile vedere come DeepSpeed riesce a fare ciò.

Come ZeRO sta per Zero Redundancy Optimizer, è facile vedere che è all’altezza del suo nome.

Oltre al supporto previsto per la partizione dei parametri del modello in DeepSpeed, sono già state rilasciate nuove funzionalità che non abbiamo ancora esplorato. Queste includono DeepSpeed Sparse Attention e 1-bit Adam, che dovrebbero ridurre l’utilizzo della memoria e ridurre drasticamente il sovraccarico di comunicazione tra le GPU, il che dovrebbe portare a un addestramento ancora più veloce e supportare modelli ancora più grandi.

Sono convinto che vedremo nuovi contributi anche dal team di FairScale. Credo che stiano lavorando anche alla fase 3 di ZeRO.

Ancora più eccitante, ZeRO viene integrato in pytorch .

Se hai trovato interessanti i risultati condivisi in questo articolo del blog, procedi qui per i dettagli su come utilizzare DeepSpeed e FairScale con il transformers Trainer.

Puoi, naturalmente, modificare il tuo trainer per integrare DeepSpeed e FairScale, in base alle istruzioni di ciascun progetto, oppure puoi “barare” e vedere come l’abbiamo fatto nel Trainer di transformers. Se scegli quest’ultima opzione, per orientarti, utilizza grep nel codice sorgente per cercare deepspeed e/o sharded_ddp .

La buona notizia è che ZeRO non richiede alcuna modifica al modello. Le uniche modifiche necessarie riguardano il codice di addestramento.

Se incontri problemi con la parte di integrazione di uno di questi progetti, apri una segnalazione in transformers .

Ma se hai problemi con l’installazione, la configurazione e la distribuzione di DeepSpeed e FairScale, devi chiedere ai esperti nei rispettivi ambiti, quindi utilizza invece DeepSpeed Issue o FairScale Issue.

Mentre non è necessario comprendere come funzionino questi progetti e puoi semplicemente distribuirli tramite il transformers Trainer, se desideri capire il motivo e il modo in cui funzionano, ti preghiamo di consultare le seguenti risorse.

• FairScale GitHub

• DeepSpeed GitHub

• Documento: ZeRO: Ottimizzazioni della memoria per la formazione di modelli a trilioni di parametri. Il documento è molto interessante, ma è molto conciso.

• Ecco una buona discussione video del documento con immagini.

• Documento: ZeRO-Offload: Democratizzazione della formazione di modelli su scala miliardaria. Appena pubblicato – questo va nei dettagli della funzionalità ZeRO Offload.

• Configurazione e tutorial di DeepSpeed

• Oltre al documento, consiglio vivamente di leggere i seguenti articoli di blog dettagliati con diagrammi:

  • DeepSpeed: Formazione di modelli su scala estrema per tutti
  • ZeRO & DeepSpeed: Nuove ottimizzazioni di sistema consentono la formazione di modelli con oltre 100 miliardi di parametri
  • Turing-NLG: Un modello di linguaggio da 17 miliardi di parametri di Microsoft

• Esempi DeepSpeed su GitHub

Siamo rimasti sorpresi dal sorprendente livello di supporto che abbiamo ricevuto dai team di sviluppo di FairScale e DeepSpeed durante il lavoro di integrazione di questi progetti in transformers.

In particolare, vorrei ringraziare:

• Benjamin Lefaudeux @blefaudeux
• Mandeep Baines @msbaines

del team FairScale e:

• Jeff Rasley @jeffra
• Olatunji Ruwase @tjruwase
• Samyam Rajbhandari @samyam

del team DeepSpeed per il vostro generoso e premuroso supporto e la rapida risoluzione dei problemi che abbiamo riscontrato.

E a HuggingFace per aver fornito l’accesso all’hardware su cui sono stati eseguiti i benchmark.

Sylvain Gugger @sgugger e Stas Bekman @stas00 hanno lavorato sull’integrazione di questi progetti.