Come 🤗 Accelerate esegue modelli molto grandi grazie a PyTorch

'🤗 Accelerate esegue modelli grandi con PyTorch'

Caricare ed eseguire modelli di grandi dimensioni

Meta AI e BigScience hanno recentemente reso open source modelli di linguaggio molto grandi che non possono essere inseriti nella memoria (RAM o GPU) della maggior parte dell’hardware per consumatori. Presso Hugging Face, parte della nostra missione è rendere accessibili anche questi modelli di grandi dimensioni, quindi abbiamo sviluppato strumenti che ti permettono di eseguire questi modelli anche se non possiedi un supercomputer. Tutti gli esempi selezionati in questo post del blog sono eseguiti su un’istanza Colab gratuita (con RAM e spazio su disco limitati), se hai accesso a uno spazio su disco maggiore, non esitare a selezionare checkpoint più grandi.

Ecco come possiamo eseguire OPT-6.7B:

import torch
from transformers import pipeline

# Questo funziona su un'istanza Colab base.
# Se hai tempo da aspettare e abbastanza spazio su disco, seleziona un checkpoint più grande!
checkpoint = "facebook/opt-6.7b"
generator = pipeline("text-generation", model=checkpoint, device_map="auto", torch_dtype=torch.float16)

# Esegui l'inferezza
generator("Sempre più modelli di linguaggio di grandi dimensioni vengono resi open source, quindi Hugging Face ha")

Spiegheremo cosa fanno ciascuno di questi argomenti tra un momento, ma prima considera semplicemente il normale processo di caricamento del modello in PyTorch: di solito consiste in:

1. Creare il modello
2. Caricare in memoria i suoi pesi (in un oggetto chiamato di solito state_dict)
3. Caricare quei pesi nel modello creato
4. Spostare il modello sul dispositivo per l’inferezza

Sebbene questo abbia funzionato abbastanza bene negli ultimi anni, i modelli molto grandi rendono questo approccio impegnativo. Qui il modello selezionato ha 6.7 miliardi di parametri. Nella precisione predefinita, ciò significa che solo il passaggio 1 (creazione del modello) richiederà circa 26,8 GB di RAM (un parametro in float32 occupa 4 byte in memoria). Questo non può nemmeno essere inserito nella RAM disponibile su Colab.

• Classificazione delle immagini con AutoTrain
• Modelli di linguaggio molto grandi e come valutarli
• Diffusione stabile giapponese

Quindi il passaggio 2 caricherà in memoria una seconda copia del modello (quindi altri 26,8 GB di RAM nella precisione predefinita). Se stessi cercando di caricare i modelli più grandi, ad esempio BLOOM o OPT-176B (che entrambi hanno 176 miliardi di parametri), in questo modo avresti bisogno di 1,4 terabyte di RAM della CPU. Questo è un po’ eccessivo! E tutto questo solo per spostare il modello su una (o più) GPU al passaggio 4.

Chiaramente abbiamo bisogno di qualcosa di più intelligente. In questo post del blog, spiegheremo come Accelerate sfrutta le funzionalità di PyTorch per caricare ed eseguire l’inferezza con modelli molto grandi, anche se non si adattano alla RAM o a una GPU. In poche parole, cambia il processo sopra come segue:

1. Creare un modello vuoto (ad esempio senza pesi)
2. Decidere dove andrà ogni livello (quando sono disponibili più dispositivi)
3. Caricare in memoria parti dei suoi pesi
4. Caricare quei pesi nel modello vuoto
5. Spostare i pesi sul dispositivo per l’inferezza
6. Ripetere dal passaggio 3 per i pesi successivi fino a quando tutti i pesi sono stati caricati

Creare un modello vuoto

PyTorch 1.9 ha introdotto un nuovo tipo di dispositivo chiamato dispositivo meta. Questo ci consente di creare tensori senza alcun dato associato: un tensore sul dispositivo meta ha solo bisogno di una forma. Finché ti trovi sul dispositivo meta, puoi quindi creare tensori di dimensioni arbitrariamente grandi senza preoccuparti della RAM della CPU (o della GPU).

Ad esempio, il seguente codice causerà un errore su Colab:

import torch

large_tensor = torch.randn(100000, 100000)

poiché questo grande tensore richiede 4 * 10**10 byte (la precisione predefinita è FP32, quindi ogni elemento del tensore occupa 4 byte), quindi 40 GB di RAM. Tuttavia, sul dispositivo meta funziona perfettamente:

import torch

large_tensor = torch.randn(100000, 100000, device="meta")

