Incredibilmente veloce inferenza BLOOM con DeepSpeed e Accelerate

'Incredibly fast BLOOM inference with DeepSpeed and Accelerate.'

Questo articolo mostra come ottenere un throughput incredibilmente veloce per token durante la generazione con il modello BLOOM a 176 miliardi di parametri.

Dato che il modello richiede 352 GB di pesi in formato bf16 (bfloat16) (176*2), la configurazione più efficiente è con 8 GPU A100 da 80 GB ciascuna. È anche possibile utilizzare 2 GPU A100 da 8×40 GB o 2 GPU A6000 da 8×48 GB. Il motivo principale per utilizzare queste GPU è che, al momento della stesura di questo articolo, forniscono la memoria GPU più ampia, ma è possibile utilizzare anche altre GPU. Ad esempio, si possono utilizzare 24 V100 da 32 GB.

L’utilizzo di un singolo nodo di solito garantisce un throughput più veloce, poiché nella maggior parte dei casi l’hardware di collegamento GPU intra-nodo è più veloce di quello inter-nodo, ma non è sempre così.

Se non si dispone di così tanto hardware, è comunque possibile eseguire l’inferenza di BLOOM su GPU più piccole, utilizzando il supporto della CPU o NVMe, ma ovviamente il tempo di generazione sarà molto più lento.

• Newsletter Etica e Società #1
• SetFit Apprendimento Efficienti a Pochi Esempi Senza Promemoria
• Come 🤗 Accelerate esegue modelli molto grandi grazie a PyTorch

Discuteremo anche le soluzioni quantizzate a 8 bit, che richiedono la metà della memoria GPU a costo di un throughput leggermente più lento. Esamineremo le librerie BitsAndBytes e Deepspeed-Inference in quel contesto.

Benchmarks

Senza ulteriori indugi, mostriamo alcuni numeri.

Per garantire la coerenza, a meno che non venga specificato diversamente, i benchmark in questo articolo sono stati eseguiti tutti sullo stesso nodo A100 da 8×80 GB con 512 GB di memoria CPU su Jean Zay HPC. Gli utenti di JeanZay HPC godono di un’IO molto rapida con una velocità di lettura di circa 3 GB/s (GPFS). Questo è importante per il tempo di caricamento del checkpoint. Un disco lento comporterà un tempo di caricamento lento. Soprattutto perché stiamo facendo IO contemporaneamente in più processi.

Tutti i benchmark generano in modo avido 100 token di output:

Genera argomenti {'max_length': 100, 'do_sample': False}

Il prompt di input è composto solo da alcuni token. È anche abilitata la memorizzazione nella cache del token precedente, poiché sarebbe molto lento ricalcolarli ogni volta.

Iniziamo con un’occhiata veloce a quanto tempo è stato necessario per prepararsi alla generazione, ovvero quanto tempo è stato necessario per caricare e preparare il modello:

Deepspeed-Inference viene fornito con repository di pesi pre-shardizzati e il caricamento richiede circa 1 minuto. Anche il tempo di caricamento di Accelerate è eccellente, appena circa 2 minuti. Le altre soluzioni sono molto più lente in questa fase.

Il tempo di caricamento può essere o meno importante, poiché una volta caricati è possibile generare continuamente i token più e più volte senza un ulteriore overhead di caricamento.

Successivamente, il benchmark più importante è il throughput della generazione di token. La metrica del throughput qui è semplice: quanto tempo è stato necessario per generare 100 nuovi token diviso per 100 e per la dimensione del batch (ovvero diviso per il numero totale di token generati).

Ecco il throughput in millisecondi su 8 GPU da 80 GB:

dove OOM == condizione Out of Memory, in cui la dimensione del batch era troppo grande per entrare nella memoria GPU.

Un throughput inferiore a 1 millisecondo con il Parallelismo Tensoriale (TP) di Deepspeed-Inference e i kernel CUDA personalizzati fusi! È assolutamente incredibile! Tuttavia, utilizzare questa soluzione per altri modelli su cui non è stata testata potrebbe richiedere del tempo agli sviluppatori per farla funzionare.

Anche Accelerate è molto veloce. Utilizza un approccio molto semplice di Parallelismo in Pipeline (PP) e poiché è molto semplice, dovrebbe funzionare con qualsiasi modello senza problemi.

Poiché Deepspeed-ZeRO può elaborare più flussi di generazione in parallelo, il suo throughput può essere ulteriormente diviso per 8 o 16, a seconda che si utilizzino 8 o 16 GPU durante la chiamata generate. E, naturalmente, ciò significa che può elaborare una dimensione del batch di 64 nel caso di 8×80 A100 (la tabella sopra) e quindi il throughput è di circa 4 millisecondi – quindi tutte e 3 le soluzioni sono molto vicine tra loro.

