Ottimizzazione della memoria per l'inferenza e la messa a punto di modelli linguistici di grandi dimensioni

I modelli linguistici di grandi dimensioni (LLM) come GPT-4, Bloom e LLaMA hanno raggiunto notevoli capacità scalando fino a miliardi di parametri. Tuttavia, l'implementazione di questi enormi modelli per l'inferenza o la messa a punto è impegnativa a causa dei loro immensi requisiti di memoria. In questo blog tecnico esploreremo le tecniche per stimare e ottimizzare il consumo di memoria durante l'inferenza LLM e la messa a punto tra varie configurazioni hardware.

Comprensione dei requisiti di memoria

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

• Il caricamento di un modello con X miliardi di parametri richiede circa 4 GB di VRAM in 32-bit precisione del galleggiante
• Il caricamento di un modello con X miliardi di parametri richiede circa 2 GB di VRAM in 16-bit precisione bfloat16/float16

Ad esempio, il caricamento del modello GPT-175 con parametro 3B richiederebbe circa 350 GB di VRAM con precisione bfloat16. Ad oggi, le più grandi GPU disponibili in commercio come NVIDIA A100 e H100 offrono solo 80 GB di VRAM, richiedendo tecniche di parallelismo del tensore e di parallelismo del modello.

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

Quantificazione della memoria di inferenza

Quantifichiamo i requisiti di memoria per l'inferenza utilizzando il modello OctoCode, che ha circa 15 miliardi di parametri in formato bfloat16 (~ 31 GB). Useremo il Libreria Transformers per caricare il modello e generare testo:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

model = AutoModelForCausalLM.from_pretrained("bigcode/octocoder",
torch_dtype=torch.bfloat16,
device_map="auto",
pad_token_id=0)
tokenizer = AutoTokenizer.from_pretrained("bigcode/octocoder")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = "Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]

def bytes_to_gigabytes(bytes):
return bytes / 1024 / 1024 / 1024

bytes_to_gigabytes(torch.cuda.max_memory_allocated())

Produzione:

L'utilizzo massimo della memoria della GPU è di circa 29 GB, in linea con la nostra stima di 31 GB per il caricamento dei parametri del modello in formato bfloat16.

Ottimizzazione della memoria di inferenza con la quantizzazione

Sebbene bfloat16 sia la precisione comune utilizzata per l'addestramento degli LLM, i ricercatori hanno scoperto che quantizzare i pesi del modello su tipi di dati a precisione inferiore come interi a 8 bit (int8) o interi a 4 bit può ridurre significativamente l'utilizzo della memoria con una perdita minima di precisione per attività di inferenza come generazione del testo.

Vediamo il risparmio di memoria derivante dalla quantizzazione a 8 e 4 bit del modello OctoCode:

</div>
# 8-bit quantization
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>

Produzione:

# 4-bit quantization
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())

Produzione:

Con la quantizzazione a 8 bit, il requisito di memoria scende da 31 GB a 15 GB, mentre la quantizzazione a 4 bit lo riduce ulteriormente a soli 9.5 GB! Ciò consente di eseguire il modello OctoCode con parametri 15B su GPU consumer come RTX 3090 (VRAM da 24 GB).

Tuttavia, è importante notare che una quantizzazione più aggressiva, come quella a 4 bit, può talvolta comportare un degrado della precisione rispetto alla precisione a 8 bit o bfloat16. Esiste un compromesso tra risparmio di memoria e precisione che gli utenti dovrebbero valutare in base al loro caso d'uso.

La quantizzazione è una tecnica potente che può consentire l'implementazione LLM in ambienti con risorse limitate come istanze cloud, dispositivi edge o persino telefoni cellulari riducendo drasticamente l'ingombro della memoria.

Stima della memoria per la messa a punto

Mentre la quantizzazione viene utilizzata principalmente per un'inferenza efficiente, tecniche come il parallelismo del tensore e il parallelismo del modello sono cruciali per la gestione dei requisiti di memoria durante l'addestramento o ritocchi di grandi modelli linguistici.

Il consumo massimo di memoria durante la regolazione fine è in genere 3-4 volte superiore all'inferenza a causa dei requisiti di memoria aggiuntivi per:

• Sfumature
• Stati dell'ottimizzatore
• Attivazioni dal passaggio in avanti archiviate per la propagazione all'indietro

Una stima prudente è che la messa a punto di un LLM con X miliardi di parametri richiede circa 4 * (2X) = 8 GB di VRAM con precisione bfloat16.

Ad esempio, la messa a punto del modello LLaMA con parametri 7B richiederebbe circa 7 * 8 = 56 GB di VRAM per GPU con precisione bfloat16. Ciò supera la capacità di memoria delle attuali GPU, rendendo necessarie tecniche di messa a punto distribuite.

Tecniche di fine tuning distribuito

Sono stati proposti diversi metodi di regolazione fine distribuita 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 suddividendo e distribuendo i batch di dati di training. Ciò riduce il tempo di addestramento in modo lineare con il numero di GPU, ma non riduce i requisiti di memoria di picco su ciascuna GPU.
2. ZeRO Fase 3: una forma avanzata di parallelismo dei dati che suddivide i parametri del modello, i gradienti e gli stati dell'ottimizzatore tra le GPU. Riduce la memoria rispetto al classico parallelismo dei dati mantenendo solo i dati partizionati richiesti su ciascuna GPU durante le diverse fasi di training.
3. Parallelismo tensoriale: invece di replicare il modello, il parallelismo tensore divide i parametri del modello in righe o colonne e li distribuisce tra le GPU. Ogni GPU opera su un insieme partizionato di parametri, gradienti e stati di ottimizzazione, con conseguente notevole risparmio di memoria.
4. Parallelismo delle condutture: questa tecnica suddivide i livelli del modello tra diverse GPU/worker, con ciascun dispositivo che esegue un sottoinsieme di livelli. Le attivazioni vengono passate tra i lavoratori, riducendo i picchi di memoria ma aumentando il sovraccarico di comunicazione.

La stima dell'utilizzo della memoria per questi metodi distribuiti non è banale poiché la distribuzione di parametri, gradienti, attivazioni e stati dell'ottimizzatore varia a seconda delle tecniche. Inoltre, diversi componenti come il corpo del trasformatore e la testa di modellazione del linguaggio possono mostrare comportamenti diversi di allocazione della memoria.

La soluzione LLMem

I ricercatori hanno recentemente proposto LLMem, una soluzione che stima accuratamente il consumo di memoria della GPU quando si applicano metodi di ottimizzazione distribuiti a LLM su più GPU.

Stima dell'utilizzo della memoria della GPU per la messa a punto di LLM preaddestrato

LLMem considera fattori come la ricombinazione dei parametri prima del calcolo (ZeRO Stage 3), la raccolta dell'output nel passaggio all'indietro (parallelismo del tensore) 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 il picco di utilizzo della memoria GPU per la messa a punto di LLM su una singola GPU con tassi di errore fino all'1.6%, superando il tasso di errore medio di DNNMem all'avanguardia di 42.6%. Quando si applicano metodi di regolazione fine distribuiti a LLM con oltre un miliardo di parametri su più GPU, LLMem raggiunge un impressionante tasso di errore medio di 3.0%.