Ottimizzazione della Memoria per l’Inferenza e la Fine-Tuning dei Modelli Linguistici di Grande Scala

I modelli linguistici di grande scala (LLM) come GPT-4, Bloom e LLaMA hanno raggiunto capacità notevoli aumentando la scala a miliardi di parametri. Tuttavia, il deploy di questi modelli massicci per l’inferenza o la fine-tuning è impegnativo a causa delle loro immense esigenze di memoria. In questo blog tecnico, esploreremo le tecniche per stimare e ottimizzare il consumo di memoria durante l’inferenza e la fine-tuning degli LLM su vari setup hardware.

Comprensione delle Esigenze di Memoria

La memoria richiesta per caricare un LLM è determinata principalmente dal numero di parametri e dalla precisione numerica utilizzata per memorizzare i parametri. Una regola empirica è:

• Caricare un modello con X miliardi di parametri richiede circa 4X GB di VRAM in precisione float a 32 bit
• Caricare un modello con X miliardi di parametri richiede circa 2X GB di VRAM in precisione bfloat16/float16

Ad esempio, caricare il modello GPT-3 con 175 miliardi di parametri richiederebbe circa 350 GB di VRAM in precisione bfloat16. Al momento, le più grandi GPU commerciali come l’NVIDIA A100 e H100 offrono solo 80 GB di VRAM, rendendo necessarie tecniche di parallelismo dei tensori e dei modelli.

Durante l’inferenza, l’impronta di memoria è dominata dai parametri del modello e dai tensori di attivazione temporanei prodotti. Una stima di alto livello per l’utilizzo di memoria di picco durante l’inferenza è la somma della memoria richiesta per caricare i parametri del modello e la memoria per le attivazioni.

Quantificazione della Memoria di Inferenza

Quantifichiamo le esigenze di memoria per l’inferenza utilizzando il modello OctoCode, che ha circa 15 miliardi di parametri in formato bfloat16 (~ 31 GB). Utilizzeremo la libreria Transformers per caricare il modello e generare testo:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>model = AutoModelForCausalLM.from_pretrained(&quot;bigcode/octocoder&quot;,
torch_dtype=torch.bfloat16,
device_map=&quot;auto&quot;,
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained(&quot;bigcode/octocoder&quot;)
pipe = pipeline(&quot;text-generation&quot;, model=model, tokenizer=tokenizer)</p>

<p>prompt = &quot;Domanda: Scrivi una funzione Python per convertire byte in gigabyte.\n\nRisposta:&quot;
result = pipe(prompt, max_new_tokens=60)[0][&quot;generated_text&quot;][len(prompt):]</p>

<p>def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024</p>

<p>bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Output:

L’utilizzo di memoria di picco è di circa 29 GB, che si allinea con la nostra stima di 31 GB per caricare i parametri del modello in formato bfloat16.

Ottimizzazione della Memoria di Inferenza con Quantizzazione

Mentre la precisione bfloat16 è comune per l’addestramento degli LLM, i ricercatori hanno scoperto che quantizzare i pesi del modello a precisioni più basse come interi a 8 bit (int8) o 4 bit può ridurre notevolmente l’utilizzo di memoria con una perdita di accuratezza minima per le attività di inferenza come la generazione di testo.

Vediamo i risparmi di memoria dalla quantizzazione a 8 bit e 4 bit del modello OctoCode:

</div>
# Quantizzazione a 8 bit
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_8bit=True,
pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())</pre>

Output:

# Quantizzazione a 4 bit
model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder", load_in_4bit=True,
low_cpu_mem_usage=True, pad_token_id=0)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]
bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Output:

Con la quantizzazione a 8 bit, l’esigenza di memoria scende da 31 GB a 15 GB, mentre la quantizzazione a 4 bit la riduce ulteriormente a solo 9,5 GB! Ciò consente di eseguire il modello OctoCode da 15 miliardi di parametri su GPU consumer come l’RTX 3090 (24 GB di VRAM).

Tuttavia, si noti che una quantizzazione più aggressiva come la quantizzazione a 4 bit può talvolta portare a una degradazione dell’accuratezza rispetto alla precisione a 8 bit o bfloat16. C’è un compromesso tra risparmio di memoria e accuratezza che gli utenti devono valutare per il loro caso d’uso.