Rivediamo ancora una volta come sono stati calcolati questi numeri. Per generare 100 nuovi token per una dimensione del batch di 128, sono stati necessari 8832 millisecondi in tempo reale quando si utilizza Deepspeed-Inference in modalità fp16. Quindi ora per calcolare il throughput abbiamo fatto: tempo di esecuzione/(dimensione del batch*nuovi token) o 8832/(128*100) = 0.69.

Ora vediamo la potenza dei modelli basati su int8 quantizzati forniti da Deepspeed-Inference e BitsAndBytes, in quanto richiedono solo la metà della memoria GPU originale dell’inferenza in bfloat16 o float16.

Throughput in msecs 4x80GB A100:

Per riprodurre i risultati del benchmark, aggiungi semplicemente --benchmark a uno di questi 3 script di cui parleremo di seguito.

Soluzioni

Prima controlla il repository demo:

git clone https://github.com/huggingface/transformers-bloom-inference
cd transformers-bloom-inference

In questo articolo utilizzeremo 3 script situati in bloom-inference-scripts/.

Le soluzioni specifiche per il framework sono presentate in ordine alfabetico:

HuggingFace Accelerate

Accelerate

Accelerate gestisce i modelli di grandi dimensioni per l’inferenza nel seguente modo:

1. Istanzia il modello con pesi vuoti.
2. Analizza la dimensione di ogni layer e lo spazio disponibile su ogni dispositivo (GPU, CPU) per decidere dove posizionare ciascun layer.
3. Carica il checkpoint del modello bit per bit e colloca ciascun peso sul suo dispositivo

Successivamente assicura che il modello funzioni correttamente con hook che trasferiscono gli input e gli output sul dispositivo corretto e che i pesi del modello scaricati sulla CPU (o persino sul disco) vengano caricati su una GPU proprio prima del passaggio in avanti, prima di essere scaricati nuovamente una volta terminato il passaggio in avanti.

In una situazione in cui ci sono più GPU con spazio sufficiente per ospitare l’intero modello, il controllo passa da una GPU alla successiva finché tutti i layer non sono stati eseguiti. Solo una GPU funziona in un dato momento, il che può sembrare inefficiente, ma produce comunque un throughput decente nonostante l’inattività delle GPU.

È anche molto flessibile poiché lo stesso codice può essere eseguito su qualsiasi configurazione data. Accelerate utilizzerà prima tutte le GPU disponibili, quindi scaricherà sulla CPU fino a quando la RAM non sarà piena e infine sul disco. Lo scaricamento sulla CPU o sul disco renderà le cose più lente. Ad esempio, gli utenti hanno segnalato di eseguire BLOOM senza modifiche del codice solo su 2 A100 con un throughput di 15s per token rispetto a 10 msecs su 8×80 A100.

Puoi saperne di più su questa soluzione nella documentazione di Accelerate.

Preparazione

pip install transformers>=4.21.3 accelerate>=0.12.0

Esecuzione

L’esecuzione semplice è:

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --batch_size 1 --benchmark

Per attivare la soluzione con quantizzazione a 8 bit di BitsAndBytes, installa prima bitsandbytes:

pip install bitsandbytes

e quindi aggiungi --dtype int8 alla riga di comando precedente:

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

se hai più di 4 GPU, puoi indicargli di utilizzare solo 4 con:

CUDA_VISIBLE_DEVICES=0,1,2,3 python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

La dimensione del batch più alta che siamo riusciti ad eseguire senza OOM è stata 40 in questo caso. Se guardi all’interno dello script, abbiamo dovuto modificare la mappa di allocazione della memoria per liberare la prima GPU in modo che gestisca solo attivazioni e la cache dei token precedenti.

DeepSpeed-Inference

DeepSpeed-Inference utilizza la parallelizzazione dei tensori e i kernel CUDA fusi efficienti per fornire un’infereza super veloce di meno di 1msec per token su un batch di dimensione 128.

Preparazione

pip install deepspeed>=0.7.3

Esecuzione

1. L’approccio più veloce consiste nell’utilizzare un checkpoint pre-sharded con TP (Tensor Parallel) che richiede solo ~1 minuto per il caricamento, rispetto a 10 minuti per il checkpoint bloom non pre-sharded:

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-fp16

1a. se si desidera eseguire il checkpoint originale di bloom, che una volta caricato funzionerà con la stessa velocità di esecuzione della soluzione precedente, ma il caricamento richiederà 10-20 minuti:

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name bigscience/bloom

