Optimalisatie van geheugen voor grote taalmodellen en fijnaanpassing

Grote taalmodellen (LLM’s) zoals GPT-4, Bloom en LLaMA hebben opmerkelijke capaciteiten bereikt door op te schalen tot miljarden parameters. Het inzetten van deze enorme modellen voor inferentie of fijnaanpassing is echter uitdagend vanwege hun enorme geheugeneisen. In deze technische blog zullen we technieken onderzoeken voor het schatten en optimaliseren van geheugengebruik tijdens LLM-inferentie en fijnaanpassing op verschillende hardware-configuraties.

Geheugeneisen begrijpen

Het geheugen dat nodig is om een LLM te laden, wordt voornamelijk bepaald door het aantal parameters en de numerieke precisie die wordt gebruikt om de parameters op te slaan. Een eenvoudige vuistregel is:

• Het laden van een model met X miljard parameters vereist ongeveer 4X GB VRAM in 32-bit float precisie
• Het laden van een model met X miljard parameters vereist ongeveer 2X GB VRAM in 16-bit bfloat16/float16 precisie

Bijvoorbeeld, het laden van het 175B parameter GPT-3 model zou ongeveer 350GB VRAM vereisen in bfloat16 precisie. Op dit moment bieden de grootste commercieel beschikbare GPU’s, zoals de NVIDIA A100 en H100, slechts 80GB VRAM, waardoor tensor parallelisme en model parallelisme technieken noodzakelijk zijn.

Tijdens inferentie wordt het geheugengebruik gedomineerd door de modelparameters en de tijdelijke activatietensors die worden gegenereerd. Een hoog niveau schatting voor het maximale geheugengebruik tijdens inferentie is de som van het geheugen dat nodig is om de modelparameters te laden en het geheugen voor activaties.

Inferentiegeheugen kwantificeren

Laten we de geheugeneisen voor inferentie kwantificeren met het OctoCode model, dat ongeveer 15 miljard parameters heeft in bfloat16 formaat (~ 31GB). We zullen de Transformers bibliotheek gebruiken om het model te laden en tekst te genereren:

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

<p>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)</p>

<p>prompt = "Vraag: Schrijf een Python-functie om bytes naar gigabytes om te zetten.\n\nAntwoord:"
result = pipe(prompt, max_new_tokens=60)[0]["generated_text"][len(prompt):]</p>

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

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

Uitvoer:

Het maximale GPU-geheugengebruik is ongeveer 29GB, wat overeenkomt met onze schatting van 31GB voor het laden van de modelparameters in bfloat16 formaat.

Inferentiegeheugen optimaliseren met kwantificatie

Terwijl bfloat16 de gebruikelijke precisie is die wordt gebruikt voor het trainen van LLM’s, hebben onderzoekers ontdekt dat het kwantificeren van de modelgewichten naar lagere precisie gegevenstypen zoals 8-bits integers (int8) of 4-bits integers het geheugengebruik aanzienlijk kan verminderen met minimale nauwkeurigheidsverlies voor inferentietaken zoals tekstgeneratie.

Laten we de geheugensaving zien van 8-bits en 4-bits kwantificatie van het OctoCode model:

</div>
# 8-bits kwantificatie
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>

Uitvoer:

# 4-bits kwantificatie
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())

Uitvoer:

Met 8-bits kwantificatie daalt het geheugengebruik van 31GB naar 15GB, terwijl 4-bits kwantificatie het nog verder vermindert tot slechts 9,5GB! Dit maakt het mogelijk om het 15B parameter OctoCode model uit te voeren op consumenten-GPU’s zoals de RTX 3090 (24GB VRAM).

Houd er echter rekening mee dat agressievere kwantificatie zoals 4-bits soms kan leiden tot nauwkeurigheidsverlies in vergelijking met 8-bits of bfloat16 precisie. Er is een compromis tussen geheugensaving en nauwkeurigheid dat gebruikers moeten evalueren voor hun specifieke gebruik.

Kwantificatie is een krachtige techniek die LLM-implementatie mogelijk maakt op resource-beperkte omgevingen zoals cloud-exemplaren, edge-apparaten of zelfs mobiele telefoons door het geheugengebruik drastisch te verminderen.