La quantizzazione è una tecnica potente che può consentire il deploy degli LLM in ambienti con risorse limitate come istanze cloud, dispositivi edge o addirittura telefoni cellulari riducendo drasticamente l’impronta di memoria.

Stima della Memoria per la Fine-Tuning

Mentre la quantizzazione è utilizzata principalmente per l’inferenza efficiente, tecniche come il parallelismo dei tensori e il parallelismo dei modelli sono cruciali per la gestione delle esigenze di memoria durante l’addestramento o la fine-tuning degli LLM.

L’utilizzo di memoria di picco durante la fine-tuning è solitamente 3-4 volte superiore all’inferenza a causa di esigenze di memoria aggiuntive per:

• Gradienti
• Stati dell’ottimizzatore
• Attivazioni del passo forward memorizzate per la backpropagation

Una stima conservativa è che la fine-tuning di un LLM con X miliardi di parametri richieda circa 4 * (2X) = 8X GB di VRAM in precisione bfloat16.

Ad esempio, la fine-tuning del modello LLaMA da 7 miliardi di parametri richiederebbe circa 7 * 8 = 56 GB di VRAM per GPU in precisione bfloat16. Ciò supera la capacità di memoria delle attuali GPU, rendendo necessarie tecniche di fine-tuning distribuite.

Tecniche di Fine-Tuning Distribuite

Sono state proposte diverse tecniche di fine-tuning distribuite per superare i vincoli di memoria della GPU per modelli di grandi dimensioni:

1. Parallelismo dei Dati: l’approccio classico del parallelismo dei dati replica l’intero modello su più GPU mentre divide e distribuisce i batch di dati di addestramento. Ciò riduce il tempo di addestramento linearmente con il numero di GPU, ma non riduce l’esigenza di memoria di picco su ogni GPU.
2. ZeRO Stage 3: una forma avanzata di parallelismo dei dati che partiziona i parametri del modello, i gradienti e gli stati dell’ottimizzatore su più GPU. Ciò riduce la memoria rispetto al parallelismo dei dati classico mantenendo solo i dati partizionati necessari su ogni GPU durante le diverse fasi dell’addestramento.
3. Parallelismo dei Tensori: invece di replicare il modello, il parallelismo dei tensori divide i parametri del modello in righe o colonne e li distribuisce su più GPU. Ogni GPU opera su un set di parametri, gradienti e stati dell’ottimizzatore partizionati, portando a risparmi di memoria sostanziali.
4. Parallelismo a Pipeline: questa tecnica partiziona i layer del modello su più GPU/lavoratori, con ogni dispositivo che esegue un subset dei layer. Le attivazioni vengono passate tra i lavoratori, riducendo la memoria di picco ma aumentando l’overhead di comunicazione.

Stimare l’utilizzo di memoria per questi metodi distribuiti non è banale, poiché la distribuzione dei parametri, dei gradienti, delle attivazioni e degli stati dell’ottimizzatore varia tra le tecniche. Inoltre, componenti diverse come il corpo del trasformatore e la testa di modellazione del linguaggio possono esibire comportamenti di allocazione della memoria diversi.

La Soluzione LLMem

I ricercatori hanno recentemente proposto LLMem, una soluzione che stima con precisione l’utilizzo di memoria della GPU quando si applicano metodi di fine-tuning distribuiti agli LLM su più GPU.

Stima dell’Utilizzo di Memoria della GPU per la Fine-Tuning di LLM Pre-Addestrati

LLMem considera fattori come la ricombinazione dei parametri prima del calcolo (ZeRO Stage 3), la raccolta dei dati di output nel passo backward (parallelismo dei tensori) e le diverse strategie di allocazione della memoria per il corpo del trasformatore e la testa di modellazione del linguaggio.

I risultati sperimentali mostrano che LLMem può stimare l’utilizzo di memoria di picco della GPU per la fine-tuning di LLM su una singola GPU con tassi di errore fino al 1,6%, superando il tasso di errore medio del 42,6% di DNNMem. Quando si applicano metodi di fine-tuning distribuiti agli LLM con oltre un miliardo di parametri su più GPU, LLMem raggiunge un impressionante tasso di errore medio del 3,0%.