Optimering af hukommelse til stor sprogmodelinference og finjustering

Store sprogmodeller (LLM’er) som GPT-4, Bloom og LLaMA har opnået bemærkelsesværdige evner ved at skala op til milliarder af parametre. however, at deployere disse massive modeller til inference eller finjustering er udfordrende på grund af deres enorme hukommelseskrav. I denne tekniske blog vil vi udforske teknikker til at estimere og optimere hukommelsesforbrug under LLM-inference og finjustering på tværs af forskellige hardware-konfigurationer.

Forståelse af hukommelseskrav

Hukommelsen, der kræves for at indlæse en LLM, bestemmes primært af antallet af parametre og den numeriske præcision, der bruges til at gemme parametrene. En simpel regel er:

• At indlæse en model med X milliarder parametre kræver omtrent 4X GB VRAM i 32-bit float-præcision
• At indlæse en model med X milliarder parametre kræver omtrent 2X GB VRAM i 16-bit bfloat16/float16-præcision

For eksempel ville indlæsning af 175B parameter GPT-3 model kræve omtrent 350GB VRAM i bfloat16-præcision. Som det er i dag, tilbyder de største kommercielt tilgængelige GPU’er som NVIDIA A100 og H100 kun 80GB VRAM, hvilket nødvendiggør tensor-parallellisme og model-parallellisme-teknikker.

Under inference domineres hukommelsesaftrykket af modelparametrene og de midlertidige aktiverings-tensorer, der produceres. En højniveau-estimation for peak-hukommelsesforbrug under inference er summen af hukommelsen, der kræves for at indlæse modelparametrene, og hukommelsen for aktiveringer.

Kvantificering af inference-hukommelse

Lad os kvantificere hukommelseskravene for inference ved hjælp af OctoCode-modellen, der har omkring 15 milliarder parametre i bfloat16-format (~ 31GB). Vi vil bruge Transformers-biblioteket til at indlæse modellen og generere tekst:

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;Spørgsmål: Skriv en Python-funktion til at konvertere bytes til gigabyte.\n\nSvar:&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:

Peak-GPU-hukommelsesforbrug er omkring 29GB, hvilket stemmer overens med vores estimation på 31GB for at indlæse modelparametrene i bfloat16-format.

Optimering af inference-hukommelse med kvantificering

Selvom bfloat16 er den almindelige præcision, der bruges til at træne LLM’er, har forskere fundet, at kvantificering af modelvægte til lavere præcisionsdata typer som 8-bit integers (int8) eller 4-bit integers kan betydeligt reducere hukommelsesforbrug med minimal tab af nøjagtighed for inference-opgaver som tekstgenerering.

Lad os se hukommelsesbesparelsen fra 8-bit og 4-bit kvantificering af OctoCode-modellen:

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

# 4-bit kvantificering
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:

Med 8-bit kvantificering falder hukommelseskravet fra 31GB til 15GB, mens 4-bit kvantificering reducerer det yderligere til kun 9,5GB! Dette tillader kørsel af 15B parameter OctoCode-modellen på forbruger-GPU’er som RTX 3090 (24GB VRAM).

Men bemærk, at mere aggressive kvantificering som 4-bit kan nogle gange føre til en nedgang i nøjagtighed sammenlignet med 8-bit eller bfloat16-præcision. Der er en balance mellem hukommelsesbesparelse og nøjagtighed, som brugere skal evaluere for deres brugstilfælde.

Kvantificering er en kraftfuld teknik, der kan aktivere LLM-deployment på ressource-begrænsede miljøer som skyen, edge-enheder eller endda mobiltelefoner ved at drastisk reducere hukommelsesaftrykket.

Estimering af hukommelse til finjustering

Selvom kvantificering primært bruges til effektiv inference, er teknikker som tensor-parallellisme og model-parallellisme afgørende for at håndtere hukommelseskrav under træning eller finjustering af store sprogmodeller.

Peak-hukommelsesforbrug under finjustering er typisk 3-4 gange højere end under inference på grund af ekstra hukommelseskrav til:

• Gradients
• Optimizer-tilstande
• Aktiveringer fra forward-passen gemt til backpropagation

En konservativ estimation er, at finjustering af en LLM med X milliarder parametre kræver omtrent 4 * (2X) = 8X GB VRAM i bfloat16-præcision.

For eksempel ville finjustering af 7B parameter LLaMA-modellen kræve omtrent 7 * 8 = 56GB VRAM per GPU i bfloat16-præcision. Dette overstiger hukommelseskapaciteten af nuværende GPU’er, hvilket nødvendiggør distribueret finjusteringsteknikker.

Distribuerede finjusteringsteknikker

Forskere har foreslået flere distribuerede finjusteringmetoder for at overvinde GPU-hukommelsesbegrænsninger for store modeller:

1. Data-parallellisme: Den klassiske data-parallellisme-tilgang replikerer hele modellen på tværs af flere GPU’er, mens træningsdata-batchene deles og distribueres. Dette reducerer træningstiden lineært med antallet af GPU’er, men reducerer ikke peak-hukommelseskravet på hver GPU.
2. ZeRO Stage 3: En avanceret form for data-parallellisme, der partitionerer modelparametre, grader og optimizer-tilstande på tværs af GPU’er. Det reducerer hukommelse sammenlignet med klassisk data-parallellisme ved at holde kun den nødvendige partitionerede data på hver GPU under forskellige faser af træning.
3. Tensor-parallellisme: I stedet for at replikere modellen, partitionerer tensor-parallellisme modelparametre i rækker eller kolonner og distribuerer dem på tværs af GPU’er. Hver GPU opererer på en partitioneret samling af parametre, grader og optimizer-tilstande, hvilket fører til betydelige hukommelsesbesparelser.
4. Rørlednings-parallellisme: Denne teknik partitionerer model-lag på tværs af forskellige GPU’er/arbejdere, hvor hver enhed udfører en undermængde af lagene. Aktiveringer overføres mellem arbejdere, hvilket reducerer peak-hukommelse, men øger kommunikations-overhead.

At estimere hukommelsesforbrug for disse distribuerede metoder er ikke-trivialt, da distributionen af parametre, grader, aktiveringer og optimizer-tilstande varierer på tværs af teknikker. Desuden kan forskellige komponenter som transformer-kroppen og sprogmodel-hovedet vise forskellige hukommelses-allokering-adfærd.

LLMem-løsningen

Forskere har foreslået LLMem, en løsning, der nøjagtigt estimerer GPU-hukommelsesforbrug, når distribuerede finjusteringsteknikker anvendes på store sprogmodeller på tværs af flere GPU’er.

Estimering af GPU-hukommelsesforbrug til finjustering af forudtrænede LLM’er

LLMem tager hensyn til faktorer som gensammensætning af parametre før beregning (ZeRO Stage 3), output-samling i backward-passet (tensor-parallellisme) og forskellige hukommelses-allokering-strategier for transformer-kroppen og sprogmodel-hovedet.

Eksperimentelle resultater viser, at LLMem kan estimere peak-GPU-hukommelsesforbrug for finjustering af LLM’er på en enkelt GPU med fejl-rater på op til 1,6%, hvilket overgår den gennemsnitlige fejl-rate på 42,6% for den hidtidige DNNMem. Når distribuerede finjusteringsteknikker anvendes på LLM’er med over en milliard parametre på flere GPU’er, opnår LLMem en imponerende gennemsnitlig fejl-rate på 3,0%.