El bloque Transformer y la clase GPT completa

Ya tienes las dos piezas con información: embeddings (qué token, en qué posición) y atención causal (mezclar información entre posiciones). Falta la maquinaria que hace que esto se pueda apilar decenas de veces y siga entrenando: conexiones residuales, normalización y una red feed-forward por posición. Al final de esta lección tendrás la clase GPT completa.

Anatomía de un bloque

Un bloque Transformer estilo GPT-2 hace dos cosas, en orden:

1. Comunicar: self-attention causal. Los tokens intercambian información.
2. Computar: una MLP aplicada a cada posición de forma independiente. Cada token "piensa" sobre lo que leyó.

Cada sub-capa se envuelve con una LayerNorm antes y una suma residual después. Esta variante se llama Pre-LN:

normalizar, pasar por la atención y sumarle el resultado a lo que ya había; después, lo mismo con el MLP

El Transformer original usaba Post-LN: $x' = \text{LN}(x + \text{Attn}(x))$. Parece un cambio cosmético, pero Xiong et al. (2020) mostraron que con Post-LN los gradientes de los parámetros cercanos a la salida son grandes al inicializar, lo que vuelve imprescindible un warmup cuidadoso de la tasa de aprendizaje; con Pre-LN los gradientes están mejor comportados. GPT-2 y la gran mayoría de los LLM actuales usan Pre-LN (o variantes como RMSNorm en la misma posición).

La vista del residual stream

Hay una forma muy útil de leer la arquitectura Pre-LN. Desenrolla las sumas: el vector de salida de la última capa es la entrada más todo lo que cada sub-capa le sumó:

La MLP

Expande a $4d$, aplica una no linealidad y vuelve a $d$:

agrandar el vector a 4 veces su tamaño, doblarlo con GELU y volver a achicarlo; $W_1$, $W_2$, $b_1$ y $b_2$ son números que el modelo aprende

GELU es una versión suave de ReLU que usa GPT-2. El factor 4 es una convención del paper original, no una ley. Aunque la atención se lleva la fama, en un bloque estándar dos tercios de los parámetros están en la MLP: $8d^2$ frente a $4d^2$ de la atención.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
# CausalSelfAttention y GPTConfig vienen de las lecciones anteriores


class MLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.fc = nn.Linear(cfg.n_embd, 4 * cfg.n_embd)
        self.proj = nn.Linear(4 * cfg.n_embd, cfg.n_embd)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x):
        return self.drop(self.proj(F.gelu(self.fc(x))))


class Block(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.n_embd)
        self.attn = CausalSelfAttention(cfg)
        self.ln2 = nn.LayerNorm(cfg.n_embd)
        self.mlp = MLP(cfg)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # comunicar
        x = x + self.mlp(self.ln2(x))    # computar
        return x

La clase GPT

Ahora ensamblamos todo: embeddings, $L$ bloques, LayerNorm final y la capa de salida atada a los embeddings.

class GPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_emb = nn.Embedding(cfg.vocab_size, cfg.n_embd)
        self.pos_emb = nn.Embedding(cfg.block_size, cfg.n_embd)
        self.drop = nn.Dropout(cfg.dropout)
        self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)])
        self.ln_f = nn.LayerNorm(cfg.n_embd)
        self.lm_head = nn.Linear(cfg.n_embd, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.tok_emb.weight          # weight tying

        self.apply(self._init_weights)
        # proyecciones que escriben en el residual stream: init más chica
        for name, p in self.named_parameters():
            if name.endswith("proj.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * cfg.n_layer))

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        B, T = idx.shape
        assert T <= self.cfg.block_size, f"secuencia de {T} > block_size"
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        for block in self.blocks:
            x = block(x)
        logits = self.lm_head(self.ln_f(x))                  # (B, T, V)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

El escalado $0.02/\sqrt{2L}$ en las proyecciones de salida (attn.proj y mlp.proj) viene del paper de GPT-2: cada bloque suma dos contribuciones al residual stream, y con $2L$ sumas la varianza crecería con la profundidad si no se compensara.

Contar parámetros

Por bloque, ignorando biases y LayerNorm: atención $4d^2$ (Q, K, V y salida) + MLP $8d^2$ = $12d^2$. Para todo el modelo:

parámetros totales: $12d^2$ por cada una de las $L$ capas, más la tabla de tokens ($V \cdot d$) y la de posiciones ($T_{max} \cdot d$)

Comprueba la fórmula con GPT-2 small ($L=12$, $d=768$, $V=50257$, $T_{max}=1024$): $12 \cdot 12 \cdot 768^2 \approx 84{,}9\text{M}$, más $38{,}6\text{M}$ de embeddings de tokens y $0{,}8\text{M}$ de posiciones: unos 124 M, el número oficial.

Con nuestra configuración de juguete:

cfg = GPTConfig(vocab_size=512, block_size=64, n_layer=4, n_head=4, n_embd=128)
model = GPT(cfg)
n_params = sum(p.numel() for p in model.parameters())
print(f"{n_params:,}")      # 867,072

idx = torch.randint(0, cfg.vocab_size, (2, 16))
logits, loss = model(idx, targets=idx)
print(logits.shape, loss.item())

model.parameters() no cuenta dos veces el tensor compartido entre tok_emb y lm_head, por eso el total ya refleja el weight tying.

Lo que te llevas

• Bloque = comunicar (atención) + computar (MLP), cada uno con Pre-LN y suma residual.
• El residual stream es un canal compartido al que cada sub-capa le suma su aporte.
• Un bloque tiene ~$12d^2$ parámetros, dos tercios en la MLP.
• Ya tienes un GPT que produce logits y loss. Próxima lección: entrenarlo.