Optimizarea Memoriei pentru Inferența și Reglarea Modelelor Lingvistice Mari

Modelele lingvistice mari (LLM) precum GPT-4, Bloom și LLaMA au atins capacități remarcabile prin escaladarea la miliarde de parametri. Cu toate acestea, implementarea acestor modele masive pentru inferență sau reglare este dificilă din cauza cerințelor lor imense de memorie. În acest blog tehnic, vom explora tehnici pentru estimarea și optimizarea consumului de memorie în timpul inferenței și reglării LLM pe diverse configurații de hardware.

Înțelegerea Cerințelor de Memorare

Memoria necesară pentru încărcarea unui LLM este determinată în primul rând de numărul de parametri și precizia numerică utilizată pentru stocarea parametrilor. O regulă simplă este:

• Încărcarea unui model cu X miliarde de parametri necesită aproximativ 4X GB de VRAM în precizie 32-bit float
• Încărcarea unui model cu X miliarde de parametri necesită aproximativ 2X GB de VRAM în precizie 16-bit bfloat16/float16

De exemplu, încărcarea modelului GPT-3 cu 175 de miliarde de parametri ar necesita aproximativ 350 GB de VRAM în precizie bfloat16. În prezent, cele mai mari GPU-uri comerciale disponibile, cum ar fi NVIDIA A100 și H100, oferă doar 80 GB de VRAM, necesitând tehnici de paralelism tensorial și model paralel.

În timpul inferenței, amprenta de memorie este dominată de parametrii modelului și tensorii de activare temporari produși. O estimare de nivel înalt pentru utilizarea maximă de memorie în timpul inferenței este suma memoriei necesare pentru încărcarea parametrilor modelului și a memoriei pentru activări.

Cuantificarea Memoriei de Inferență

Să cuantificăm cerințele de memorie pentru inferență utilizând modelul OctoCode, care are aproximativ 15 miliarde de parametri în format bfloat16 (~ 31 GB). Vom utiliza biblioteca Transformers pentru a încărca modelul și a genera text:

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;Întrebare: Vă rog să scrieți o funcție Python pentru a converti octeți în gigaocteți.\n\nRăspuns:&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())

Ieșire:

Utilizarea maximă de memorie GPU este de aproximativ 29 GB, ceea ce se aliniază cu estimarea noastră de 31 GB pentru încărcarea parametrilor modelului în format bfloat16.

Optimizarea Memoriei de Inferență cu Cuantificare

Deși bfloat16 este precizia comună utilizată pentru antrenarea LLM, cercetătorii au descoperit că cuantificarea greutăților modelului la tipuri de date cu precizie mai mică, cum ar fi numerele întregi de 8 biți (int8) sau numerele întregi de 4 biți, poate reduce semnificativ utilizarea memoriei cu o pierdere minimă de precizie pentru sarcini de inferență, cum ar fi generarea de text.

Să vedem economiile de memorie de la cuantificarea de 8 biți și 4 biți a modelului OctoCode:

</div>
# Cuantificare de 8 biți
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>

Ieșire:

# Cuantificare de 4 biți
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())

Ieșire:

Cu cuantificarea de 8 biți, cerința de memorie scade de la 31 GB la 15 GB, în timp ce cuantificarea de 4 biți o reduce și mai mult la doar 9,5 GB! Acest lucru permite rularea modelului OctoCode de 15 miliarde de parametri pe GPU-uri de consum, cum ar fi RTX 3090 (24 GB VRAM).

Cu toate acestea, trebuie remarcat că o cuantificare mai agresivă, cum ar fi cea de 4 biți, poate duce uneori la o scădere a preciziei în comparație cu precizia de 8 biți sau bfloat16. Există un compromis între economiile de memorie și precizie pe care utilizatorii trebuie să îl evalueze pentru cazul lor de utilizare.

Cuantificarea este o tehnică puternică care poate permite implementarea LLM pe medii cu resurse limitate, cum ar fi instanțe cloud, dispozitive edge sau chiar telefoane mobile, prin reducerea drastică a amprentei de memorie.