Geheugen schatten voor fijnaanpassing

Terwijl kwantificatie voornamelijk wordt gebruikt voor efficiënte inferentie, zijn technieken zoals tensor parallelisme en model parallelisme cruciaal voor het beheren van geheugeneisen tijdens het trainen of fijnaanpassen van grote taalmodellen.

Het maximale geheugengebruik tijdens fijnaanpassing is typisch 3-4 keer hoger dan bij inferentie vanwege extra geheugeneisen voor:

• Gradients
• Optimizer statussen
• Activaties van de voorwaartse pas die worden opgeslagen voor terugpropagatie

Een conservatieve schatting is dat fijnaanpassen van een LLM met X miljard parameters ongeveer 4 * (2X) = 8X GB VRAM vereist in bfloat16 precisie.

Bijvoorbeeld, fijnaanpassen van het 7B parameter LLaMA model zou ongeveer 7 * 8 = 56GB van VRAM per GPU vereisen in bfloat16 precisie. Dit overschrijdt de geheugencapaciteit van huidige GPU’s, waardoor gedistribueerde fijnaanpastechnieken noodzakelijk zijn.

Gedistribueerde fijnaanpastechnieken

Verschillende gedistribueerde fijnaanpassingsmethoden zijn voorgesteld om GPU-geheugensbeperkingen voor grote modellen te overwinnen:

1. Gegevensparallelisme: De klassieke gegevensparallelisme-benadering repliceert het hele model over meerdere GPU’s terwijl de trainingsgegevensbatches worden gesplitst en gedistribueerd. Dit vermindert de trainingsduur lineair met het aantal GPU’s, maar vermindert het maximale geheugengebruik per GPU niet.
2. ZeRO Stage 3: Een geavanceerde vorm van gegevensparallelisme die modelparameters, gradients en optimizer statussen over GPU’s partitioneert. Het vermindert geheugen in vergelijking met klassiek gegevensparallelisme door alleen de vereiste gepartitioneerde gegevens op elke GPU te houden tijdens verschillende fasen van het trainen.
3. Tensorparallelisme: In plaats van het model te repliceren, deelt tensorparallelisme de modelparameters in rijen of kolommen en distribueert ze over GPU’s. Elke GPU werkt met een gepartitioneerde set parameters, gradients en optimizer statussen, wat aanzienlijke geheugensaving oplevert.
4. Pijpparallelisme: Deze techniek partitioneert de modellagen over verschillende GPU’s/werkers, waarbij elke apparaat een subset van de lagen uitvoert. Activaties worden tussen werkers doorgegeven, waardoor het maximale geheugengebruik wordt verlaagd, maar de communicatie-overhead toeneemt.

Het schatten van geheugengebruik voor deze gedistribueerde methoden is niet triviaal, omdat de verdeling van parameters, gradients, activaties en optimizer statussen varieert over technieken. Bovendien kunnen verschillende componenten zoals de transformer body en de language modeling head verschillende geheugentoewijzingsgedrag vertonen.

De LLMem-oplossing

Onderzoekers hebben onlangs LLMem voorgesteld, een oplossing die nauwkeurig het GPU-geheugengebruik schat bij het toepassen van gedistribueerde fijnaanpastechnieken op LLM’s over meerdere GPU’s.

Estimating GPU Memory Usage for Fine-Tuning Pre-Trained LLM

LLMem houdt rekening met factoren zoals het hercombineren van parameters vóór berekening (ZeRO Stage 3), outputverzameling in de achterwaartse pas (tensorparallelisme) en de verschillende geheugentoewijzingsstrategieën voor de transformer body en de language modeling head.

Experimentele resultaten laten zien dat LLMem het maximale GPU-geheugengebruik voor fijnaanpassen van LLM’s op een enkele GPU kan schatten met foutpercentages tot 1,6%, waarmee het de state-of-the-art DNNMem’s gemiddelde foutpercentage van 42,6% overtreft. Wanneer gedistribueerde fijnaanpastechnieken worden toegepast op LLM’s met meer dan een miljard parameters op meerdere GPU’s, bereikt LLMem een indrukwekkend gemiddeld foutpercentage van 3,0%.