2a. La versione quantizzata a 8 bit richiede solo la metà della memoria GPU rispetto alla versione di precisione dimezzata normale:

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

Qui abbiamo utilizzato microsoft/bloom-deepspeed-inference-int8 e abbiamo anche indicato allo script di eseguirsi in int8.

E ovviamente, solo 4x80GB di GPU A100 sono ora sufficienti:

deepspeed --num_gpus 4 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

La dimensione del batch più alta che siamo riusciti a eseguire senza OOM è stata 128 in questo caso.

Puoi vedere due fattori che contribuiscono a una migliore performance qui.

1. La velocità di esecuzione è stata migliorata utilizzando il parallelismo dei tensori (Tensor Parallelism, TP) invece del parallelismo dei pipeline (Pipeline Parallelism, PP) di Accelerate. Poiché Accelerate è progettato per essere molto generico, è anche sfortunatamente difficile massimizzare l’utilizzo della GPU. Tutti i calcoli vengono eseguiti prima sulla GPU 0, quindi sulla GPU 1, ecc. fino alla GPU 8, il che significa che 7 GPU sono inattive tutto il tempo. DeepSpeed-Inference, d’altra parte, utilizza TP, il che significa che invierà i tensori a tutte le GPU, calcolerà parte della generazione su ciascuna GPU e poi tutte le GPU si scambieranno i risultati tra loro, per poi passare al livello successivo. Ciò significa che tutte le GPU sono attive contemporaneamente, ma devono comunicare molto di più.

2. DeepSpeed-Inference utilizza anche kernel CUDA personalizzati per evitare di allocare troppa memoria e copiare i tensori da e verso le GPU. L’effetto di ciò è una minore richiesta di memoria e meno avvii del kernel, il che migliora la velocità di esecuzione e consente batch size più grandi, aumentando così la velocità di esecuzione complessiva.

Se sei interessato ad altri esempi, puoi dare un’occhiata all’esecuzione di GPT-J con Accelerate-Inference su GPU o all’esecuzione di BERT con Accelerate-Inference su GPU.

Deepspeed ZeRO-Inference

Deepspeed ZeRO utilizza un approccio di suddivisione magica che può prendere quasi qualsiasi modello e scalare su poche o centinaia di GPU per addestramento o inferenza.

Configurazione

pip install deepspeed

Esegui

Si noti che lo script esegue attualmente gli stessi input su tutte le GPU, ma è possibile eseguire un flusso diverso su ciascuna GPU e ottenere un throughput n_gpu volte più veloce. Ciò non è possibile con Deepspeed-Inference.

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 1 --benchmark

Ricorda che con ZeRO l’utente può generare più flussi unici contemporaneamente; pertanto, le prestazioni complessive dovrebbero essere la velocità di esecuzione in secondi per token divisa per il numero di GPU partecipanti, quindi da 8x a 16x più veloce a seconda che siano state utilizzate 8 o 16 GPU!

Puoi anche provare le soluzioni di offloading con una sola GPU di piccole dimensioni, che richiederanno molto tempo per eseguire, ma se non hai 8 GPU enormi, questa è la soluzione migliore possibile.

Offload CPU (1x GPU):

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark

Offload NVMe (1x GPU):

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --nvme_offload_path=/path/to/nvme_offload --benchmark

assicurati di adattare /percorso/a/nvme_offload a un’area in cui hai circa 400 GB di memoria libera su un’unità NVMe veloce.

Soluzioni client e server aggiuntive

Su transformers-bloom-inference troverai altre soluzioni molto efficienti, inclusi server solutions.

Ecco alcuni anteprime.

Soluzioni server:

• Mayank Mishra ha preso tutti gli script di demo discussi in questo post del blog e li ha trasformati in un pacchetto per un webserver, che puoi scaricare da qui

• Nicolas Patry ha sviluppato una soluzione per un webserver super efficiente basata su Rust.

Altre soluzioni lato client:

• Thomas Wang sta sviluppando un kernel CUDA personalizzato molto veloce per il modello BLOOM.

• Il team JAX @HuggingFace ha sviluppato una soluzione basata su JAX

Poiché questo post del blog potrebbe diventare obsoleto se lo leggi mesi dopo la sua pubblicazione, ti preghiamo di utilizzare transformers-bloom-inference per trovare le soluzioni più aggiornate.

Crediti del blog

Un enorme ringraziamento alle persone gentili seguenti che hanno posto domande pertinenti e contribuito a migliorare la leggibilità dell’articolo: Olatunji Ruwase e Philipp Schmid.