Entrenar el modelo: datos, loss, AdamW y schedules

Tienes un GPT que produce logits y una loss. Ahora hay que hacer que esa loss baje. Esta lección cubre el pipeline completo: cómo se arman los batches, qué minimiza exactamente la cross-entropy, por qué AdamW, cómo se programa el learning rate y cómo leer lo que pasa durante el entrenamiento.

Datos: una cinta larga de tokens

El preprocesamiento es simple: tokenizas todo el corpus una vez y lo guardas como un único tensor 1D. Para formar un ejemplo, tomas una ventana de $T+1$ tokens a partir de una posición aleatoria; la entrada son los primeros $T$ y el objetivo, los mismos desplazados en uno.

import torch

# `tok` es el BPETokenizer de la lección 1
text = open("input.txt", encoding="utf-8").read()
data = torch.tensor(tok.encode(text), dtype=torch.long)

n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

def get_batch(split, batch_size, block_size, device):
    src = train_data if split == "train" else val_data
    ix = torch.randint(len(src) - block_size - 1, (batch_size,))
    x = torch.stack([src[i : i + block_size] for i in ix])
    y = torch.stack([src[i + 1 : i + 1 + block_size] for i in ix])
    return x.to(device), y.to(device)

Dos decisiones a notar. Primero, la separación train/validación es contigua (el último 10 % del texto), no aleatoria por ventana: las ventanas superpuestas harían que la validación vea casi el mismo texto que el entrenamiento. Segundo, muestrear posiciones al azar equivale a recorrer el corpus en desorden, sin el concepto rígido de "época".

Qué minimiza la cross-entropy

Para cada posición, el modelo produce una distribución $p_\theta(\cdot \mid x_{<t})$ sobre el vocabulario. La loss es la log-verosimilitud negativa promedio del token correcto:

promediado sobre $B$ secuencias de $T$ tokens: qué probabilidad le dio el modelo al token que de verdad venía después, en logaritmo y con signo menos. Cuanto más baja, mejor predice

Una forma de interpretarla: $e^{\mathcal{L}}$ es la perplejidad, el número "efectivo" de tokens entre los que el modelo duda en promedio. Con $V = 512$ y predicciones uniformes, la perplejidad es exactamente 512; el entrenamiento consiste en reducir ese número. Nunca llega a 1, porque el lenguaje tiene incertidumbre genuina: muchos siguientes tokens son legítimos.

El optimizador: AdamW

Adam mantiene, para cada parámetro, promedios móviles del gradiente ($m$) y de su cuadrado ($v$), y da pasos normalizados por $\sqrt{v}$: los parámetros con gradientes ruidosos o grandes avanzan con más cautela. AdamW (Loshchilov y Hutter) desacopla el weight decay del gradiente: en vez de sumar $\lambda\theta$ al gradiente (donde Adam lo normalizaría), lo aplica directamente sobre los pesos:

el paso de cada peso: el gradiente promediado ($\hat m_t$) dividido por su tamaño típico ($\sqrt{\hat v_t}$), más un tirón chico hacia cero ($\lambda\theta_t$); $\eta_t$ marca qué tan largo es el paso

Convención habitual: aplicar weight decay solo a las matrices (tensores 2D: lineales y embeddings) y no a biases ni a los parámetros de LayerNorm, que son vectores 1D.

def configure_optimizer(model, lr, weight_decay=0.1):
    decay = [p for p in model.parameters() if p.requires_grad and p.dim() >= 2]
    no_decay = [p for p in model.parameters() if p.requires_grad and p.dim() < 2]
    groups = [
        {"params": decay, "weight_decay": weight_decay},
        {"params": no_decay, "weight_decay": 0.0},
    ]
    return torch.optim.AdamW(groups, lr=lr, betas=(0.9, 0.95))

model.parameters() devuelve una sola vez el tensor atado entre tok_emb y lm_head, así que no queda duplicado en los grupos.

Schedule: warmup + cosine decay

Al inicio, las estimaciones de $v$ de Adam son poco fiables y los gradientes pueden ser grandes; arrancar con el learning rate máximo puede desestabilizar el entrenamiento. Se usa un warmup lineal y luego un decaimiento en forma de coseno hasta un mínimo:

el largo del paso: sube en línea recta durante los primeros $t_w$ pasos y después baja suave, con forma de coseno, hasta $\eta_{min}$

import math

def get_lr(step, max_lr, min_lr, warmup, total):
    if step < warmup:
        return max_lr * (step + 1) / warmup
    progress = min(1.0, (step - warmup) / max(1, total - warmup))
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))

El loop completo

from contextlib import nullcontext

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(1337)

cfg = GPTConfig(vocab_size=512, block_size=64, n_layer=4, n_head=4, n_embd=128, dropout=0.1)
model = GPT(cfg).to(device)
opt = configure_optimizer(model, lr=3e-4)

max_steps, warmup, batch_size = 3000, 100, 32

@torch.no_grad()
def estimate_loss(iters=50):
    model.eval()
    out = {}
    for split in ("train", "val"):
        losses = torch.zeros(iters)
        for k in range(iters):
            x, y = get_batch(split, batch_size, cfg.block_size, device)
            _, loss = model(x, y)
            losses[k] = loss.item()
        out[split] = losses.mean().item()
    model.train()
    return out

for step in range(max_steps):
    lr = get_lr(step, max_lr=3e-4, min_lr=3e-5, warmup=warmup, total=max_steps)
    for g in opt.param_groups:
        g["lr"] = lr

    x, y = get_batch("train", batch_size, cfg.block_size, device)
    _, loss = model(x, y)

    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    opt.step()

    if step % 250 == 0 or step == max_steps - 1:
        l = estimate_loss()
        print(f"step {step:5d} | lr {lr:.2e} | train {l['train']:.3f} | val {l['val']:.3f}")

torch.save(model.state_dict(), "mini_gpt.pt")

estimate_loss promedia varios batches en modo eval() (sin dropout): la loss de un solo batch es demasiado ruidosa para tomar decisiones.

Sanity check: overfittear un batch

Antes de lanzar un entrenamiento largo, haz esta prueba de 30 segundos: toma un solo batch y entrena sobre él repetidamente, sin dropout. Un modelo correcto debe memorizarlo y llevar la loss cerca de cero.

model = GPT(GPTConfig(dropout=0.0)).to(device)
opt = configure_optimizer(model, lr=1e-3, weight_decay=0.0)
x, y = get_batch("train", 8, 64, device)
for i in range(300):
    _, loss = model(x, y)
    opt.zero_grad(set_to_none=True)
    loss.backward()
    opt.step()
    if i % 50 == 0:
        print(i, round(loss.item(), 4))

Si la loss no baja, el bug está en el modelo o en los datos (objetivos desalineados, máscara rota, gradientes que no llegan), no en los hiperparámetros. Y si baja "demasiado rápido" en el entrenamiento real, sospecha de una fuga: vuelve al test de la máscara causal.

Leer las curvas de loss

Con un corpus pequeño (un par de libros) y este modelo de menos de un millón de parámetros, es esperable ver la brecha entre train y val abrirse con relativa rapidez. Ese es justamente el experimento: prueba dropout=0.0 contra dropout=0.2 y compara las curvas de validación.

Lo que te llevas

• Un ejemplo es una ventana de $T+1$ tokens: entrada y objetivo desplazado en uno.
• Cross-entropy = log-verosimilitud negativa; $e^{\mathcal{L}}$ es la perplejidad.
• AdamW + weight decay solo en matrices + warmup + cosine + clipping: la receta estándar.
• Antes de entrenar en serio, overfittea un batch. Después, mira siempre la validación.