Un GPT desde cero + nanoGPT

Objetivo de maestría

ensamblar un Transformer GPT completo entendiendo cada componente, y entrenarlo a nivel de carácter. Tras esta lección, cuando leas el código de cualquier LLM moderno reconocerás todas las piezas.

2.1Anatomía de un bloque GPT

Un GPT es una pila de bloques transformer idénticos. Cada bloque tiene dos sub-capas, cada una con conexión residual y normalización:

x → LayerNorm → Multi-Head Attention → (+x)  → LayerNorm → MLP → (+x)

Las piezas y su porqué:

Embeddings de token + posicionales: convierten ids en vectores y añaden información de posición (la atención por sí sola no sabe el orden).
Multi-Head Attention (N5·L1): mezcla información entre posiciones.
MLP (feed-forward, normalmente 4× la dimensión): procesa cada posición independientemente; aporta la mayor parte de los parámetros y la "memoria".
Conexión residual (+x): deja que el gradiente fluya sin degradarse en redes profundas; sin ella, no entrenarías 12+ capas.
LayerNorm: estabiliza las activaciones. Pre-norm (norm antes de la sub-capa, como arriba) es el estándar moderno: entrena más estable que post-norm.

2.2Laboratorio L2.1 — GPT mínimo en PyTorch (código completo)

python

1# lab_n5l2_gpt.py — un GPT entrenable desde cero, estilo nanoGPT, comentado
2import torch, torch.nn as nn, torch.nn.functional as F, math
3
4class Block(nn.Module):
5    def __init__(self, d, n_heads, block_size):
6        super().__init__()
7        self.ln1 = nn.LayerNorm(d)
8        self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True)
9        self.ln2 = nn.LayerNorm(d)
10        self.mlp = nn.Sequential(
11            nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))  # MLP 4x con GELU
12        # máscara causal precomputada (un token no ve el futuro)
13        self.register_buffer("mask", torch.triu(
14            torch.ones(block_size, block_size), diagonal=1).bool())
15
16    def forward(self, x):
17        T = x.size(1)
18        a = self.ln1(x)
19        # attn_mask causal: True = posición bloqueada
20        a, _ = self.attn(a, a, a, attn_mask=self.mask[:T, :T], need_weights=False)
21        x = x + a                                    # residual 1
22        x = x + self.mlp(self.ln2(x))                # residual 2 (pre-norm)
23        return x
24
25class GPT(nn.Module):
26    def __init__(self, vocab, d=384, n_heads=6, n_layers=6, block_size=256):
27        super().__init__()
28        self.block_size = block_size
29        self.tok_emb = nn.Embedding(vocab, d)              # id -> vector
30        self.pos_emb = nn.Embedding(block_size, d)         # posición -> vector
31        self.blocks = nn.ModuleList([Block(d, n_heads, block_size) for _ in range(n_layers)])
32        self.ln_f = nn.LayerNorm(d)
33        self.head = nn.Linear(d, vocab, bias=False)        # proyecta a logits de vocabulario
34        self.head.weight = self.tok_emb.weight             # weight tying (comparte pesos)
35
36    def forward(self, idx, targets=None):
37        B, T = idx.shape
38        pos = torch.arange(T, device=idx.device)
39        x = self.tok_emb(idx) + self.pos_emb(pos)          # suma embeddings de token + posición
40        for blk in self.blocks: x = blk(x)
41        logits = self.head(self.ln_f(x))                   # [B, T, vocab]
42        if targets is None: return logits, None
43        loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
44        return logits, loss
45
46    @torch.no_grad()
47    def generate(self, idx, max_new_tokens, temperature=1.0):
48        for _ in range(max_new_tokens):
49            idx_cond = idx[:, -self.block_size:]           # recorta al contexto máximo
50            logits, _ = self(idx_cond)
51            logits = logits[:, -1, :] / temperature        # solo el último paso
52            probs = F.softmax(logits, dim=-1)
53            nxt = torch.multinomial(probs, 1)              # muestrea el siguiente token
54            idx = torch.cat([idx, nxt], dim=1)
55        return idx

Líneas no triviales explicadas:

