Reproducir GPT-2 124M en tu 5090 (Checkpoint C5·a)

3.1El objetivo y por qué es alcanzable en local

El estándar moderno del "speedrun" de GPT-2: entrenar un modelo de 124M parámetros activos hasta ~3.28 de val loss en FineWeb (el target que iguala la reproducción de GPT-2 small de Karpathy). En 8×H100 con la receta de modded-nanogpt esto baja de los 45 min del baseline llm.c a ~4.4 minutos usando solo 0.75B tokens (en vez de 10B).

En una sola GPU de consumo, los forks single-GPU del speedrun (FineWeb-Edu 10BT, ~3B tokens) reproducen ~3.25 de loss en ~90 minutos en una 4090 — y tu 5090 es más rápida. Es decir: el pretraining de un GPT-2 entero cabe en una tarde en tu hardware. Esto era impensable hace dos años; las técnicas modernas lo hicieron posible.

Nota sobre targets: hay dos referencias comunes. El OWT clásico de nanoGPT (~2.85–2.9 val loss con el eval original de GPT-2) y el FineWeb 3.28 del speedrun moderno. Usa el del repo que sigas; documenta cuál.

3.2Las técnicas modernas que lo aceleran (entiéndelas, no solo las copies)

El speedrun no es "más GPU"; es mejor ML. Las piezas clave de modded-nanogpt y por qué funcionan:

• Optimizador Muon (Momentum Orthogonalized by Newton-Schulz): para las matrices de las capas ocultas, ortogonaliza la actualización del gradiente mediante una iteración de Newton-Schulz (aproxima reemplazar el gradiente G por U·Vᵀ de su SVD). Esto da un condicionamiento mucho mejor → converge en muchos menos pasos. Embeddings y lm_head siguen con AdamW (Muon es para matrices, no para embeddings).
• Arquitectura modernizada: RoPE (posiciones rotatorias), QK-Norm (normaliza Q y K antes de la atención → estabilidad), ReLU² en el MLP.
• Value residual y embedding shortcut: conexiones extra que mejoran el flujo de gradiente (capacidad sin FLOPs).
• Tanh logit capping (logits = 30·tanh(logits/30)): estabiliza el entrenamiento en BF16/FP8 evitando logits desbocados (siguiendo Gemma 2).
• FlexAttention con warmup de ventana: atención eficiente con ventana de contexto que crece durante el entrenamiento.
• Inicialización muP-like (capas de proyección/clasificación a cero): mejor escalado.

Cada una vale una entrada en tu lab notebook: ¿cuántos pasos/tiempo te ahorró? El speedrun es el mejor sitio para aprender optimización de entrenamiento moderna.

3.3Laboratorio L3.1 — Ejecutar el speedrun en tu 5090

1# Opción A: el repo canónico de Keller Jordan (pensado para 8xH100 pero estudiable y adaptable)
2git clone https://github.com/KellerJordan/modded-nanogpt && cd modded-nanogpt
3pip install -r requirements.txt
4python data/cached_fineweb10B.py 18        # descarga shards de FineWeb (tokenizados)
5# ajusta run.sh a 1 GPU (reduce batch global, sube grad_accum) y lánzalo
6
7# Opción B (recomendada para una 5090): un fork single-GPU del speedrun
8git clone https://github.com/Deveraux-Parker/nanoGPT_1GPU_SPEEDRUN && cd nanoGPT_1GPU_SPEEDRUN
9# Instala PyTorch nightly (FlexAttention/TorchAO) + numpy tqdm datasets tiktoken torchao
10python data_prep.py                        # tokeniza FineWeb-Edu sample-10BT (~3B tokens)
11python train_gpt2_single.py                # preconfigurado para 24GB; sube batch en tu 32GB

Qué observar durante el run (y entender):

• Pasos 0–300: la loss cae rápido (10.0 → ~4.5), throughput alto (~160k tok/s en una 4090; más en la 5090).
• "Lag spikes" en ~300 y ~600: la ventana de contexto se duplica → recompilación de FlexAttention. Normal.
• Fase de cooldown final: el learning rate baja y la loss cae de ~3.3 a ~3.25. Aquí se "asienta" el modelo.