Se provi a visualizzare questo tensore, ecco cosa stamperà PyTorch:

tensor(..., device='meta', size=(100000, 100000))

Come abbiamo detto prima, non ci sono dati associati a questo tensore, solo una forma.

Puoi istanziare un modello direttamente sul dispositivo meta:

large_model = torch.nn.Linear(100000, 100000, device="meta")

Ma per un modello esistente, questa sintassi richiederebbe di riscrivere tutto il codice di modellazione in modo che ogni sottomodulo accetti e passi un argomento chiave device. Poiché ciò era impraticabile per i 150 modelli della libreria Transformers, abbiamo sviluppato un gestore di contesto che istanzierà un modello vuoto per te.

Ecco come puoi istanziare una versione vuota di BLOOM:

from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("bigscience/bloom")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

Questo funziona su qualsiasi modello, ma otterrai indietro un guscio che non puoi utilizzare direttamente: alcune operazioni sono implementate per il dispositivo meta, ma non tutte ancora. Qui ad esempio, puoi utilizzare il large_model definito sopra con un input, ma non il modello BLOOM. Anche quando lo si utilizza, l’output sarà un tensore del dispositivo meta, quindi otterrai la forma del risultato, ma nient’altro.

Come ulteriore lavoro su questo, il team di PyTorch sta lavorando su una nuova classe FakeTensor, che è un po’ come i tensori sul dispositivo meta, ma con le informazioni sul dispositivo (oltre alla forma e al dtype)

Dato che conosciamo la forma di ogni peso, possiamo tuttavia sapere quanto spazio di memoria consumeranno tutti una volta che carichiamo completamente i tensori preaddestrati. Pertanto, possiamo prendere una decisione su come suddividere il nostro modello tra CPU e GPU.

Calcolo di una mappa dei dispositivi

Prima di iniziare a caricare i pesi preaddestrati, dovremo sapere dove vogliamo posizionarli. In questo modo possiamo liberare la RAM della CPU ogni volta che abbiamo posizionato un peso nel suo posto giusto. Ciò può essere fatto con il modello vuoto sul dispositivo meta, poiché abbiamo solo bisogno di conoscere la forma di ciascun tensore e il suo dtype per calcolare quanto spazio occuperà in memoria.

Accelerate fornisce una funzione per determinare automaticamente una mappa dei dispositivi da un modello vuoto. Cercherà di massimizzare l’utilizzo di tutte le GPU disponibili, quindi della RAM della CPU, e infine segnalerà i pesi che non entrano, per il trasferimento su disco. Vediamo un esempio usando OPT-13b.

from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

device_map = infer_auto_device_map(model)

Questo restituirà un dizionario che mappa moduli o pesi a un dispositivo. Su una macchina con una Titan RTX, ad esempio, otteniamo quanto segue:

{'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.9': 0,
 'model.decoder.layers.10.self_attn': 0,
 'model.decoder.layers.10.activation_fn': 0,
 'model.decoder.layers.10.self_attn_layer_norm': 0,
 'model.decoder.layers.10.fc1': 'cpu',
 'model.decoder.layers.10.fc2': 'cpu',
 'model.decoder.layers.10.final_layer_norm': 'cpu',
 'model.decoder.layers.11': 'cpu',
 ...
 'model.decoder.layers.17': 'cpu',
 'model.decoder.layers.18.self_attn': 'cpu',
 'model.decoder.layers.18.activation_fn': 'cpu',
 'model.decoder.layers.18.self_attn_layer_norm': 'cpu',
 'model.decoder.layers.18.fc1': 'disk',
 'model.decoder.layers.18.fc2': 'disk',
 'model.decoder.layers.18.final_layer_norm': 'disk',
 'model.decoder.layers.19': 'disk',
 ...
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

Accelerate ha valutato che gli embedding e il decoder fino al 9° blocco potevano essere tutti inseriti nella GPU (dispositivo 0), quindi una parte del 10° blocco deve essere sulla CPU, così come i pesi successivi fino al 17° strato. Quindi il 18° strato è diviso tra la CPU e il disco e i successivi strati devono tutti essere trasferiti su disco

In realtà, utilizzare questa mappa del dispositivo in seguito non funzionerà, perché i livelli che compongono questo modello hanno connessioni residue (dove l’input del blocco viene aggiunto all’output del blocco), quindi tutti i livelli di un dato strato dovrebbero essere sullo stesso dispositivo. Possiamo indicare questo ad Accelerate passando un elenco di nomi di moduli che non devono essere suddivisi con l’argomento chiave no_split_module_classes:

device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"])

Questo restituirà:

'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.9': 0,
 'model.decoder.layers.10': 'cpu',
 'model.decoder.layers.11': 'cpu',
 ...
 'model.decoder.layers.17': 'cpu',
 'model.decoder.layers.18': 'disk',
 ...
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

Ora, ogni livello è sempre sullo stesso dispositivo.

In Transformers, quando si utilizza device_map nel metodo from_pretrained() o in una pipeline, le classi di blocchi da lasciare sullo stesso dispositivo sono fornite automaticamente, quindi non è necessario preoccuparsene. Si noti che hai le seguenti opzioni per device_map (rilevanti solo quando si dispone di più di una GPU):

• "auto" o "balanced": Accelerate dividerà i pesi in modo che ogni GPU venga utilizzata in modo uguale;
• "balanced_low_0": Accelerate dividerà i pesi in modo che ogni GPU venga utilizzata in modo uguale, ad eccezione della prima, dove cercherà di avere il minor numero di pesi possibile (utile quando si desidera lavorare con gli output del modello su una GPU, ad esempio quando si utilizza la funzione generate);
• "sequential": Accelerate riempirà le GPU in ordine (quindi le ultime potrebbero non essere utilizzate affatto).

È anche possibile passare il proprio device_map purché segua il formato che abbiamo visto prima (dizionario dei nomi dei layer/moduli al dispositivo).

Infine, si noti che i risultati del device_map che si ricevono dipendono dal tipo di dati selezionato (poiché diversi tipi di float occupano una quantità diversa di spazio). Fornendo dtype="float16" otterremo risultati diversi:

device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")

In questa precisione, possiamo adattare il modello fino al livello 21 sulla GPU:

{'model.decoder.embed_tokens': 0,
 'model.decoder.embed_positions': 0,
 'model.decoder.final_layer_norm': 0,
 'model.decoder.layers.0': 0,
 'model.decoder.layers.1': 0,
 ...
 'model.decoder.layers.21': 0,
 'model.decoder.layers.22': 'cpu',
 ...
 'model.decoder.layers.37': 'cpu',
 'model.decoder.layers.38': 'disk',
 'model.decoder.layers.39': 'disk',
 'lm_head': 'disk'}

Ora che sappiamo dove ogni peso deve andare, possiamo caricare progressivamente i pesi preaddestrati all’interno del modello.

Sharding state dicts

Tradizionalmente, i modelli PyTorch vengono salvati in un unico file contenente una mappa dal nome del parametro al peso. Questa mappa viene spesso chiamata state_dict. Ecco un estratto della documentazione di PyTorch sul salvataggio e il caricamento:

# Salva i pesi del modello
torch.save(my_model.state_dict(), 'model_weights.pth')

# Ricaricali
new_model = ModelClass()
new_model.load_state_dict(torch.load('model_weights.pth'))

Questo funziona molto bene per modelli con meno di 1 miliardo di parametri, ma per modelli più grandi, questo richiede molta memoria RAM. Il modello BLOOM ha 176 miliardi di parametri; anche con i pesi salvati in bfloat16 per risparmiare spazio, rappresenta comunque 352 GB nel complesso. Mentre il supercomputer che ha addestrato questo modello potrebbe avere questa quantità di memoria disponibile, richiederla per l’uso in inferenza è irrealistico.

Questo è il motivo per cui i modelli di grandi dimensioni nell’Hugging Face Hub non vengono salvati e condivisi con un unico file che contiene tutti i pesi, ma diversi di essi. Se si va alla pagina del modello BLOOM, ad esempio, si vedranno 72 file chiamati pytorch_model_xxxxx-of-00072.bin, ognuno dei quali contiene una parte dei pesi del modello. Utilizzando questo formato, possiamo caricare una parte del dizionario di stato in memoria, inserire i pesi all’interno del modello, spostarli sul dispositivo corretto, quindi eliminare questa parte del dizionario di stato prima di passare alla successiva. Invece di richiedere una quantità di RAM sufficiente per ospitare l’intero modello, abbiamo bisogno solo di una quantità di RAM sufficiente per ottenere la parte più grande del checkpoint, che chiamiamo shard, quindi 7,19 GB nel caso di BLOOM.

Chiamiamo i checkpoint salvati in diversi file come checkpoint frammentati di BLOOM, e abbiamo standardizzato il loro formato come segue:

• Un file (chiamato pytorch_model.bin.index.json) contiene alcuni metadati e una mappa dei nomi dei parametri ai nomi dei file, indicando dove trovare ogni peso
• Tutti gli altri file sono dizionari di stato standard di PyTorch, contengono solo una parte del modello invece dell’intero. Puoi dare un’occhiata al contenuto del file di indice qui .

Per caricare un checkpoint frammentato in un modello, è sufficiente iterare sui vari frammenti. Accelerate fornisce una funzione chiamata load_checkpoint_in_model che farà questo per te se hai clonato uno dei repository del Hub, oppure puoi utilizzare direttamente il metodo from_pretrained di Transformers, che gestirà il download e la memorizzazione nella cache per te:

import torch
from transformers import AutoModelForCausalLM

# Darà un errore
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.float16)

