Optimering av minne för stora språkmodeller för inferens och finjustering

Stora språkmodeller (LLM) som GPT-4, Bloom och LLaMA har uppnått anmärkningsvärda förmågor genom att skala upp till miljarder parametrar. Det är dock utmanande att distribuera dessa massiva modeller för inferens eller finjustering på grund av deras enorma minneskrav. I den här tekniska bloggen kommer vi att utforska tekniker för att uppskatta och optimera minnesanvändning under LLM-inferens och finjustering över olika maskinvarukonfigurationer.

Förstå minneskrav

Minnet som krävs för att ladda en LLM bestäms främst av antalet parametrar och den numeriska precision som används för att lagra parametrarna. En enkel tumregel är:

• Att ladda en modell med X miljarder parametrar kräver ungefär 4X GB VRAM i 32-bit flyttalsprecision
• Att ladda en modell med X miljarder parametrar kräver ungefär 2X GB VRAM i 16-bit bfloat16/float16 precision

Till exempel skulle laddning av 175B parameter GPT-3-modellen kräva cirka 350 GB VRAM i bfloat16 precision. För närvarande erbjuder de största kommersiellt tillgängliga GPU:erna, som NVIDIA A100 och H100, endast 80 GB VRAM, vilket kräver tensorparallelism och modellparallelismstekniker.

Under inferens domineras minnesavtrycket av modellparametrarna och de tillfälliga aktiverings tensorer som produceras. En högnivåuppskattning för toppminnesanvändning under inferens är summan av minnet som krävs för att ladda modellparametrarna och minnet för aktiveringar.

Kvantifiering av inferensminne

Låt oss kvantifiera minneskraven för inferens med hjälp av OctoCode-modellen, som har cirka 15 miljarder parametrar i bfloat16-format (~ 31 GB). Vi kommer att använda Transformers-biblioteket för att ladda modellen och generera 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;Fråga: Vänligen skriv en Python-funktion för att konvertera byte till 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())

Utdata:

Toppminnesanvändningen på GPU är cirka 29 GB, vilket överensstämmer med vår uppskattning på 31 GB för att ladda modellparametrarna i bfloat16-format.

Optimering av inferensminne med kvantisering

Medan bfloat16 är den vanliga precisionen som används för att träna LLM, har forskare funnit att kvantisering av modellvikter till lägre precision datatyper som 8-bitars heltal (int8) eller 4-bitars heltal kan minska minnesanvändningen avsevärt med minimal förlust av noggrannhet för inferenstyper som textgenerering.

Låt oss se minnessparandet från 8-bitars och 4-bitars kvantisering av OctoCode-modellen:

</div>
# 8-bitars 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-bitars 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-bitars kvantisering minskar minneskravet från 31 GB till 15 GB, medan 4-bitars kvantisering minskar det ytterligare till endast 9,5 GB! Detta möjliggör körning av 15B parameter OctoCode-modellen på konsument-GPU:er som RTX 3090 (24 GB VRAM).

Men observera att mer aggressiv kvantisering som 4-bitars kan ibland leda till försämrad noggrannhet jämfört med 8-bitars eller bfloat16 precision. Det finns en avvägning mellan minnessparande och noggrannhet som användare bör utvärdera för sin användning.

Kvantisering är en kraftfull teknik som kan möjliggöra distribution av LLM på resursbegränsade miljöer som molninstanser, edge-enheter eller till och med mobiltelefoner genom att drastiskt minska minnesavtrycket.

Uppskattning av minne för finjustering

Medan kvantisering främst används för effektiv inferens, är tekniker som tensorparallelism och modellparallelism avgörande för att hantera minneskrav under träning eller finjustering av stora språkmodeller.

Toppminnesanvändningen under finjustering är vanligtvis 3-4 gånger högre än under inferens på grund av ytterligare minneskrav för:

• Gradient
• Optimeringslägen
• Aktiveringar från framåtpassagen som lagras för bakåtpassagen

En konservativ uppskattning är att finjustering av en LLM med X miljarder parametrar kräver cirka 4 * (2X) = 8X GB VRAM i bfloat16 precision.

Till exempel skulle finjustering av 7B parameter LLaMA-modellen kräva cirka 7 * 8 = 56 GB VRAM per GPU i bfloat16 precision. Detta överstiger minneskapaciteten hos nuvarande GPU:er, vilket kräver distribuerad finjusteringsteknik.

Distribuerad finjusteringsteknik

Flera distribuerade finjusteringsmetoder har föreslagits för att övervinna GPU-minnesbegränsningar för stora modeller:

1. Data Parallelism: Den klassiska data parallelism-metoden replikerar hela modellen över flera GPU:er medan datatoken är uppdelade och distribuerade. Detta minskar tränings­tiden linjärt med antalet GPU:er men minskar inte toppminneskravet på varje GPU.
2. ZeRO Stage 3: En avancerad form av data parallelism som partitionerar modellparametrar, gradient och optimeringslägen över GPU:er. Det minskar minnet jämfört med klassisk data parallelism genom att endast behålla den partitionerade datan som behövs på varje GPU under olika träningsfaser.
3. Tensor Parallelism: Istället för att replikera modellen, delar tensorparallelism modellparametrarna i rader eller kolumner och distribuerar dem över GPU:er. Varje GPU arbetar med en partitionerad uppsättning parametrar, gradient och optimeringslägen, vilket leder till betydande minnessparande.
4. Pipeline Parallelism: Denna teknik partitionerar modell­lagren över olika GPU:er/arbetare, där varje enhet kör en delmängd av lagren. Aktiveringar skickas mellan arbetare, vilket minskar toppminnet men ökar kommunikations­över­huvudet.

Att uppskatta minnesanvändning för dessa distribuerade metoder är inte trivialt, eftersom fördelningen av parametrar, gradient, aktiveringar och optimeringslägen varierar mellan tekniker. Dessutom kan olika komponenter som transformer­kroppen och språk­modell­huvudet visa olika minnes­allokerings­beteende.

LLMem-lösningen

Forskare har nyligen föreslagit LLMem, en lösning som noggrant uppskattar GPU-minnesanvändning när distribuerade finjusteringsmetoder tillämpas på LLM:er över flera GPU:er.

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

LLMem tar hänsyn till faktorer som omkombinering av parametrar före beräkning (ZeRO Stage 3), utgångssamling i bakåtpassagen (tensorparallelism) och olika minnesallokeringsstrategier för transformer­kroppen och språk­modell­huvudet.

Experimentella resultat visar att LLMem kan uppskatta topp-GPU-minnesanvändning för finjustering av LLM:er på en enda GPU med fel­frekvenser på upp till 1,6 %, vilket överträffar den tidigare bästa DNNMem:s genomsnittliga fel­frekvens på 42,6 %. När distribuerad finjustering tillämpas på LLM:er med över en miljard parametrar på flera GPU:er uppnår LLMem en imponerande genomsnittlig fel­frekvens på 3,0 %.