microgpt - entrenamiento e inferencia de GPT implementados en 200 líneas de Python puro

• Proyecto artístico publicado por Karpathy. Implementa todo el algoritmo de GPT en un solo archivo de 200 líneas, sin dependencias externas
• La diferencia con los LLM de producción es solo de escala y eficiencia; el núcleo es el mismo, y entender este código equivale a entender la esencia algorítmica de GPT
• Incluye dataset, tokenizador, motor de autograd, arquitectura Transformer similar a GPT-2, optimizador Adam y hasta el bucle de entrenamiento e inferencia
• Como resultado final de 10 años de trabajo de simplificación de LLM, integrando proyectos previos como micrograd, makemore y nanogpt, condensa la esencia de GPT en su forma mínima, imposible de simplificar más
• Entrena con 32,000 datos de nombres para generar nombres nuevos verosímiles, realizando todos los cálculos directamente con autograd a nivel escalar
• El proceso de entrenamiento se compone de cálculo de pérdida → backpropagation → actualización con Adam, y puede ejecutarse en alrededor de 1 minuto

---

Resumen de microgpt

• microgpt es un script de Python de 200 líneas que implementa de forma completa el proceso de entrenamiento e inferencia de un modelo GPT
  • Sin librerías externas, incluye dataset, tokenizador, autograd, modelo, optimizador y bucle de entrenamiento
• Reúne en un solo archivo proyectos anteriores como micrograd, makemore y nanogpt
• Es una implementación que deja solo el núcleo algorítmico, al nivel de “ya no se puede simplificar más”
• El código completo está disponible en GitHub Gist, página web y Google Colab

Composición del dataset

• El combustible de los modelos de lenguaje a gran escala es un flujo de datos de texto; en producción se usan páginas web de internet, pero en microgpt se usa un ejemplo simple con 32,000 nombres puestos uno por línea
• Cada nombre se trata como un “documento”, y el objetivo del modelo es aprender los patrones estadísticos de los datos para generar documentos nuevos similares
• Tras completar el entrenamiento, el modelo “alucina” nombres nuevos plausibles como "kamon", "karai" y "vialan"
• Desde la perspectiva de ChatGPT, una conversación con el usuario también es solo un “documento con una forma peculiar”; si el documento se inicializa con un prompt, la respuesta del modelo corresponde a una completación estadística de documento

Tokenizador

• Como la red neuronal no opera con caracteres sino con números, se necesita una forma de convertir el texto en una secuencia de IDs de tokens enteros y luego reconstruirlo
• Los tokenizadores de producción como tiktoken (usado por GPT-4) trabajan por fragmentos de caracteres para mayor eficiencia, pero el tokenizador más simple asigna un entero a cada carácter único del dataset
• Se ordenan las letras minúsculas de la a a la z y se asigna a cada carácter un ID como índice; el valor entero en sí no tiene significado y cada token es un símbolo discreto independiente
• Se agrega un token especial BOS (Beginning of Sequence) para indicar que “empieza/termina un documento nuevo”, y "emma" se envuelve como [BOS, e, m, m, a, BOS]
• El vocabulario final tiene 27 elementos (26 letras minúsculas + 1 BOS)

Diferenciación automática (Autograd)

• Para entrenar una red neuronal se necesitan gradientes: para cada parámetro hay que saber “si subo un poco este valor, ¿la pérdida sube o baja, y cuánto?”
• El grafo computacional tiene muchas entradas (parámetros del modelo y tokens de entrada), pero converge en una única salida escalar: la pérdida (loss)
• La backpropagation comienza en la salida y recorre el grafo en sentido inverso, apoyándose en la regla de la cadena del cálculo para obtener el gradiente de la pérdida respecto de todas las entradas
• Se implementa con la clase Value: cada Value envuelve un solo escalar (.data) y rastrea cómo fue calculado
  • Al hacer operaciones como suma o multiplicación, el nuevo Value recuerda las entradas (_children) y las derivadas locales de esa operación (_local_grads)
  • Ejemplo: __mul__ registra ∂(a·b)/∂a=b y ∂(a·b)/∂b=a
