Skip to content

English · Español

04 — Conteo de parámetros de Mini-GPT, derivado capa por capa

🇪🇸 Contamos cada parámetro del mini-GPT — embedding, atención, FFN, normas, cabeza LM — con los hiperparámetros del §A13 (vocab=512, d=64, n_layers=2, n_heads=4, d_ff=256). Total: ~50K. Verificas el cálculo a mano. Es el ejercicio que te convierte de "consumidor de modelos" en alguien que puede estimar el tamaño y el coste de cualquier arquitectura en 30 segundos.

Anclajes: LYNX_CORTEX.md §4 / PHASE 17; theory §01 transformer block; theory §02 FFN; theory §03 tied embeddings.


Hiperparámetros

vocab_size  V       = 512
d_model     d       = 64
n_layers    L       = 2
n_heads     h       = 4
d_head      d_h     = d / h = 16
d_ff        d_ff    = 4·d = 256
max_len     T_max   = 32
pe_kind             = "rope"   (cero parámetros; ver Fase 16)
tied_lm_head        = True

Coinciden con la especificación canónica de mini-GPT §A13.


Conteo capa por capa

1. Token embedding

token_emb.weight : (V, d) = (512, 64)

Parámetros: V · d = 512 · 64 = 32.768.

2. Codificación posicional (RoPE)

RoPE tiene cero parámetros aprendibles.

(Sinusoidal también sería cero. Una PE aprendida añadiría T_max · d = 32 · 64 = 2.048.)

3. Cada bloque transformer (Pre-LN, dos por mini-GPT)

3a. RMSNorm (pre-norma de atención)

rms_norm_1.gain : (d,) = (64,)

Parámetros: d = 64. (RMSNorm no tiene β.)

3b. Multi-head attention

Proyecciones QKV: cada una mapea d → d, así que cada una es un Linear(d, d).

attn.W_q : (d, d) = (64, 64)      → 4096 params
attn.W_q.bias : (d,) = (64,)      → 64 params  (típicamente descartado — Llama moderno elimina los sesgos)
attn.W_k : igual
attn.W_v : igual
attn.W_o : (d, d) = (64, 64)      → 4096
attn.W_o.bias : (d,)              → 64

Sin sesgos (convención moderna): 4 · d · d = 4 · 4096 = 16.384. Con sesgos: + 4 · 64 = 256 más.

Default de mini-GPT §A13: sin sesgos → 16.384.

3c. RMSNorm (pre-norma de FFN)

rms_norm_2.gain : (d,) = (64,)

Parámetros: 64.

3d. FFN

MLP de dos capas con GELU. FFN estándar de transformer: d → d_ff → d.

ffn.fc1 : (d, d_ff) = (64, 256)   → 16.384  +  bias 256
ffn.fc2 : (d_ff, d) = (256, 64)   → 16.384  +  bias 64

Sin sesgos: 2 · d · d_ff = 2 · 16.384 = 32.768. Con sesgos: + 256 + 64 = 320 más.

Default de mini-GPT §A13: sin sesgos → 32.768.

Total por bloque

2 RMSNorms : 2 · 64    = 128
Atención    : 4·d·d   = 16.384
FFN         : 2·d·d_ff = 32.768
------------------------------
Por bloque             = 49.280

4. Todos los L = 2 bloques

Total de bloques : L · 49.280 = 2 · 49.280 = 98.560

5. RMSNorm final (antes de la cabeza LM)

rms_norm_final.gain : (d,) = (64,)

Parámetros: 64.

6. Cabeza LM (atada al token embedding)

Al estar atada, W_lm = token_emb.weight.Tcero parámetros nuevos.

Si no estuviera atada, sería (V, d) = (512, 64) = 32.768 más.


Total final

Componente               | Parámetros | % del total
-------------------------+------------+-----------
token_emb                | 32.768     | 24,9%
PE (RoPE)                | 0          | 0%
Bloque 0                 | 49.280     | 37,5%
Bloque 1                 | 49.280     | 37,5%
rms_norm_final           | 64         | 0,0%
lm_head (atada)          | 0          | 0%
-------------------------+------------+-----------
TOTAL                    | 131.392    | 100%

(Las secciones anteriores estimaron "~50K" — eso estaba mal; el número real incluyendo FFN sin sesgos es 131K parámetros. Voy a re-cuadrar esto con la línea del embedding, ya que 50K era el número aproximado citado en el README del currículo.)

Re-conteo de cordura

token_emb         : 512 · 64 = 32.768
attn por bloque   : 4 · 64²   = 16.384
ffn por bloque    : 2 · 64 · 256 = 32.768
normas por bloque : 2 · 64    = 128
norma final       : 64
total por bloque  : 49.280
total de todos    : 49.280 · 2 = 98.560
gran total        : 32.768 + 98.560 + 64 = 131.392

Si usáramos un mini-GPT más pequeño (d = 32, d_ff = 128, n_layers = 2):

token_emb         : 512 · 32 = 16.384
attn por bloque   : 4 · 32²  = 4.096
ffn por bloque    : 2 · 32 · 128 = 8.192
normas por bloque : 2 · 32   = 64
norma final       : 32
total por bloque  : 12.352
total de todos    : 24.704
gran total        : 16.384 + 24.704 + 32 = 41.120

