Minnesbehov för transformermodell: så beräknar du VRAM

Minnesbehov för transformermodell: så beräknar du VRAM

För att driftsätta och köra stora språkmodeller krävs en noggrann uppskattning av vilket minnesbehov för transformermodell som krävs i grafikkortets minne (VRAM). Det räcker inte att enbart titta på antalet parametrar. Vid praktisk inferens och beräkning fördelas det totala minnesbehovet mellan modellvikter, cache för uppmärksamhet (KV-cache), aktiveringar och systemets arbetsbuffertar.

I vår grundläggande guide om neur.se visar vi hur du räknar fram vikterna i parametrar i en transformermodell. I denna artikel går vi igenom hur du beräknar det faktiska dynamiska minnesbehovet när modellen är i drift.

Modellvikternas minnesbehov vid olika precisioner

Den statiska grunden för minnesbehov för transformermodell avgörs av antalet parametrar och den numeriska precision som används vid lagring. Varje parameter representeras av ett flyttal eller ett kvantiserat heltal, där precisionen styr hur många byte som förbrukas per parameter:

• FP32 (32-bitars flyttal): 4 byte per parameter. En modell med 7 miljarder parametrar tar $7 \times 10^9 \times 4 \text{ byte} \approx 28 \text{ GB}$ VRAM.
• FP16 / BF16 (16-bitars flyttal): 2 byte per parameter. En modell med 7 miljarder parametrar tar $\approx 14 \text{ GB}$ VRAM.
• INT8 (8-bitars kvantisering): 1 byte per parameter. En 7B-modell tar $\approx 7 \text{ GB}$ VRAM.
• INT4 (4-bitars kvantisering, såsom AWQ eller GPTQ): 0,5 byte per parameter. En 7B-modell tar $\approx 3,5 \text{ GB}$ VRAM.

Dessa vikter måste laddas in fullständigt i VRAM för att grafikkortets beräkningskärnor ska kunna utföra matrismultiplikationer utan att flaskhalsas av PCIe-bussen.

KV-cache: det dynamiska minnet under sekvensgenerering

När en transformermodell genererar text token för token (autoregressiv avkodning) behöver den inte beräkna om alla tidigare tokens representationer. Istället sparas nyckel- och värdevektorerna (Keys och Values) från uppmärksamhetslagren i ett minnesutrymme som kallas KV-cache (Key-Value cache).

Detta gör att minnesbehov för transformermodell växer linjärt med kontextlängden och antalet samtidiga användare (batchstorlek). Formeln för att beräkna KV-cache per sekvens lyder:

$$\text{KV-cache (byte)} = 2 \times n_{\text{lager}} \times n_{\text{kv_huvuden}} \times d_{\text{huvud}} \times s \times b \times p$$

Där parametrarna definieras som:

• $2$: En faktor två för att både nycklar ($K$) och värden ($V$) sparas.
• $n_{\text{lager}}$: Antalet transformatorlager i modellen.
• $n_{\text{kv_huvuden}}$: Antalet uppmärksamhetshuvuden för nycklar och värden. I klassisk Multi-Head Attention (MHA) är detta lika med det totala antalet huvuden. I Grouped-Query Attention (GQA) delas huvuden i grupper vilket drastiskt minskar cachen.
• $d_{\text{huvud}}$: Dimensionen per uppmärksamhetshuvud (ofta dold dimension $d_{\text{modell}} / n_{\text{huvuden}}$).
• $s$: Sekvenslängd (antal tokens i prompt plus genererad text).
• $b$: Batchstorlek (antal samtidiga sekvenser).
• $p$: Precision i byte (vanligen 2 byte för FP16/BF16).

Räkneexempel för KV-cache i en 7B-modell

Låt oss räkna på en typisk 7B-modell med 32 lager, dold dimension 4096, 32 uppmärksamhetshuvuden och FP16-precision ($p=2$):

1. Per token och sekvens: $2 \times 32 \times 4096 \times 2 = 524{,}288 \text{ byte} \approx 0{,}5 \text{ MB}$.
2. Vid 4 096 tokens kontext: $4096 \times 0{,}5 \text{ MB} \approx 2 \text{ GB}$ per samtidig ström.
3. Vid 32 768 tokens (32k): $32768 \times 0{,}5 \text{ MB} \approx 16 \text{ GB}$ per ström.

Vid längre kontexter överstiger alltså KV-cachen själva vikternas storlek. Arkitekturer som använder Grouped-Query Attention (GQA), exempelvis Mistral och Llama, minskar detta minnesbehov avsevärt genom att använda färre KV-huvuden än frågehuvuden.

Aktiveringar och buffertminne vid beräkning

Utöver modellvikter och KV-cache kräver grafikkortet minne för aktiveringar (mellanresultat i matrismultiplikationer, Softmax och LayerNorm) samt minnesallokerarens interna fragmentering i CUDA.

Vid ren inferens sparas aktiveringar endast under den pågående beräkningen av ett lager och skrivs snabbt över, vilket gör att aktiveringsminnet sällan överskrider ett par hundra megabyte till någon gigabyte för mindre batchstorlekar. Det är dock god praxis att lägga till en säkerhetsmarginal på minst 15–20 procent utöver vikter och KV-cache för att undvika Out of Memory (OOM)-krascher orsakade av fragmentering.

För fördjupad information om maskininlärning och neurala nätverk finns översikter på Wikipedia.

Sammanställning: dimensionera rätt hårdvara

För att få fram det fullständiga minnesbehovet för en transformermodell summerar du:

1. Statisk viktlagring: Parametrar $\times$ byte per precision.
2. KV-cache: $2 \times \text{lager} \times \text{kv-dimension} \times \text{kontextlängd} \times \text{batch} \times \text{byte per precision}$.
3. Buffert för CUDA och aktiveringar: Cirka 20 procent extra marginal.

Med denna beräkningsmodell kan du med säkerhet välja rätt grafikkort och kvantisering för din servermiljö utan att drabbas av överraskande minnesbrist.