• Bloques de operación soportados: suma, multiplicación, potencia, log, exp y ReLU
• El método backward() recorre el grafo en orden topológico inverso y aplica la regla de la cadena en cada paso
  • Comienza en el nodo de pérdida con self.grad = 1 (∂L/∂L=1)
  • Multiplica los gradientes locales a lo largo de cada trayectoria y los propaga hasta los parámetros
• Se acumula con += (no se asigna): cuando el grafo se bifurca, los gradientes fluyen de manera independiente por cada rama y luego deben sumarse (resultado de la regla de la cadena multivariable)
• Es algorítmicamente idéntico a .backward() de PyTorch, pero en vez de tensores opera a nivel escalar, por lo que es mucho más simple aunque menos eficiente

Inicialización de parámetros

• Los parámetros son el conocimiento del modelo: un gran conjunto de números de punto flotante que comienzan aleatoriamente y se optimizan de forma iterativa durante el entrenamiento
• Se inicializan con valores aleatorios pequeños tomados de una distribución gaussiana
• Están organizados como matrices nombradas en state_dict: tabla de embeddings, pesos de atención, pesos del MLP y proyección final de salida
• Configuración de hiperparámetros:
  • n_embd = 16: dimensión del embedding
  • n_head = 4: número de cabezas de atención
  • n_layer = 1: número de capas
  • block_size = 16: longitud máxima de secuencia
• En este modelo pequeño hay 4,192 parámetros (GPT-2 tiene 1.6 mil millones, y los LLM modernos tienen cientos de miles de millones)

Arquitectura

• La arquitectura del modelo es una función sin estado: recibe tokens, posiciones, parámetros y claves/valores en caché de posiciones anteriores, y devuelve los logits (puntajes) del siguiente token
• Sigue a GPT-2, pero con una ligera simplificación: RMSNorm (en lugar de LayerNorm), sin sesgos, ReLU (en lugar de GeLU)

• Funciones auxiliares

  • linear: mediante multiplicación matriz-vector calcula un producto punto por cada fila de la matriz de pesos; es la transformación lineal aprendida que sirve como bloque básico de una red neuronal
  • softmax: convierte puntajes crudos (logits) en una distribución de probabilidad; todos los valores quedan en el rango [0,1] y suman 1; para estabilidad numérica primero se resta el valor máximo
  • rmsnorm: reescala el vector para que tenga una raíz de la media cuadrática unitaria, evitando que las activaciones crezcan o se reduzcan al pasar por la red y estabilizando el entrenamiento

• Estructura del modelo

  • Embeddings: el ID del token y el ID de posición consultan filas en sus respectivas tablas de embeddings (wte, wpe); luego se suman ambos vectores para codificar al mismo tiempo qué es el token y dónde está en la secuencia
    • Los LLM modernos suelen omitir los embeddings de posición y usar técnicas de posicionamiento relativo como RoPE
  • Bloque de atención: proyecta el token actual en tres vectores: Q (query), K (key) y V (value)
    • Query: "¿qué estoy buscando?", key: "¿qué contengo?", value: "¿qué entrego si soy seleccionado?"
    • Ejemplo: al predecir lo siguiente después de la segunda "m" en "emma", puede aprender una query como "¿qué vocal apareció recientemente?"; la "e" anterior coincide bien con esa query y obtiene un peso de atención alto
    • Las keys y values se agregan a la KV cache, lo que permite referenciar posiciones anteriores
    • Cada cabeza de atención calcula el producto punto entre la query y todas las keys en caché (escalado por √d_head), obtiene pesos de atención con softmax y calcula una suma ponderada de los values en caché
    • Las salidas de todas las cabezas se concatenan y luego se proyectan con attn_wo
    • El bloque de atención es el único lugar donde el token en la posición t puede "ver" los tokens pasados 0..t-1; la atención es el mecanismo de comunicación entre tokens
  • Bloque MLP: una red feedforward de 2 capas: expande a 4 veces la dimensión del embedding → aplica ReLU → vuelve a reducir
    • Aquí ocurre la mayor parte del "razonamiento" por posición
    • A diferencia de la atención, en el tiempo t es un cálculo completamente local
    • El Transformer alterna comunicación (atención) y cómputo (MLP)
  • Conexiones residuales: tanto el bloque de atención como el bloque MLP vuelven a sumar su salida a la entrada
    • Permiten que los gradientes atraviesen la red directamente y hacen entrenables los modelos profundos
  • Salida: el estado oculto final se proyecta con lm_head al tamaño del vocabulario para generar un logit por token (aquí, 27 números); un logit alto = mayor probabilidad de que ese token venga después
  • Particularidad de la KV cache: aunque durante el entrenamiento es poco común usar KV cache, como microgpt procesa un solo token a la vez, se construye de forma explícita; las keys y values en caché son nodos Value activos del grafo computacional y participan en la retropropagación