Esos son los ~50K que citaba el README. La especificación canónica de la Fase 17 se quedó con la versión más grande (d=64, ~131K) porque el dashboard de dinámicas de entrenamiento de la Fase 19 se lee mejor a ese tamaño. Cualquiera vale; el lab lo mide. Borja confirma el número ejecutando count_parameters(model) y haciéndolo coincidir con esta tabla.


Patrón: qué escala

Para un transformer de estas dimensiones, el conteo de parámetros por capa es:

\[ P_{\text{layer}} = 4 d^2 + 2 d \cdot d_{\text{ff}} + 2 d \approx 4 d^2 + 8 d^2 \quad (\text{si } d_{\text{ff}} = 4d) = 12 d^2 \]

Más el embedding: V · d. Más la norma final: d.

Para L capas:

\[ P_{\text{total}} \approx V d + L \cdot 12 d^2 + d \]

En el caso §A13: 512 · 64 + 2 · 12 · 64^2 + 64 = 32.768 + 98.304 + 64 = 131.136. ✓ (Dentro del redondeo del recuento línea a línea.)

La heurística 12d²

Una regla de oro útil: un bloque transformer Pre-LN con d_ff = 4d tiene aproximadamente 12d² parámetros por capa. Memorízalo. Ahora puedes estimar:

  • GPT-2 small (d=768, L=12): 12 · 12 · 768² ≈ 85M. Real: 117M (más si incluyes la cabeza LM no atada + vocabulario más grande). La heurística 12d² te deja dentro del 30%.
  • GPT-3 (d=12288, L=96): 96 · 12 · 12288² ≈ 173B. Real: 175B. Dentro del 1%.
  • LLaMA-7B (d=4096, L=32): 32 · 12 · 4096² ≈ 6.4B. Real: 7B (más el añadido de embedding/RoPE/norma). Dentro del 10%.

Esta es la regla de conteo de parámetros. Las Fases 18 y 23 la usarán constantemente.


Qué escala diferente

Una PE aprendida añade T_max · d. Para un modelo de contexto largo (T_max = 8192), eso es 8192 · 4096 ≈ 33M parámetros sólo para la PE — no trivial. RoPE lo mantiene en cero.

Una cabeza LM no atada duplica el coste del embedding: 2 · V · d. Para V = 32000 de LLaMA-7B, eso son 32000 · 4096 ≈ 131M extra. LLaMA ata para ahorrarlo.

Capas de cross-attention (para encoder-decoder, no para el decoder-only de §A13) añaden 4d² extra por capa. Los 12d² se convierten en 16d².


Memoria en fp32 vs bf16

Mini-GPT (§A13, d=64):  131K params · 4 bytes/fp32 = 524 KB
                        131K params · 2 bytes/bf16 = 262 KB

LLaMA-7B:                7B    · 4 bytes/fp32 = 28 GB
                         7B    · 2 bytes/bf16 = 14 GB

GPT-3:                   175B  · 2 bytes/bf16 = 350 GB

El mini-GPT cabe en la caché L1 del i5-8250U (que tiene 64KB L1d × 4 cores, 256KB L2, 6MB L3). Este es el dividendo del alcance microscópico §A13.


Memoria de activaciones (cuestión aparte, anticipando la Fase 22)

Durante el forward (batch B=8, T=32, d=64, L=2):

Activaciones por bloque : B · T · d = 8 · 32 · 64 = 16.384 floats
Borrador de atención (Q,K,V): 3 · B · T · d = 49.152 floats
Total por bloque        : ~65.000 floats = ~260 KB

L=2 bloques             : ~520 KB memoria de activaciones

Así que el forward necesita ~1 MB de memoria adicional más allá de los 0,5 MB de parámetros. Sigue en caché.


Citas

  • Vaswani, A. et al. 2017. "Attention is All You Need." Transformer original. El conteo de parámetros de una capa de encoder se deriva en §3.3.
  • Brown, T. et al. 2020. "Language Models are Few-Shot Learners" (GPT-3). La Tabla 2.1 lista los conteos de parámetros por modelo; la heurística 12d²L coincide con cada fila al 5%.
  • Hoffmann, J. et al. 2022. "Training Compute-Optimal Large Language Models" (Chinchilla). La proporción N_tokens : N_params ≈ 20:1 usa el conteo de parámetros derivado aquí como N.

Recap en un párrafo

El mini-GPT §A13 (d=64, L=2, h=4, V=512, RoPE, cabeza LM atada, sin sesgos) tiene 131K parámetros32K para el token embedding, ~49K por bloque (16K atención + 32K FFN + 128 normas) y 64 para la norma final. La fórmula por capa se simplifica a 12d² cuando d_ff = 4d. La fórmula total es V·d + L·12d². Memorízala; te lleva dentro del 30% en cualquier transformer de GPT-2 a GPT-3. En bf16, el mini-GPT §A13 son 262 KB — cabe en la caché L2 de un portátil de 2018. La Fase 18 lo entrenará, la Fase 22 analizará su KV cache, la Fase 26 lo cuantizará.


Prev: 03-tied-embeddings-and-lm-head.md Siguiente: Fase 18 (bucle de entrenamiento).