Estimarea Memoriei pentru Reglare

În timp ce cuantificarea este utilizată în primul rând pentru inferență eficientă, tehnici precum paralelismul tensorial și paralelismul modelului sunt cruciale pentru gestionarea cerințelor de memorie în timpul antrenării sau reglării modelelor lingvistice mari.

Consumul maxim de memorie în timpul reglării este, în general, de 3-4 ori mai mare decât în timpul inferenței, din cauza cerințelor suplimentare de memorie pentru:

• Gradient
• Stări de optimizator
• Activări din trecerea înainte stocate pentru retropropagare

O estimare conservatoare este că reglarea unui LLM cu X miliarde de parametri necesită aproximativ 4 * (2X) = 8X GB de VRAM în precizie bfloat16.

De exemplu, reglarea modelului LLaMA de 7 miliarde de parametri ar necesita aproximativ 7 * 8 = 56 GB de VRAM pe GPU în precizie bfloat16. Acest lucru depășește capacitatea de memorie a GPU-urilor actuale, necesitând tehnici de reglare distribuite.

Tehnici de Reglare Distribuite

Au fost propuse mai multe metode de reglare distribuite pentru a depăși constrângerile de memorie ale GPU pentru modele mari:

1. Paralelismul de date: Abordarea clasică de paralelism de date replică întregul model pe mai multe GPU-uri, în timp ce împarte și distribuie loturile de date de antrenare. Acest lucru reduce timpul de antrenare liniar cu numărul de GPU-uri, dar nu reduce cerința maximă de memorie pe fiecare GPU.
2. Etapa ZeRO 3: O formă avansată de paralelism de date care împarte parametrii modelului, gradientul și stările optimizatorului pe GPU-uri. Acesta reduce memoria în comparație cu paralelismul de date clasic, păstrând doar datele partitionate necesare pe fiecare GPU în timpul diferitelor faze de antrenare.
3. Paralelismul tensorial: În loc de a replica modelul, paralelismul tensorial împarte parametrii modelului în rânduri sau coloane și le distribuie pe GPU-uri. Fiecare GPU operează pe un set de parametri, gradient și stări de optimizator partitionate, ceea ce conduce la economii substanțiale de memorie.
4. Paralelismul de conductă: Această tehnică împarte straturile modelului pe diferite GPU-uri/lucrători, fiecare dispozitiv executând un subset de straturi. Activările sunt transmise între lucrători, reducând memoria maximă, dar crescând supărarea de comunicare.

Estimarea utilizării de memorie pentru aceste metode distribuite nu este trivială, deoarece distribuția parametrilor, gradientului, activărilor și stărilor de optimizator variază între tehnici. Mai mult, componente diferite, cum ar fi corpul transformatorului și capul de modelare a limbajului, pot prezenta comportamente de alocare a memoriei diferite.

Soluția LLMem

Cercetătorii au propus recent LLMem, o soluție care estimează cu acuratețe consumul de memorie GPU atunci când se aplică metode de reglare distribuite la LLM pe mai multe GPU-uri.

Estimarea Utilizării de Memorare GPU pentru Reglarea Preantrenată a LLM

LLMem ia în considerare factori precum recombinația parametrilor înainte de calcul (Etapa ZeRO 3), colectarea ieșirilor în trecerea inversă (paralelism tensorial) și strategiile diferite de alocare a memoriei pentru corpul transformatorului și capul de modelare a limbajului.

Rezultatele experimentale arată că LLMem poate estima utilizarea maximă de memorie GPU pentru reglarea LLM pe o singură GPU cu rate de eroare de până la 1,6%, depășind rata medie de eroare de 42,6% a lui DNNMem. Atunci când se aplică metode de reglare distribuite la LLM cu peste un miliard de parametri pe mai multe GPU-uri, LLMem obține o rată medie de eroare de 3,0%.