Bucle de entrenamiento

• El bucle de entrenamiento repite: (1) seleccionar un documento → (2) ejecutar la pasada hacia adelante del modelo sobre los tokens → (3) calcular la pérdida → (4) obtener gradientes por retropropagación → (5) actualizar los parámetros

• Tokenización

  • En cada paso de entrenamiento se elige un documento y se envuelve con BOS en ambos extremos: "emma" → [BOS, e, m, m, a, BOS]
  • El objetivo del modelo es predecir cada siguiente token dados los tokens anteriores

• Forward pass y pérdida

  • Los tokens se alimentan al modelo uno por uno mientras se construye la KV cache
  • En cada posición, el modelo produce 27 logits, que se convierten en probabilidades con softmax
  • La pérdida en cada posición es la log-probabilidad negativa del siguiente token correcto: −log p(target); a esto se le llama pérdida de entropía cruzada
  • La pérdida mide qué tan sorprendido está el modelo por lo que realmente viene: si asigna probabilidad 1.0, la pérdida es 0; si asigna una probabilidad cercana a 0, la pérdida es +∞
  • Se promedian las pérdidas de todas las posiciones del documento para obtener una única pérdida escalar

• Backward pass

  • Con una sola llamada a loss.backward() se ejecuta la retropropagación sobre todo el grafo computacional
  • Después, el .grad de cada parámetro indica cómo debe cambiar para reducir la pérdida

• Optimizador Adam

  • En lugar de usar descenso de gradiente simple (p.data -= lr * p.grad), se usa Adam
  • Mantiene dos promedios móviles por parámetro:
    • m: promedio de gradientes recientes (momentum)
    • v: promedio de los gradientes al cuadrado recientes (adaptación de la tasa de aprendizaje por parámetro)
  • m_hat y v_hat son la corrección de sesgo de m y v, inicializados en 0
  • La tasa de aprendizaje disminuye linealmente durante el entrenamiento
  • Después de la actualización, se reinicia con .grad = 0

• Resultados del entrenamiento

  • Durante 1,000 pasos, la pérdida baja de aproximadamente 3.3 (adivinar al azar entre 27 tokens: −log(1/27)≈3.3) a cerca de 2.37
  • Como mientras más baja mejor, y el mínimo es 0 (predicción perfecta), todavía hay margen de mejora, pero está claro que el modelo está aprendiendo los patrones estadísticos de los nombres

Inferencia

• Una vez terminado el entrenamiento, se pueden muestrear nombres nuevos del modelo; se fijan los parámetros y se ejecuta la pasada hacia adelante en bucle, realimentando cada token generado como la siguiente entrada

• Proceso de muestreo

  • Cada muestra comienza con el token BOS ("inicio de un nombre nuevo")
  • El modelo genera 27 logits → se convierten en probabilidades → se muestrea aleatoriamente un token según esas probabilidades
  • Ese token se realimenta como siguiente entrada, y se repite hasta que el modelo vuelva a generar BOS ("terminado") o se alcance la longitud máxima de secuencia

• Temperatura (Temperature)

  • Antes de softmax, los logits se dividen por la temperatura
  • Temperatura 1.0: se muestrea directamente de la distribución que aprendió el modelo
  • Temperatura baja (por ejemplo, 0.5): hace más aguda la distribución, por lo que es más probable que el modelo sea más conservador y elija las opciones principales
  • Temperatura cercana a 0: siempre elige el único token con mayor probabilidad (decodificación codiciosa)
  • Temperatura alta: aplana la distribución y produce salidas más diversas pero menos consistentes

