Optimering af hukommelse til stor sprogmodelslutning og finjustering

Store sprogmodeller (LLM'er) som GPT-4, Bloom og LLaMA har opnået bemærkelsesværdige egenskaber ved at skalere op til milliarder af parametre. Det er imidlertid en udfordring at implementere disse massive modeller til inferens eller finjustering på grund af deres enorme hukommelseskrav. I denne tekniske blog vil vi udforske teknikker til at estimere og optimere hukommelsesforbrug under LLM-inferens og finjustering på tværs af forskellige hardwareopsætninger.

Forståelse af hukommelseskrav

Den nødvendige hukommelse for at indlæse en LLM er primært bestemt af antallet af parametre og den numeriske præcision, der bruges til at gemme parametrene. En simpel tommelfingerregel er:

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

For eksempel vil indlæsning af 175B parameter GPT-3-modellen kræve cirka 350 GB VRAM i bfloat16-præcision. I dag tilbyder de største kommercielt tilgængelige GPU'er som NVIDIA A100 og H100 kun 80 GB VRAM, hvilket nødvendiggør tensorparallelisme og modelparallelismeteknikker.

Under inferens domineres hukommelsesfodaftrykket af modelparametrene og de producerede midlertidige aktiveringstensorer. Et estimat på højt niveau for den maksimale hukommelsesbrug under inferens er summen af ​​den hukommelse, der kræves for at indlæse modelparametrene og hukommelsen til aktiveringer.

Kvantificering af inferenshukommelse

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

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())

Output:

Det maksimale GPU-hukommelsesforbrug er omkring 29 GB, hvilket stemmer overens med vores estimat på 31 GB for indlæsning af modelparametrene i bfloat16-format.

Optimering af inferenshukommelse med kvantisering

Mens bfloat16 er den almindelige præcision, der bruges til træning af LLM'er, har forskere fundet ud af, at kvantificering af modelvægtene for at sænke præcisionsdatatyper som 8-bit heltal (int8) eller 4-bit heltal kan reducere hukommelsesforbruget med minimalt nøjagtighedstab for inferensopgaver som f.eks. tekstgenerering.

Lad os se hukommelsesbesparelserne fra 8-bit og 4-bit kvantisering af OctoCode-modellen:

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

Output:

# 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())

Output:

Med 8-bit kvantisering falder hukommelsesbehovet fra 31 GB til 15 GB, mens 4-bit kvantisering reducerer det yderligere til kun 9.5 GB! Dette gør det muligt at køre 15B parameter OctoCode-modellen på forbruger-GPU'er som RTX 3090 (24GB VRAM).

Bemærk dog, at mere aggressiv kvantisering som 4-bit nogle gange kan føre til nøjagtighedsforringelse sammenlignet med 8-bit eller bfloat16-præcision. Der er en afvejning mellem hukommelsesbesparelser og nøjagtighed, som brugerne bør vurdere for deres brugssituation.

Kvantisering er en kraftfuld teknik, der kan aktivere LLM-implementering i ressourcebegrænsede miljøer som cloud-instanser, edge-enheder eller endda mobiltelefoner ved at reducere hukommelsesfodaftrykket drastisk.

Estimering af hukommelse til finjustering

Mens kvantisering primært bruges til effektiv inferens, er teknikker som tensorparallelisme og modelparallelisme afgørende for håndtering af hukommelseskrav under træningen eller finjustering af store sprogmodeller.

Det maksimale hukommelsesforbrug under finjustering er typisk 3-4 gange højere end inferens på grund af yderligere hukommelseskrav til:

• farveforløb
• Optimizer angiver
• Aktiveringer fra det fremadgående gennemløb er gemt til backpropagation

Et konservativt skøn er, at finjustering af en LLM med X milliarder parametre kræver ca 4 * (2X) = 8X GB af VRAM i bfloat16-præcision.

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

Distribuerede finjusteringsteknikker

Adskillige distribuerede finjusteringsmetoder er blevet foreslået for at overvinde GPU-hukommelsesbegrænsninger for store modeller:

1. Dataparallelisme: Den klassiske dataparallelisme-tilgang replikerer hele modellen på tværs af flere GPU'er, mens træningsdatabatchene opdeles og distribueres. Dette reducerer træningstiden lineært med antallet af GPU'er, men reducerer ikke peak-hukommelseskravet på hver GPU.
2. ZeRO trin 3: En avanceret form for dataparallelisme, der opdeler modelparametre, gradienter og optimeringstilstande på tværs af GPU'er. Det reducerer hukommelsen sammenlignet med klassisk dataparallelisme ved kun at beholde de nødvendige opdelte data på hver GPU under forskellige faser af træningen.
3. Tensor parallelisme: I stedet for at replikere modellen opdeler tensorparallelisme modelparametrene i rækker eller kolonner og fordeler dem på tværs af GPU'er. Hver GPU fungerer på et opdelt sæt af parametre, gradienter og optimeringstilstande, hvilket fører til betydelige hukommelsesbesparelser.
4. Pipeline parallelisme: Denne teknik opdeler modellagene på tværs af forskellige GPU'er/arbejdere, hvor hver enhed udfører en delmængde af lagene. Aktiveringer overføres mellem arbejdere, hvilket reducerer spidsbelastningshukommelsen, men øger kommunikationsomkostningerne.

Det er ikke-trivielt at estimere hukommelsesbrug for disse distribuerede metoder, da fordelingen af ​​parametre, gradienter, aktiveringer og optimeringstilstande varierer på tværs af teknikker. Desuden kan forskellige komponenter som transformatorkroppen og sprogmodelleringshovedet udvise forskellig hukommelsesallokeringsadfærd.

LLMem-løsningen

Forskere foreslog for nylig LLMem, en løsning, der nøjagtigt estimerer GPU-hukommelsesforbrug, når der anvendes distribuerede finjusteringsmetoder til LLM'er på tværs af flere GPU'er.

Estimering af GPU-hukommelsesforbrug til finjustering af præ-trænet LLM

LLMem overvejer faktorer som rekombination af parametre før beregning (ZeRO Stage 3), outputsamling i baglæns passage (tensorparallelisme) og de forskellige hukommelsesallokeringsstrategier for transformatorkroppen og sprogmodelleringshovedet.

Eksperimentelle resultater viser, at LLMem kan estimere maksimal GPU-hukommelsesbrug til finjustering af LLM'er på en enkelt GPU med fejlrater på op til 1.6 %, hvilket overgår den avancerede DNNMems gennemsnitlige fejlrate på 42.6 %. Når man anvender distribuerede finjusteringsmetoder til LLM'er med over en milliard parametre på flere GPU'er, opnår LLMem en imponerende gennemsnitlig fejlrate på 3.0 %.