Optimering av minne for stor skalamodellinference og finjustering

Storskalamodeller (LLM) som GPT-4, Bloom og LLaMA har oppnådd bemerkelsesverdige evner ved å skala opp til milliarder av parametre. Imidlertid er det utfordrende å deployere disse massive modellene for inference eller finjustering på grunn av deres enorme minnekrev. I denne tekniske bloggen vil vi utforske teknikkene for å estimere og optimere minneforbruk under LLM-inference og finjustering på ulike hardvareoppsett.

Forståelse av minnekrev

Minnet som kreves for å laste en LLM bestemmes hovedsakelig av antallet parametre og den numeriske presisjonen som brukes til å lagre parameterne. En enkel regel er:

• Lasting av en modell med X milliarder parametre krever omtrent 4X GB av VRAM i 32-bit flytende punkt presisjon
• Lasting av en modell med X milliarder parametre krever omtrent 2X GB av VRAM i 16-bit bfloat16/float16 presisjon

For eksempel ville lasting av 175B parameter GPT-3-modellen kreve omtrent 350GB av VRAM i bfloat16 presisjon. Per i dag tilbyr de største kommersielle tilgjengelige GPU-ene, som NVIDIA A100 og H100, bare 80GB av VRAM, noe som nødvendiggjør tensor-parallellisme og modell-parallellisme-teknikker.

Under inference domineres minneavtrykket av modellparameterne og de midlertidige aktiverings-tensorer som produseres. En høynivå-estimat for toppminnebruk under inference er summen av minnet som kreves for å laste modellparameterne og minnet for aktiveringer.

Kvantifisering av inferensminne

La oss kvantifisere minnekrev for inference ved å bruke OctoCode-modellen, som har omtrent 15 milliarder parametre i bfloat16-format (~ 31GB). Vi vil bruke Transformers-biblioteket til å laste 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;Question: Please write a Python function to convert bytes to gigabytes.\n\nAnswer:&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())

Utdata:

Topp-GPU-minnebruk er omtrent 29GB, som stemmer overens med vår estimat på 31GB for å laste modellparameterne i bfloat16-format.

Optimering av inferensminne med kvantisering

Selv om bfloat16 er den vanlige presisjonen som brukes for å trene LLM, har forskere funnet at kvantisering av modellvektene til lavere presisjon datatyper som 8-bit heltall (int8) eller 4-bit heltall kan redusere minnebruk betydelig med minimal tap av nøyaktighet for inferensoppgaver som tekstgenerering.

La oss se på minnesparingene fra 8-bit og 4-bit kvantisering av OctoCode-modellen:

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

Utdata:

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

Utdata:

Med 8-bit kvantisering reduseres minnekrevet fra 31GB til 15GB, mens 4-bit kvantisering reduserer det ytterligere til bare 9,5GB! Dette gjør det mulig å kjøre 15B parameter OctoCode-modellen på forbruker-GPU-er som RTX 3090 (24GB VRAM).

Imidlertid bør det merkes at mer aggressive kvantisering som 4-bit kan noen ganger føre til nøyaktighetsnedgang sammenlignet med 8-bit eller bfloat16 presisjon. Det er en avveining mellom minnesparing og nøyaktighet som brukerne må vurdere for sin brukssak.

Kvantisering er en kraftfull teknikk som kan muliggjøre LLM-deployering på ressursbegrensede miljøer som skytjenester, edge-enheter eller selv mobiltelefoner ved å redusere minneavtrykket drastisk.

Estimering av minne for finjustering

Selv om kvantisering hovedsakelig brukes for effektiv inference, er teknikkene som tensor-parallellisme og modell-parallellisme avgjørende for å håndtere minnekrev under trening eller finjustering av store språkmodeller.

Toppminneforbruk under finjustering er vanligvis 3-4 ganger høyere enn under inference på grunn av ekstra minnekrev for:

• Gradient
• Optimerer-tilstander
• Aktiveringer fra fremover-passasje lagret for bakover-propagasjon

En konservativ estimat er at finjustering av en LLM med X milliarder parametre krever omtrent 4 * (2X) = 8X GB av VRAM i bfloat16 presisjon.

For eksempel ville finjustering av 7B parameter LLaMA-modellen kreve omtrent 7 * 8 = 56GB av VRAM per GPU i bfloat16 presisjon. Dette overstiger minnekapasiteten til nåværende GPU-er, noe som nødvendiggjør distribuert finjusteringsteknikker.

Distribuert finjusteringsteknikker

Flere distribuerte finjusteringsteknikker er blitt foreslått for å overvinne GPU-minnebegrensninger for store modeller:

1. Data-parallellisme: Den klassiske data-parallellisme-tilnærmingen repliserer hele modellen på flere GPU-er mens den splitter og distribuerer treningdata-batchene. Dette reduserer treningstiden lineært med antallet GPU-er, men reduserer ikke toppminnekrevet på hver enkelt GPU.
2. ZeRO Stage 3: En avansert form for data-parallellisme som partitionerer modellparameterne, gradientene og optimerer-tilstandene på tvers av GPU-er. Den reduserer minne sammenlignet med klassisk data-parallellisme ved å bare holde den nødvendige partitionerte data på hver GPU under ulike faser av trening.
3. Tensor-parallellisme: I stedet for å replikere modellen, deler tensor-parallellisme modellparameterne i rader eller kolonner og distribuerer dem på flere GPU-er. Hver GPU opererer på en partitionert mengde parametre, gradient og optimerer-tilstander, noe som fører til betydelige minnesparing.
4. Pipeline-parallellisme: Denne teknikken partitionerer modelllagene på tvers av ulike GPU-er/arbeidere, hvor hver enhet kjører en undermengde av lagene. Aktiveringer sendes mellom arbeidere, noe som reduserer toppminne, men øker kommunikasjons-overhodet.

Estimering av minnebruk for disse distribuerte metodene er ikke-trivielt, da distribusjonen av parametre, gradient, aktiveringer og optimerer-tilstander varierer på tvers av teknikker. I tillegg kan ulike komponenter som transformer-kroppen og språk-modell-hodet vise forskjellige minne-allokering-atferd.

LLMem-løsningen

Forskere har nylig foreslått LLMem, en løsning som nøyaktig estimerer GPU-minneforbruk når distribuerte finjusteringsteknikker brukes på LLM-er på flere GPU-er.

Estimering av GPU-minnebruk for finjustering av forhånds-trente LLM

LLMem tar hensyn til faktorer som gjenkobling av parametre før beregning (ZeRO Stage 3), utgangssamling i bakover-passasje (tensor-parallellisme) og ulike minne-allokering-strategier for transformer-kroppen og språk-modell-hodet.

Eksperimentelle resultater viser at LLMem kan estimere topp-GPU-minnebruk for finjustering av LLM-er på en enkelt GPU med feilrater på opptil 1,6%, og overgår den gjeldende DNNMem-feilraten på 42,6%. Når distribuerte finjusteringsteknikker brukes på LLM-er med over en milliard parametre på flere GPU-er, oppnår LLMem en imponerende gjennomsnittlig feilrate på 3,0%.