Cómo ejecutarlo

• Solo se necesita Python (sin pip install, sin dependencias): python train.py
• Toma alrededor de 1 minuto en una MacBook
• Imprime la pérdida en cada paso: baja de ~3.3 (aleatorio) a ~2.37
• Al terminar el entrenamiento, genera nombres nuevos alucinados: "kamon", "ann", "karai", etc.
• También puede ejecutarse en un notebook de Google Colab, y se le pueden hacer preguntas a Gemini
• Se puede probar con otros datasets, entrenar más tiempo aumentando num_steps, o mejorar los resultados aumentando el tamaño del modelo

Etapas del código

Archivo	Contenido agregado
train0.py	Tabla de conteo de bigramas — sin red neuronal, sin gradientes
train1.py	MLP + gradientes manuales (numéricos y analíticos) + SGD
train2.py	Autograd (clase Value) — reemplaza los gradientes manuales
train3.py	Embeddings de posición + atención de una sola cabeza + rmsnorm + residual
train4.py	Atención multi-head + loop de capas — arquitectura GPT completa
train5.py	Optimizador Adam — este es train.py

• En las Revisions del Gist build_microgpt.py se pueden revisar todas las versiones y el diff entre cada paso

Diferencias con los LLM de producción

• microgpt incluye la esencia algorítmica completa del entrenamiento y la ejecución de GPT; la diferencia con un LLM de producción como ChatGPT no cambia el algoritmo central, sino los elementos que permiten que funcione a escala

• Datos

  • En lugar de 32K nombres cortos, se entrena con billones de tokens de texto de internet (páginas web, libros, código, etc.)
  • Eliminación de duplicados, filtrado de calidad y mezcla cuidadosa entre dominios

• Tokenizador

  • En lugar de caracteres individuales, se usa un tokenizador de subpalabras como BPE (Byte Pair Encoding)
  • Fusiona secuencias de caracteres que aparecen juntas con frecuencia en un solo token; palabras comunes como "the" son un solo token, mientras que las raras se dividen en fragmentos
  • Vocabulario de ~100K tokens; como ve más contenido por posición, es mucho más eficiente

• Autograd

  • En lugar de objetos escalares Value en Python puro, usa tensores (grandes arreglos numéricos multidimensionales) y corre en GPU/TPU capaces de realizar miles de millones de operaciones de punto flotante por segundo
  • PyTorch maneja autograd sobre tensores, y kernels CUDA como FlashAttention fusionan varias operaciones
  • Las matemáticas son las mismas; muchos escalares se procesan en paralelo

• Arquitectura

  • microgpt: 4,192 parámetros; un modelo de nivel GPT-4: cientos de miles de millones
  • En general, la red neuronal Transformer es muy parecida, pero mucho más ancha (dimensión de embedding 10,000+) y mucho más profunda (100+ capas)
  • Tipos adicionales de bloques Lego y cambios en el orden:
    • RoPE (embeddings posicionales rotatorios) — en lugar de embeddings posicionales aprendidos
    • GQA (grouped query attention) — reduce el tamaño de la caché KV
    • Activación lineal con compuerta — en lugar de ReLU
    • Capas MoE (mixture of experts)
  • La estructura central se conserva bien: atención (comunicación) y MLP (cálculo) alternando sobre el flujo residual

• Entrenamiento

  • En lugar de un documento por paso, usa lotes masivos (millones de tokens por paso), acumulación de gradientes, precisión mixta (float16/bfloat16) y un ajuste cuidadoso de hiperparámetros
  • El entrenamiento de modelos de frontera corre en miles de GPU durante varios meses

• Optimización

  • microgpt: Adam + disminución lineal simple de la tasa de aprendizaje
  • A gran escala, la optimización es un campo propio: precisión reducida (bfloat16, fp8), entrenamiento en grandes clústeres de GPU
  • Los ajustes del optimizador (tasa de aprendizaje, weight decay, parámetros beta, warmup/schedule de decay) requieren un ajuste fino, y los valores correctos dependen del tamaño del modelo, el tamaño del lote y la composición del dataset
  • Las leyes de escalado (por ejemplo, Chinchilla) guían cómo asignar un presupuesto fijo de cómputo entre el tamaño del modelo y la cantidad de tokens de entrenamiento
  • Equivocarse en estos detalles a gran escala puede desperdiciar millones de dólares en cómputo, por lo que los equipos hacen amplios experimentos pequeños antes de una corrida completa de entrenamiento

