Parametrar i en transformermodell: så räknar du

Antalet parametrar i en transformermodell får du genom att lägga ihop vikterna i embedding-lagret och i varje transformerlager. Ett transformerlager har fyra attention-matriser (query, key, value och output) som är lika stora som embedding-dimensionen i kvadrat, och ett feedforward-nätverk vars dolda lager är fyra gånger så brett som embedding-dimensionen. Multiplicerar du antalet parametrar med antalet bytes per parameter får du minnesbehovet för vikterna.

Så räknar du parametrar i ett transformerlager

Beteckna embedding-dimensionen med d. Varje attention-projektion (query, key, value och output) är en matris av storleken d × d, alltså d² parametrar. Fyra sådana ger 4d².

Feedforward-nätverket har två linjära lager. Enligt Wikipedia är det dolda lagrets storlek i både GPT-2 och BERT fyra gånger embedding-storleken. Det ger två matriser på d × 4d, alltså 8d² parametrar.

Tillsammans blir ett transformerlager ungefär 12d² parametrar. Det här räknar bort bias-termer och layer normalization, som är små jämfört med matriserna. Om din modell använder en annan storlek på det dolda lagret ändrar du faktorn 4 i beräkningen.

Embedding-lagret

Embedding-lagret översätter varje token-identifierare till en vektor. Det är en matris med en rad per token i vokabulären och en kolumn per dimension, alltså |V| × d parametrar. För en vokabulär på 50 000 tokens och d = 768 blir det 38,4 miljoner parametrar, bara i embedding-lagret.

Det utgångslager som gör om vektorerna till sannolikheter över vokabulären har samma form. Enligt Wikipedia använder vissa arkitekturer transponatet av embedding-matrisen som utgångsvikt. Då delar lagren vikter, och du ska inte räkna in båda. Kontrollera alltid i modellens konfiguration om vikterna är bundna.

Ett räkneexempel på en hel modell

Anta en modell med 12 lager, d = 768 och vokabulär 50 000, med delade vikter mellan in- och utgång.

• Embedding-lagret: 50 000 × 768 = 38 400 000
• Ett transformerlager: 12 × 768² = 12 × 589 824 = 7 077 888
• Tolv lager: 12 × 7 077 888 = 84 934 656
• Totalt: 38 400 000 + 84 934 656 = 123 334 656

Modellen har alltså ungefär 123 miljoner parametrar. Räkna med det här som en uppskattning. Den täcker inte bias-termer, positionsvikter om de finns separat, eller layer normalization.

Från parametrar till minne

Ett antal parametrar säger ingenting om minnet förrän du vet hur många bitar varje parameter tar. En half-precision-siffra (FP16) använder 16 bitar, alltså två bytes, enligt artikeln om half-precision-formatet på engelska Wikipedia. Bfloat16 använder också 16 bitar, men med 8 exponentbitar och 7 lagrade signifikantbitar, enligt Wikipedia-artikeln om bfloat16. Full precision (FP32) tar fyra bytes per parameter.

För modellen ovan ger det:

• FP32: 123 334 656 × 4 bytes, ungefär 493 MB
• FP16 eller bfloat16: 123 334 656 × 2 bytes, ungefär 247 MB

Det här är bara vikterna. Vid inferens tillkommer aktiveringar och KV-cache, och vid träning dessutom gradienter och optimerarens tillstånd. Det kan göra den totala minnesåtgången flera gånger större än vikterna.

Vanliga misstag

• Att glömma att räkna in utgångslagret när det inte delar vikter med embedding-lagret.
• Att blanda ihop MB och MiB. Här använder vi decimal MB (10⁶ bytes).
• Att räkna vikterna som om de vore i FP32 när modellen laddas i FP16.
• Att anta att alla transformrar har fyra gånger bredd i feedforward-lagret. Det är ett vanligt val, inte ett krav.
• Att blanda ihop antalet parametrar med antalet tokens som modellen kan läsa i ett svep. Kontextlängden påverkar minnet för KV-cachen, inte antalet vikter.

Vanliga frågor

Räknas bias-termer med?

Inte i den här förenklingen. De är små jämfört med matriserna. Om du vill ha exakt siffra, summera dem från modellens konfiguration.

Varför ska jag räkna i bytes?

För att se om modellen ryms i GPU-minnet. Parametrarna i bytes och minnet på kortet är de två siffror du behöver jämföra.

Vad gör jag om jag inte vet d?

Den står i modellens konfiguration eller modellkortet. Gissa inte. Utan d går det inte att räkna fram antalet parametrar.