¿Qué es el NanoGPT Speedrun?

El NanoGPT Speedrun de Andrej Karpathy es un esfuerzo comunitario para reproducir GPT-2 (124M parámetros) en OpenWebText en el menor tiempo de reloj posible, midiendo únicamente mejoras independientes del hardware — ganancias algorítmicas, no GPUs más rápidas.

💡 La métrica objetivo es la pérdida de validación en OpenWebText que iguale el checkpoint original de GPT-2 (124M): ~3,28 bits por carácter. Cada entrada en el speedrun debe entrenar desde cero.

El baseline es una reimplementación fiel de GPT-2 en PyTorch, entrenada con AdamW y un schedule de learning rate coseno. Sirve como punto de referencia contra el cual se mide cada mejora posterior.

Arquitectura

GPT-2 (124M) es un Transformer solo-decoder con los siguientes hiperparámetros:

  • Capas: 12 bloques Transformer
  • Cabezas: 12 cabezas de atención
  • Dimensión de embedding: 768
  • Longitud de contexto: 1024 tokens
  • Tamaño del vocabulario: 50.257 (tokenizador BPE de GPT-2)

Cada bloque aplica pre-layer-norm, atención causal multi-cabeza y un MLP con activaciones GELU. Las matrices de embedding y unembedding comparten pesos (weight tying).

El Mecanismo de Atención

La atención causal impide que cada token atienda a posiciones futuras. Para una secuencia de longitud $T$ y dimensión de embedding $d$, la atención escalada por producto punto es:

$$\text{Attn}(Q, K, V) = \text{softmax}\!\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

donde la máscara asegura que $\text{Attn}_{ij} = 0$ para todo $j > i$.

Configuración del Entrenamiento

El bucle de entrenamiento del baseline es directo: 

import torch
from torch.optim import AdamW

# Cosine LR schedule with linear warmup
def get_lr(step, warmup_steps, max_steps, max_lr, min_lr):
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    if step > max_steps:
        return min_lr
    decay = (step - warmup_steps) / (max_steps - warmup_steps)
    coeff = 0.5 * (1.0 + math.cos(math.pi * decay))
    return min_lr + coeff * (max_lr - min_lr)

optimizer = AdamW(model.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=0.1)

for step in range(max_steps):
    lr = get_lr(step, warmup_steps=715, max_steps=19073, max_lr=6e-4, min_lr=6e-5)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    loss = model(x, targets=y)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

Detalles clave del entrenamiento: batch size de 524.288 tokens (acumulación de gradientes), gradient clipping en 1,0 y weight decay de 0,1 aplicado a todos los parámetros 2D.

Resultados del Baseline

En un solo nodo con 8×H100, el baseline alcanza la pérdida de validación objetivo en aproximadamente ~1 hora . Las entradas posteriores del speedrun buscan igualar la misma pérdida en menos tiempo mediante mejoras puramente algorítmicas.

¿Qué parámetro controla cuán lejos puede 'mirar hacia atrás' cada token durante la atención en GPT-2?