• Post-entrenamiento

  • El modelo base que sale del entrenamiento (modelo de "preentrenamiento") es un completador de documentos, no un chatbot
  • El proceso para convertirlo en ChatGPT tiene dos pasos:
    • SFT (supervised fine-tuning): se reemplazan documentos por diálogos curados y se continúa el entrenamiento; algorítmicamente no cambia nada
    • RL (reinforcement learning): el modelo genera respuestas → se les asigna una puntuación (humanos, modelos "juez", algoritmos) → aprende a partir de esa retroalimentación
  • En el fondo sigue entrenando sobre documentos, pero ahora los documentos están compuestos por tokens generados por el propio modelo

• Inferencia

  • Servir el modelo a millones de usuarios requiere su propio stack de ingeniería: batching de solicitudes, gestión y paginación de caché KV (vLLM, etc.), decodificación especulativa para velocidad, cuantización para reducir memoria (correr en int8/int4), y distribución del modelo entre varias GPU
  • En esencia, sigue prediciendo el siguiente token de la secuencia, pero con mucho esfuerzo de ingeniería para hacerlo más rápido

FAQ

• ¿El modelo "entiende" algo?

  • Es una pregunta filosófica, pero mecánicamente no ocurre ninguna magia
  • El modelo es una gran función matemática que mapea tokens de entrada a una distribución de probabilidad sobre el siguiente token
  • Durante el entrenamiento, los parámetros se ajustan para hacer más probable el siguiente token correcto
  • Que eso constituya "entender" depende de cada quien, pero el mecanismo está completamente contenido en estas 200 líneas

• ¿Por qué funciona?

  • El modelo tiene miles de parámetros ajustables, y el optimizador los mueve un poco en cada paso para reducir la pérdida
  • A lo largo de muchos pasos, los parámetros se estabilizan en valores que capturan regularidades estadísticas de los datos
  • En el caso de los nombres: muchos empiezan con consonante, "qu" tiende a aparecer junto, tres consonantes seguidas son raras, etc.
  • El modelo aprende una distribución de probabilidad que refleja eso, no reglas explícitas

• ¿Qué relación tiene con ChatGPT?

  • ChatGPT escala enormemente este mismo bucle central (predicción del siguiente token, muestreo, repetición) y le agrega post-entrenamiento para volverlo conversacional
  • Al chatear, el system prompt, el mensaje del usuario y la respuesta son todos simplemente tokens de una secuencia
  • El modelo completa un documento un token a la vez, igual que microgpt completa nombres

• ¿Qué es una "alucinación"?

  • El modelo genera tokens muestreando desde una distribución de probabilidad
  • No tiene un concepto de verdad; solo conoce secuencias estadísticamente plausibles a la luz de los datos de entrenamiento
  • Que microgpt "alucine" un nombre como "karia" es el mismo fenómeno que ChatGPT diciendo con confianza un hecho falso
  • En ambos casos son completaciones que suenan plausibles pero no son reales

• ¿Por qué es tan lento?

  • microgpt procesa un escalar a la vez en Python puro, y un solo paso de entrenamiento tarda varios segundos
  • En una GPU, las mismas matemáticas procesan millones de escalares en paralelo, ejecutándose varios órdenes de magnitud más rápido

• ¿Se puede hacer que genere mejores nombres?

  • Sí: entrenarlo por más tiempo (aumentar num_steps), incrementar el tamaño del modelo (n_embd, n_layer, n_head), usar un dataset más grande
  • Son los mismos controles que importan a gran escala

• ¿Qué pasa si cambio el dataset?

  • El modelo aprenderá cualquier patrón que esté en los datos
  • Si lo reemplazas por nombres de ciudades, nombres de Pokémon, palabras en inglés o archivos de poemas cortos, aprenderá a generar eso en su lugar
  • No hace falta cambiar el resto del código

[mhj5730]

Gracias por el buen artículo.