register_buffer("mask", ...): la máscara causal no es un parámetro entrenable pero debe moverse con el modelo a la GPU; un buffer hace justo eso.
Pre-norm (x + mlp(ln2(x))): la norm va dentro del residual, no fuera. Es lo que estabiliza el entrenamiento profundo (GPT-2 lo usa; post-norm del transformer original era más inestable).
weight tying (head.weight = tok_emb.weight): comparte la matriz de embeddings de entrada con la de salida → menos parámetros y mejor generalización. Estándar desde GPT-2.
MLP 4×d con GELU: la expansión 4× es convención empírica; GELU es la activación estándar (más suave que ReLU).
generate: autoregresivo — predice un token, lo añade al contexto, repite. temperature controla la aleatoriedad (0 = greedy, alto = creativo).

2.2bis Laboratorio L2.2 — Entrenar a nivel de carácter (Shakespeare)

python

1# lab_n5l2_train.py — entrena el GPT a nivel de carácter en ~minutos en la 5090
2import torch
3text = open("input.txt").read()                  # descarga tinyshakespeare
4chars = sorted(set(text)); vocab = len(chars)
5stoi = {c:i for i,c in enumerate(chars)}; itos = {i:c for c,i in stoi.items()}
6data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
7n = int(0.9*len(data)); train, val = data[:n], data[n:]
8
9block_size, batch = 256, 64
10def get_batch(split):
11    d = train if split=="train" else val
12    ix = torch.randint(len(d)-block_size, (batch,))
13    x = torch.stack([d[i:i+block_size] for i in ix])
14    y = torch.stack([d[i+1:i+block_size+1] for i in ix])   # target = input desplazado 1
15    return x.cuda(), y.cuda()
16
17model = GPT(vocab).cuda()
18opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
19for step in range(5000):
20    x, y = get_batch("train")
21    _, loss = model(x, y)
22    opt.zero_grad(); loss.backward(); opt.step()
23    if step % 500 == 0:
24        model.eval()
25        with torch.no_grad(): _, vl = model(*get_batch("val"))
26        print(f"step {step}: train {loss.item():.3f}  val {vl.item():.3f}")
27        model.train()
28
29# genera texto
30ctx = torch.zeros((1,1), dtype=torch.long, device="cuda")
31print("".join(itos[i] for i in model.generate(ctx, 500)[0].tolist()))

Línea clave: y = input desplazado 1. El objetivo de cada posición es el siguiente token — esa es toda la tarea de pretraining (next-token prediction). Entender esto es entender qué aprende un LLM.

2.3De aquí a nanoGPT y nanochat

Tu GPT mínimo es nanoGPT conceptualmente. nanoGPT (karpathy/nanoGPT) lo lleva a producción (entrenamiento eficiente, checkpoints, FineWeb/OWT). nanochat (karpathy/nanochat) añade el pipeline completo de ChatGPT (tokenizer BPE → pretraining → midtraining → SFT → RL → inferencia → WebUI) en ~8.000 líneas, con un único "dial" de complejidad: la profundidad (depth) determina automáticamente anchura, nº de cabezas, learning rate y horizonte de entrenamiento. Estudiar nanochat entero es ver cómo todas las piezas del curso (Niveles 2 y 5) encajan en un sistema.

2.4Ejercicios

E1. Sustituye los embeddings posicionales aprendidos por RoPE (rotary). ¿Mejora la extrapolación a secuencias más largas que las de entrenamiento?

E2. Quita las conexiones residuales y entrena. Observa cómo la loss se estanca con más capas. (Demuestra por qué los residuales son imprescindibles.)

E3. Cambia pre-norm por post-norm. ¿Notas diferencia en estabilidad/velocidad de convergencia?

2.5Trampas comunes

Olvidar desplazar el target 1 posición → el modelo "ve" la respuesta, loss absurdamente baja.
Post-norm en redes profundas sin warmup → inestable.
No recortar al block_size en generate → error de tamaño de posición.

2.6Referencias

karpathy/nanoGPT, karpathy/nanochat, nn-zero-to-hero (lecture "Let's build GPT"). CS336 (architectures, hyperparameters).