Optimalizace paměti pro inferenci a fine-tuning velkých jazykových modelů

Velké jazykové modely (LLM) jako GPT-4, Bloom a LLaMA dosáhly pozoruhodných schopností díky navýšení počtu parametrů na miliardy. Nicméně, nasazení těchto masivních modelů pro inferenci nebo fine-tuning je náročné kvůli jejich enormním požadavkům na paměť. V tomto technickém blogu budeme prozkoumávat techniky pro odhadování a optimalizaci spotřeby paměti během inferenze a fine-tuningu LLM napříč různými hardwarovými konfiguracemi.

Porozumění požadavkům na paměť

Paměť potřebná pro načtení LLM je primárně určena počtem parametrů a numerickou přesností, která se používá pro uložení parametrů. Jednoduché pravidlo je:

• Načtení modelu s X miliardami parametrů vyžaduje přibližně 4X GB VRAM v 32-bit plovoucí přesnosti
• Načtení modelu s X miliardami parametrů vyžaduje přibližně 2X GB VRAM v 16-bit bfloat16/float16 přesnosti

Například načtení modelu GPT-3 s 175 miliardami parametrů by vyžadovalo přibližně 350 GB VRAM v bfloat16 přesnosti. V současné době největší komerčně dostupné GPU jako NVIDIA A100 a H100 nabízejí pouze 80 GB VRAM, což vyžaduje techniky tensorového paralelismu a modelového paralelismu.

Během inferenze je stopa paměti dominována modelovými parametry a dočasnými aktivacemi vygenerovanými. Vysokou úroveň odhadu pro maximální spotřebu paměti během inferenze je součet paměti vyžadované pro načtení modelových parametrů a paměti pro aktivace.

Kvantifikace paměti pro inferenci

Podívejme se na kvantifikaci požadavků na paměť pro inferenci pomocí modelu OctoCode, který má kolem 15 miliard parametrů v bfloat16 formátu (~ 31 GB). Budeme používat Transformers library pro načtení modelu a generování textu:

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

Výstup:

Maximální spotřeba paměti na GPU je kolem 29 GB, což odpovídá našemu odhadu 31 GB pro načtení modelových parametrů v bfloat16 formátu.

Optimalizace paměti pro inferenci pomocí kvantizace

Zatímco bfloat16 je běžná přesnost používaná pro trénování LLM, výzkumníci zjistili, že kvantizace modelových váh na nižší přesnost, jako jsou 8bitové celá čísla (int8) nebo 4bitové celá čísla, může výrazně snížit spotřebu paměti s minimální ztrátou přesnosti pro úkoly inferenze, jako je generování textu.

Podívejme se na úspory paměti z 8bitové a 4bitové kvantizace modelu OctoCode:

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

Výstup:

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

Výstup:

S 8bitovou kvantizací se požadavky na paměť snižují z 31 GB na 15 GB, zatímco 4bitová kvantizace dále snižuje spotřebu paměti na pouhých 9,5 GB! To umožňuje spuštění 15milionového parametrického modelu OctoCode na spotřebitelských GPU, jako je RTX 3090 (24 GB VRAM).

Nicméně, je třeba poznamenat, že agresivnější kvantizace, jako je 4bitová, může někdy vést k degradaci přesnosti ve srovnání s 8bitovou nebo bfloat16 přesností. Existuje kompromis mezi úsporami paměti a přesností, který uživatelé musí vyhodnotit pro své konkrétní použití.

Kvantizace je mocnou technikou, která může umožnit nasazení LLM v prostředích s omezenými zdroji, jako jsou cloudové instance, edge zařízení nebo dokonce mobilní telefony, výrazně snížením stopáže paměti.

Odhadování paměti pro fine-tuning

Zatímco kvantizace je primárně používána pro efektivní inferenci, techniky, jako je tensorový paralelismus a modelový paralelismus, jsou zásadní pro řízení požadavků na paměť během trénování nebo fine-tuningu velkých jazykových modelů.

Maximální spotřeba paměti během fine-tuningu je obvykle 3-4krát vyšší než během inferenze kvůli dodatečným požadavkům na paměť pro:

• Gradienty
• Stavy optimalizátoru
• Aktivace z předchozího průchodu uložené pro zpětnou propagaci

Konzervativní odhad je, že fine-tuning LLM s X miliardami parametrů vyžaduje kolem 4 * (2X) = 8X GB VRAM v bfloat16 přesnosti.

Například fine-tuning modelu LLaMA s 7 miliardami parametrů by vyžadoval přibližně 7 * 8 = 56 GB VRAM na GPU v bfloat16 přesnosti. To překračuje kapacitu současné GPU, což vyžaduje distribuované techniky fine-tuningu.

Distribuované techniky fine-tuningu

Byly navrženy several distribuované metody fine-tuningu, aby se překonaly omezení paměti GPU pro velké modely:

1. Data paralelismus: Klasický data paralelismus replikuje celý model napříč několika GPU, zatímco rozděluje a distribuuje datové dávky. To snižuje dobu trénování lineárně se počtem GPU, ale nesnižuje maximální požadavky na paměť na každém GPU.
2. ZeRO Stage 3: Pokročilá forma data paralelismu, která rozděluje modelové parametry, gradienty a stavy optimalizátoru napříč GPU. To snižuje spotřebu paměti ve srovnání s klasickým data paralelismem tak, že uchovává pouze požadované rozdělené údaje na každém GPU během různých fází trénování.
3. Tensor paralelismus: Místo replikace modelu rozděluje tensorový paralelismus modelové parametry do řádků nebo sloupců a distribuuje je napříč GPU. Každé GPU operuje na rozděleném souboru parametrů, gradientů a stavů optimalizátoru, což vede k podstatným úsporám paměti.
4. Pipeline paralelismus: Tato technika rozděluje modelové vrstvy napříč různými GPU/pracovníky, přičemž každá zařízení wykonává podmnožinu vrstev. Aktivace jsou předávány mezi pracovníky, což snižuje maximální spotřebu paměti, ale zvyšuje komunikační režii.

Odhadování spotřeby paměti pro tyto distribuované metody je netriviální, protože rozdělení parametrů, gradientů, aktivací a stavů optimalizátoru se liší napříč technikami. Kromě toho mohou různé komponenty, jako je transformerové tělo a jazykové modelovací hlava, vykazovat odlišné chování alokace paměti.

Řešení LLMem

Výzkumníci nedávno navrhli LLMem, řešení, které přesně odhaduje spotřebu paměti GPU při aplikaci distribuovaných metod fine-tuningu na LLM napříč několika GPU.

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

LLMem zohledňuje faktory, jako je rekombinace parametrů před výpočtem (ZeRO Stage 3), shromažďování výstupů v zpětném průchodu (tensorový paralelismus) a různé strategie alokace paměti pro transformerové tělo a jazykové modelovací hlavu.

Experimentální výsledky ukazují, že LLMem může odhadnout maximální spotřebu paměti GPU pro fine-tuning LLM na jednom GPU s chybovými sazbami až 1,6 %, překonávající průměrnou chybovou sazbu stávajícího DNNMem o 42,6 %. Při aplikaci distribuovaných metod fine-tuningu na LLM s více než miliardou parametrů na několik GPU dosahuje LLMem průměrné chybové sazby 3,0 %.