Líneas no triviales:

• En una 5090 (32 GB vs 24 de la 4090) puedes subir el batch size → mejor utilización y posiblemente menos tiempo. Ajústalo hasta casi llenar la VRAM (N0·L3).
• PyTorch nightly: FlexAttention y TorchAO requieren versiones recientes; en Blackwell, asegúrate del índice cu128 (N0·L2).

3.4Laboratorio L3.2 — Evaluar y validar

1# lab_n5l3_eval.py — confirma que alcanzas el target y compara con GPT-2 oficial
2# 1) val loss: el script de entrenamiento la reporta; confirma que llegas a ~3.28 (FineWeb) / ~2.85 (OWT)
3# 2) HellaSwag / sanity: evalúa tu checkpoint en una tarea sencilla y compáralo con gpt2 de HF
4from transformers import GPT2LMHeadModel
5import torch
6# carga tu checkpoint entrenado y el GPT-2 oficial; compara perplexity en un holdout común
7# (tu modelo entrenado con receta moderna debería igualar o superar al gpt2 original de 2019)

El checkpoint no es solo "llegué a la loss": es entender por qué. ¿Qué pasó si subiste el LR? ¿Y si quitaste Muon (vuelve a AdamW puro)? Esos experimentos son el aprendizaje.

3.5nanochat: el pipeline completo (capstone del nivel)

Una vez reproduces el pretraining, nanochat te lleva de un GPT a un "ChatGPT" entero: tokenizer BPE propio → pretraining → midtraining → SFT (Nivel 2) → RL (Nivel 2) → inferencia → WebUI, en ~8.000 líneas. El depth único controla todo. Correrlo (o las partes que quepan) cierra el círculo del curso: ves SFT y RL (que hiciste con modelos ajenos en el Nivel 2) aplicados sobre un modelo que tú pre-entrenaste. Es la mejor preparación para el capstone del programa maestro.

3.6CHECKPOINT C5(a) — criterio de aprobado

• Reproduces GPT-2 124M desde cero alcanzando la val loss objetivo (~3.28 FineWeb o ~2.85 OWT), reproducible.
• Un writeup que justifica cada decisión: por qué Muon para matrices y AdamW para embeddings, qué hace el warmup+cooldown del LR, por qué QK-Norm/RoPE, qué aporta el value residual.
• Has hecho al menos un experimento de ablación (p.ej. Muon vs AdamW puro) y mides el efecto.

Rúbrica: Nivel 3 si reproduces y justificas; Nivel 4 si mejoras algo del speedrun (un truco de optimización que reduzca pasos) o escalas la receta a un tamaño mayor en cloud.

3.7Ejercicios

E1. Entrena con Muon y con AdamW puro (mismo todo lo demás). ¿Cuántos pasos/tiempo te ahorra Muon para llegar al target? Conecta con 3.2.

E2. Sube el batch size aprovechando los 32 GB de la 5090. ¿Mejora el tiempo total? ¿Cambia la loss final?

E3. Quita el tanh logit capping y entrena en BF16. ¿Aparece inestabilidad? ¿Por qué el capping ayuda?

3.8Trampas comunes

• Usar Muon en embeddings/lm_head (no es para eso) → peor o inestable.
• No ajustar el batch a tu VRAma (32 GB) → infrautilizas la 5090.
• Comparar tu loss en FineWeb con un target de OWT (son escalas distintas).
• Reportar "llegué a la loss" sin entender por qué → Nivel 2, no 3.

3.9Referencias

• KellerJordan/modded-nanogpt (Muon, arquitectura, world records). Forks single-GPU (Deveraux-Parker, alexjc, sam-schorb). karpathy/nanoGPT, nanochat, llm.c. CS336 (scaling laws, training). Muon: Jeremy Bernstein / Keller Jordan writeups.