Se la mappa del dispositivo calcolata automaticamente richiede che alcuni pesi siano memorizzati su disco perché non hai abbastanza RAM GPU e CPU, otterrai un errore che indica che è necessario fornire una cartella in cui verranno memorizzati i pesi che devono essere spostati su disco:

ValueError: La `device_map` corrente ha i pesi spostati su disco. Fornisci una cartella `offload_folder` per loro.

Aggiungere questo argomento risolverà l’errore:

import torch
from transformers import AutoModelForCausalLM

# Si esaurirà la RAM su Colab
checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", offload_folder="offload", torch_dtype=torch.float16
)

Si noti che se si sta cercando di caricare un modello molto grande che richiede un trasferimento su disco oltre al trasferimento su CPU, potrebbe esaurirsi la RAM quando vengono caricati gli ultimi frammenti del checkpoint, poiché la parte del modello rimane su CPU occupando spazio. In tal caso, utilizzare l’opzione offload_state_dict=True per trasferire temporaneamente la parte del modello che rimane su CPU mentre tutti i pesi vengono caricati e ricaricarla in RAM una volta che tutti i pesi sono stati elaborati

import torch
from transformers import AutoModelForCausalLM

checkpoint = "facebook/opt-13b"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map="auto", offload_folder="offload", offload_state_dict=True, torch_dtype=torch.float16
)

Questo funzionerà su Colab, ma sarà così vicino all’utilizzo di tutta la RAM disponibile che esaurirà la RAM quando si prova a generare una previsione. Per ottenere un modello utilizzabile, è necessario spostare un altro livello su disco. È possibile farlo prendendo la `device_map` calcolata nella sezione precedente, adattandola leggermente, quindi passandola alla chiamata `from_pretrained`:

import torch
from transformers import AutoModelForCausalLM

checkpoint = "facebook/opt-13b"
device_map["model.decoder.layers.37"] = "disk"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, device_map=device_map, offload_folder="offload", offload_state_dict=True, torch_dtype=torch.float16
)

Esecuzione di un modello suddiviso su più dispositivi

Un’ultima parte che non abbiamo toccato è come Accelerate consente al tuo modello di funzionare con i suoi pesi distribuiti su diversi GPU, RAM CPU e cartelle su disco. Questo viene fatto molto semplicemente utilizzando hooks.

hooks sono un’API PyTorch che aggiunge funzioni eseguite appena prima di ogni chiamata `forward`

Non abbiamo potuto utilizzare questo direttamente poiché supportano solo modelli con argomenti regolari e nessun argomento con nome nel passaggio `forward`, ma abbiamo preso la stessa idea. Una volta che il modello è caricato, la funzione `dispatch_model` aggiungerà hooks a tutti i moduli e sottomoduli che vengono eseguiti prima e dopo ogni passaggio `forward`. Questi hooks:

• garantiscono che tutti gli input del modulo siano sullo stesso dispositivo dei pesi;
• se i pesi sono stati spostati sulla CPU, li spostano sulla GPU 0 prima del passaggio `forward` e li riportano sulla CPU subito dopo;
• se i pesi sono stati spostati su disco, li caricano in RAM e sulla GPU 0 prima del passaggio `forward` e liberano questa memoria subito dopo.

L’intero processo è riassunto nel seguente video:

In questo modo, il tuo modello può essere caricato ed eseguito anche se non hai abbastanza RAM GPU e RAM CPU. L’unica cosa di cui hai bisogno è spazio su disco (e molta pazienza!) Sebbene questa soluzione sia piuttosto ingenua se si dispone di più GPU (non c’è coinvolto alcun parallelismo intelligente della pipeline, semplicemente si utilizzano le GPU in sequenza), produce comunque risultati abbastanza decenti per BLOOM. E ti consente di eseguire il modello su configurazioni più piccole (sebbene più lentamente).

Per saperne di più sull’accelerazione dell’inferenza di modelli di grandi dimensioni, consulta la documentazione.